Initial commit of the code.
This commit is contained in:
Oleksandr Kozachuk
@@ -0,0 +1,200 @@
|
||||
import unittest
|
||||
import io
|
||||
import os
|
||||
import yaml
|
||||
import argparse
|
||||
import chatmastermind.main
|
||||
from chatmastermind.main import create_chat, ai, handle_question, save_answers
|
||||
from unittest import mock
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
|
||||
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('builtins.open')
|
||||
def test_create_chat_with_tags(self, open_mock, listdir_mock):
|
||||
listdir_mock.return_value = ['testfile.yaml']
|
||||
open_mock.return_value.__enter__.return_value = io.StringIO(yaml.dump(
|
||||
{'question': 'test_content', 'answer': 'some answer',
|
||||
'tags': ['test_tag']}))
|
||||
|
||||
test_chat = create_chat(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('builtins.open')
|
||||
def test_create_chat_with_other_tags(self, open_mock, listdir_mock):
|
||||
listdir_mock.return_value = ['testfile.yaml']
|
||||
open_mock.return_value.__enter__.return_value = io.StringIO(yaml.dump(
|
||||
{'question': 'test_content', 'answer': 'some answer',
|
||||
'tags': ['other_tag']}))
|
||||
|
||||
test_chat = create_chat(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('builtins.open')
|
||||
def test_create_chat_without_tags(self, open_mock, listdir_mock):
|
||||
listdir_mock.return_value = ['testfile.yaml', 'testfile2.yaml']
|
||||
open_mock.side_effect = (
|
||||
io.StringIO(yaml.dump({'question': 'test_content',
|
||||
'answer': 'some answer',
|
||||
'tags': ['test_tag']})),
|
||||
io.StringIO(yaml.dump({'question': 'test_content2',
|
||||
'answer': 'some answer2',
|
||||
'tags': ['test_tag2']})),
|
||||
)
|
||||
|
||||
test_chat = create_chat(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.args = argparse.Namespace(
|
||||
tags=['tag1'],
|
||||
extags=['extag1'],
|
||||
output_tags=None,
|
||||
question='test question',
|
||||
number=3
|
||||
)
|
||||
self.config = {
|
||||
'setting1': 'value1',
|
||||
'setting2': 'value2'
|
||||
}
|
||||
|
||||
@patch("chatmastermind.main.create_chat", return_value="test_chat")
|
||||
@patch("chatmastermind.main.process_tags")
|
||||
@patch("chatmastermind.main.ai", return_value=(["answer1", "answer2", "answer3"],
|
||||
"test_usage"))
|
||||
@patch("chatmastermind.main.pp")
|
||||
@patch("chatmastermind.main.print")
|
||||
@patch("chatmastermind.main.yaml.dump")
|
||||
def test_handle_question(self, _, mock_print, mock_pp, mock_ai,
|
||||
mock_process_tags, mock_create_chat):
|
||||
open_mock = MagicMock()
|
||||
with patch("chatmastermind.main.open", open_mock):
|
||||
handle_question(self.args, self.config, True)
|
||||
mock_process_tags.assert_called_once_with(self.config,
|
||||
self.args.tags,
|
||||
self.args.extags)
|
||||
mock_create_chat.assert_called_once_with(self.args.question,
|
||||
self.args.tags,
|
||||
self.args.extags,
|
||||
self.config)
|
||||
mock_pp.assert_called_once_with("test_chat")
|
||||
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 = '-' * (chatmastermind.main.terminal_width - len(title))
|
||||
expected_calls.append(((f'{title}{title_end}',),))
|
||||
expected_calls.append(((answer,),))
|
||||
expected_calls.append((("-" * chatmastermind.main.terminal_width,),))
|
||||
expected_calls.append(((f"Usage: {mock_ai.return_value[1]}",),))
|
||||
open_mock.assert_has_calls([
|
||||
mock.call(f"{num:02d}.yaml", "w") for num in range(1, 4)
|
||||
] + [mock.call().__enter__(),
|
||||
mock.call().__exit__(None, None, None)] * 3,
|
||||
any_order=True)
|
||||
self.assertEqual(mock_print.call_args_list, expected_calls)
|
||||
|
||||
|
||||
class TestSaveAnswers(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.question = "What is AI?"
|
||||
self.answers = ["AI is Artificial Intelligence",
|
||||
"AI is a simulation of human intelligence"]
|
||||
self.tags = ["ai", "definition"]
|
||||
|
||||
@patch('sys.stdout', new_callable=io.StringIO)
|
||||
def assert_stdout(self, expected_output: str, mock_stdout: io.StringIO):
|
||||
save_answers(self.question, self.answers, self.tags, None)
|
||||
self.assertEqual(mock_stdout.getvalue(), expected_output)
|
||||
|
||||
def test_save_answers(self):
|
||||
try:
|
||||
self.assert_stdout(f"-- ANSWER 1 {'-'*(chatmastermind.main.terminal_width-12)}\n"
|
||||
"AI is Artificial Intelligence\n"
|
||||
f"-- ANSWER 2 {'-'*(chatmastermind.main.terminal_width-12)}\n"
|
||||
"AI is a simulation of human intelligence\n")
|
||||
for idx, answer in enumerate(self.answers, start=1):
|
||||
with open(f"{idx:02d}.yaml", "r") as file:
|
||||
data = yaml.safe_load(file)
|
||||
self.assertEqual(data["question"], self.question)
|
||||
self.assertEqual(data["answer"], answer)
|
||||
self.assertEqual(data["tags"], self.tags)
|
||||
finally:
|
||||
for idx in range(1, len(self.answers) + 1):
|
||||
if os.path.exists(f"{idx:02d}.yaml"):
|
||||
os.remove(f"{idx:02d}.yaml")
|
||||
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user