import unittest import io import pathlib import argparse from chatmastermind.utils import terminal_width from chatmastermind.main import create_parser, ask_cmd from chatmastermind.api_client import ai from chatmastermind.configuration import Config from chatmastermind.storage import create_chat_hist, save_answers, dump_data from unittest import mock from unittest.mock import patch, MagicMock, Mock, ANY class CmmTestCase(unittest.TestCase): """ Base class for all cmm testcases. """ def dummy_config(self, db: str) -> Config: """ Creates a dummy configuration. """ return Config.from_dict( {'system': 'dummy_system', 'db': db, 'openai': {'api_key': 'dummy_key', 'model': 'dummy_model', 'max_tokens': 4000, 'temperature': 1.0, 'top_p': 1, 'frequency_penalty': 0, 'presence_penalty': 0}} ) class TestCreateChat(CmmTestCase): def setUp(self) -> None: self.config = self.dummy_config(db='test_files') self.question = "test question" self.tags = ['test_tag'] @patch('os.listdir') @patch('pathlib.Path.iterdir') @patch('builtins.open') def test_create_chat_with_tags(self, open_mock: MagicMock, iterdir_mock: MagicMock, listdir_mock: MagicMock) -> None: listdir_mock.return_value = ['testfile.txt'] iterdir_mock.return_value = [pathlib.Path(x) for x in listdir_mock.return_value] open_mock.return_value.__enter__.return_value = io.StringIO(dump_data( {'question': 'test_content', 'answer': 'some answer', 'tags': ['test_tag']})) test_chat = create_chat_hist(self.question, self.tags, None, self.config) self.assertEqual(len(test_chat), 4) self.assertEqual(test_chat[0], {'role': 'system', 'content': self.config.system}) self.assertEqual(test_chat[1], {'role': 'user', 'content': 'test_content'}) self.assertEqual(test_chat[2], {'role': 'assistant', 'content': 'some answer'}) self.assertEqual(test_chat[3], {'role': 'user', 'content': self.question}) @patch('os.listdir') @patch('pathlib.Path.iterdir') @patch('builtins.open') def test_create_chat_with_other_tags(self, open_mock: MagicMock, iterdir_mock: MagicMock, listdir_mock: MagicMock) -> None: listdir_mock.return_value = ['testfile.txt'] iterdir_mock.return_value = [pathlib.Path(x) for x in listdir_mock.return_value] open_mock.return_value.__enter__.return_value = io.StringIO(dump_data( {'question': 'test_content', 'answer': 'some answer', 'tags': ['other_tag']})) test_chat = create_chat_hist(self.question, self.tags, None, self.config) self.assertEqual(len(test_chat), 2) self.assertEqual(test_chat[0], {'role': 'system', 'content': self.config.system}) self.assertEqual(test_chat[1], {'role': 'user', 'content': self.question}) @patch('os.listdir') @patch('pathlib.Path.iterdir') @patch('builtins.open') def test_create_chat_without_tags(self, open_mock: MagicMock, iterdir_mock: MagicMock, listdir_mock: MagicMock) -> None: listdir_mock.return_value = ['testfile.txt', 'testfile2.txt'] iterdir_mock.return_value = [pathlib.Path(x) for x in listdir_mock.return_value] open_mock.side_effect = ( io.StringIO(dump_data({'question': 'test_content', 'answer': 'some answer', 'tags': ['test_tag']})), io.StringIO(dump_data({'question': 'test_content2', 'answer': 'some answer2', 'tags': ['test_tag2']})), ) test_chat = create_chat_hist(self.question, [], None, self.config) self.assertEqual(len(test_chat), 6) self.assertEqual(test_chat[0], {'role': 'system', 'content': self.config.system}) self.assertEqual(test_chat[1], {'role': 'user', 'content': 'test_content'}) self.assertEqual(test_chat[2], {'role': 'assistant', 'content': 'some answer'}) self.assertEqual(test_chat[3], {'role': 'user', 'content': 'test_content2'}) self.assertEqual(test_chat[4], {'role': 'assistant', 'content': 'some answer2'}) class TestHandleQuestion(CmmTestCase): def setUp(self) -> None: self.question = "test question" self.args = argparse.Namespace( tags=['tag1'], extags=['extag1'], output_tags=None, question=[self.question], source=None, only_source_code=False, number=3, max_tokens=None, temperature=None, model=None, match_all_tags=False, with_tags=False, with_file=False, ) self.config = self.dummy_config(db='test_files') @patch("chatmastermind.main.create_chat_hist", return_value="test_chat") @patch("chatmastermind.main.print_tag_args") @patch("chatmastermind.main.print_chat_hist") @patch("chatmastermind.main.ai", return_value=(["answer1", "answer2", "answer3"], "test_usage")) @patch("chatmastermind.utils.pp") @patch("builtins.print") def test_ask_cmd(self, mock_print: MagicMock, mock_pp: MagicMock, mock_ai: MagicMock, mock_print_chat_hist: MagicMock, mock_print_tag_args: MagicMock, mock_create_chat_hist: MagicMock) -> None: open_mock = MagicMock() with patch("chatmastermind.storage.open", open_mock): ask_cmd(self.args, self.config) mock_print_tag_args.assert_called_once_with(self.args.tags, self.args.extags, []) mock_create_chat_hist.assert_called_once_with(self.question, self.args.tags, self.args.extags, self.config, False, False, False) mock_print_chat_hist.assert_called_once_with('test_chat', False, self.args.only_source_code) mock_ai.assert_called_with("test_chat", self.config, self.args.number) expected_calls = [] for num, answer in enumerate(mock_ai.return_value[0], start=1): title = f'-- ANSWER {num} ' title_end = '-' * (terminal_width() - len(title)) expected_calls.append(((f'{title}{title_end}',),)) expected_calls.append(((answer,),)) expected_calls.append((("-" * terminal_width(),),)) expected_calls.append(((f"Usage: {mock_ai.return_value[1]}",),)) self.assertEqual(mock_print.call_args_list, expected_calls) open_expected_calls = list([mock.call(f"{num:04d}.txt", "w") for num in range(2, 5)]) open_mock.assert_has_calls(open_expected_calls, any_order=True) class TestSaveAnswers(CmmTestCase): @mock.patch('builtins.open') @mock.patch('chatmastermind.storage.print') def test_save_answers(self, print_mock: MagicMock, open_mock: MagicMock) -> None: question = "Test question?" answers = ["Answer 1", "Answer 2"] tags = ["tag1", "tag2"] otags = ["otag1", "otag2"] config = self.dummy_config(db='test_db') with mock.patch('chatmastermind.storage.pathlib.Path.exists', return_value=True), \ mock.patch('chatmastermind.storage.yaml.dump'), \ mock.patch('io.StringIO') as stringio_mock: stringio_instance = stringio_mock.return_value stringio_instance.getvalue.side_effect = ["question", "answer1", "answer2"] save_answers(question, answers, tags, otags, config) open_calls = [ mock.call(pathlib.Path('test_db/.next'), 'r'), mock.call(pathlib.Path('test_db/.next'), 'w'), ] open_mock.assert_has_calls(open_calls, any_order=True) class TestAI(CmmTestCase): @patch("openai.ChatCompletion.create") def test_ai(self, mock_create: MagicMock) -> None: mock_create.return_value = { 'choices': [ {'message': {'content': 'response_text_1'}}, {'message': {'content': 'response_text_2'}} ], 'usage': {'tokens': 10} } chat = [{"role": "system", "content": "hello ai"}] config = self.dummy_config(db='dummy') config.openai.model = "text-davinci-002" config.openai.max_tokens = 150 config.openai.temperature = 0.5 result = ai(chat, config, 2) expected_result = (['response_text_1', 'response_text_2'], {'tokens': 10}) self.assertEqual(result, expected_result) class TestCreateParser(CmmTestCase): def test_create_parser(self) -> None: with patch('argparse.ArgumentParser.add_subparsers') as mock_add_subparsers: mock_cmdparser = Mock() mock_add_subparsers.return_value = mock_cmdparser parser = create_parser() self.assertIsInstance(parser, argparse.ArgumentParser) mock_add_subparsers.assert_called_once_with(dest='command', title='commands', description='supported commands', required=True) mock_cmdparser.add_parser.assert_any_call('ask', parents=ANY, help=ANY, aliases=ANY) mock_cmdparser.add_parser.assert_any_call('hist', parents=ANY, help=ANY, aliases=ANY) mock_cmdparser.add_parser.assert_any_call('tags', help=ANY, aliases=ANY) mock_cmdparser.add_parser.assert_any_call('config', help=ANY, aliases=ANY) mock_cmdparser.add_parser.assert_any_call('print', help=ANY, aliases=ANY) self.assertTrue('.config.yaml' in parser.get_default('config'))