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( # or_tags=['tag1'], # and_tags=None, # exclude_tags=['xtag1'], # output_tags=None, # question=[self.question], # source=None, # source_code_only=False, # num_answers=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.or_tags, # self.args.exclude_tags, # []) # mock_create_chat_hist.assert_called_once_with(self.question, # self.args.or_tags, # self.args.exclude_tags, # self.config, # match_all_tags=False, # with_tags=False, # with_file=False) # mock_print_chat_hist.assert_called_once_with('test_chat', # False, # self.args.source_code_only) # mock_ai.assert_called_with("test_chat", # self.config, # self.args.num_answers) # 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'))