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.storage import create_chat_hist, save_answers, dump_data from unittest import mock from unittest.mock import patch, MagicMock, Mock, ANY class TestCreateChat(unittest.TestCase): def setUp(self): self.config = { 'system': 'System text', '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, iterdir_mock, listdir_mock): 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, iterdir_mock, listdir_mock): 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, iterdir_mock, listdir_mock): 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(unittest.TestCase): def setUp(self): 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 = { 'db': 'test_files', 'setting1': 'value1', 'setting2': 'value2', 'openai': {}, } @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, mock_pp, mock_ai, mock_print_chat_hist, mock_print_tag_args, mock_create_chat_hist): 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(unittest.TestCase): @mock.patch('builtins.open') @mock.patch('chatmastermind.storage.print') def test_save_answers(self, print_mock, open_mock): question = "Test question?" answers = ["Answer 1", "Answer 2"] tags = ["tag1", "tag2"] otags = ["otag1", "otag2"] 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(unittest.TestCase): @patch("openai.ChatCompletion.create") def test_ai(self, mock_create: MagicMock): mock_create.return_value = { 'choices': [ {'message': {'content': 'response_text_1'}}, {'message': {'content': 'response_text_2'}} ], 'usage': {'tokens': 10} } number = 2 chat = [{"role": "system", "content": "hello ai"}] config = { "openai": { "model": "text-davinci-002", "temperature": 0.5, "max_tokens": 150, "top_p": 1, "n": number, "frequency_penalty": 0, "presence_penalty": 0 } } result = ai(chat, config, number) expected_result = (['response_text_1', 'response_text_2'], {'tokens': 10}) self.assertEqual(result, expected_result) class TestCreateParser(unittest.TestCase): def test_create_parser(self): 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) mock_cmdparser.add_parser.assert_any_call('hist', parents=ANY, help=ANY) mock_cmdparser.add_parser.assert_any_call('tag', help=ANY) mock_cmdparser.add_parser.assert_any_call('config', help=ANY) mock_cmdparser.add_parser.assert_any_call('print', help=ANY) self.assertTrue('.config.yaml' in parser.get_default('config'))