diff --git a/tests/test_main.py b/tests/test_main.py index ce9121a..2493d46 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -1,14 +1,14 @@ 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 +# 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 +# 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): @@ -30,207 +30,207 @@ class CmmTestCase(unittest.TestCase): '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')) +# +# +# 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'))