Compare commits
1 Commits
e8343fde01
...
cdd9153288
| Author | SHA1 | Date | |
|---|---|---|---|
| cdd9153288 |
@ -24,7 +24,6 @@ def openai_config_valid(conf: dict[str, str | float | int]) -> bool:
|
||||
str(conf['api_key'])
|
||||
str(conf['model'])
|
||||
int(conf['max_tokens'])
|
||||
float(conf['temperature'])
|
||||
float(conf['top_p'])
|
||||
float(conf['frequency_penalty'])
|
||||
float(conf['presence_penalty'])
|
||||
|
||||
@ -5,46 +5,25 @@ 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, OpenAIConfig
|
||||
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(
|
||||
system='dummy_system',
|
||||
db=db,
|
||||
openai=OpenAIConfig(
|
||||
api_key='dummy_key',
|
||||
model='dummy_model',
|
||||
max_tokens=4000,
|
||||
temperature=1.0,
|
||||
top_p=1,
|
||||
frequency_penalty=0,
|
||||
presence_penalty=0
|
||||
)
|
||||
)
|
||||
class TestCreateChat(unittest.TestCase):
|
||||
|
||||
|
||||
class TestCreateChat(CmmTestCase):
|
||||
|
||||
def setUp(self) -> None:
|
||||
self.config = self.dummy_config(db='test_files')
|
||||
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: MagicMock, iterdir_mock: MagicMock, listdir_mock: MagicMock) -> None:
|
||||
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(
|
||||
@ -66,7 +45,7 @@ class TestCreateChat(CmmTestCase):
|
||||
@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:
|
||||
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(
|
||||
@ -84,7 +63,7 @@ class TestCreateChat(CmmTestCase):
|
||||
@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:
|
||||
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 = (
|
||||
@ -111,9 +90,9 @@ class TestCreateChat(CmmTestCase):
|
||||
{'role': 'assistant', 'content': 'some answer2'})
|
||||
|
||||
|
||||
class TestHandleQuestion(CmmTestCase):
|
||||
class TestHandleQuestion(unittest.TestCase):
|
||||
|
||||
def setUp(self) -> None:
|
||||
def setUp(self):
|
||||
self.question = "test question"
|
||||
self.args = argparse.Namespace(
|
||||
tags=['tag1'],
|
||||
@ -130,7 +109,12 @@ class TestHandleQuestion(CmmTestCase):
|
||||
with_tags=False,
|
||||
with_file=False,
|
||||
)
|
||||
self.config = self.dummy_config(db='test_files')
|
||||
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")
|
||||
@ -138,9 +122,9 @@ class TestHandleQuestion(CmmTestCase):
|
||||
@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:
|
||||
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)
|
||||
@ -171,15 +155,15 @@ class TestHandleQuestion(CmmTestCase):
|
||||
open_mock.assert_has_calls(open_expected_calls, any_order=True)
|
||||
|
||||
|
||||
class TestSaveAnswers(CmmTestCase):
|
||||
class TestSaveAnswers(unittest.TestCase):
|
||||
@mock.patch('builtins.open')
|
||||
@mock.patch('chatmastermind.storage.print')
|
||||
def test_save_answers(self, print_mock: MagicMock, open_mock: MagicMock) -> None:
|
||||
def test_save_answers(self, print_mock, open_mock):
|
||||
question = "Test question?"
|
||||
answers = ["Answer 1", "Answer 2"]
|
||||
tags = ["tag1", "tag2"]
|
||||
otags = ["otag1", "otag2"]
|
||||
config = self.dummy_config(db='test_db')
|
||||
config = {'db': 'test_db'}
|
||||
|
||||
with mock.patch('chatmastermind.storage.pathlib.Path.exists', return_value=True), \
|
||||
mock.patch('chatmastermind.storage.yaml.dump'), \
|
||||
@ -195,10 +179,10 @@ class TestSaveAnswers(CmmTestCase):
|
||||
open_mock.assert_has_calls(open_calls, any_order=True)
|
||||
|
||||
|
||||
class TestAI(CmmTestCase):
|
||||
class TestAI(unittest.TestCase):
|
||||
|
||||
@patch("openai.ChatCompletion.create")
|
||||
def test_ai(self, mock_create: MagicMock) -> None:
|
||||
def test_ai(self, mock_create: MagicMock):
|
||||
mock_create.return_value = {
|
||||
'choices': [
|
||||
{'message': {'content': 'response_text_1'}},
|
||||
@ -207,20 +191,28 @@ class TestAI(CmmTestCase):
|
||||
'usage': {'tokens': 10}
|
||||
}
|
||||
|
||||
number = 2
|
||||
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
|
||||
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, 2)
|
||||
result = ai(chat, config, number)
|
||||
expected_result = (['response_text_1', 'response_text_2'],
|
||||
{'tokens': 10})
|
||||
self.assertEqual(result, expected_result)
|
||||
|
||||
|
||||
class TestCreateParser(CmmTestCase):
|
||||
def test_create_parser(self) -> None:
|
||||
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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user