237 lines
11 KiB
Python
237 lines
11 KiB
Python
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'))
|