question_cmd: added testclass for the 'question_cmd()' function
This commit is contained in:
parent
84961e5d0d
commit
bcb3a2d45a
@ -3,10 +3,15 @@ import unittest
|
|||||||
import argparse
|
import argparse
|
||||||
import tempfile
|
import tempfile
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import MagicMock
|
from unittest import mock
|
||||||
from chatmastermind.commands.question import create_message
|
from unittest.mock import MagicMock, call, ANY
|
||||||
|
from typing import Optional
|
||||||
|
from chatmastermind.configuration import Config
|
||||||
|
from chatmastermind.commands.question import create_message, question_cmd
|
||||||
|
from chatmastermind.tags import Tag
|
||||||
from chatmastermind.message import Message, Question, Answer
|
from chatmastermind.message import Message, Question, Answer
|
||||||
from chatmastermind.chat import ChatDB
|
from chatmastermind.chat import Chat, ChatDB
|
||||||
|
from chatmastermind.ai import AI, AIResponse, Tokens
|
||||||
|
|
||||||
|
|
||||||
class TestMessageCreate(unittest.TestCase):
|
class TestMessageCreate(unittest.TestCase):
|
||||||
@ -16,10 +21,10 @@ class TestMessageCreate(unittest.TestCase):
|
|||||||
"""
|
"""
|
||||||
def setUp(self) -> None:
|
def setUp(self) -> None:
|
||||||
# create ChatDB structure
|
# create ChatDB structure
|
||||||
self.db_path = tempfile.TemporaryDirectory()
|
self.db_dir = tempfile.TemporaryDirectory()
|
||||||
self.cache_path = tempfile.TemporaryDirectory()
|
self.cache_dir = tempfile.TemporaryDirectory()
|
||||||
self.chat = ChatDB.from_dir(cache_path=Path(self.cache_path.name),
|
self.chat = ChatDB.from_dir(cache_path=Path(self.cache_dir.name),
|
||||||
db_path=Path(self.db_path.name))
|
db_path=Path(self.db_dir.name))
|
||||||
# create some messages
|
# create some messages
|
||||||
self.message_text = Message(Question("What is this?"),
|
self.message_text = Message(Question("What is this?"),
|
||||||
Answer("It is pure text"))
|
Answer("It is pure text"))
|
||||||
@ -74,6 +79,7 @@ Aaaand again some text."""
|
|||||||
os.remove(self.source_file1.name)
|
os.remove(self.source_file1.name)
|
||||||
os.remove(self.source_file2.name)
|
os.remove(self.source_file2.name)
|
||||||
os.remove(self.source_file3.name)
|
os.remove(self.source_file3.name)
|
||||||
|
os.remove(self.source_file4.name)
|
||||||
|
|
||||||
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
|
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
|
||||||
# exclude '.next'
|
# exclude '.next'
|
||||||
@ -81,10 +87,10 @@ Aaaand again some text."""
|
|||||||
|
|
||||||
def test_message_file_created(self) -> None:
|
def test_message_file_created(self) -> None:
|
||||||
self.args.ask = ["What is this?"]
|
self.args.ask = ["What is this?"]
|
||||||
cache_dir_files = self.message_list(self.cache_path)
|
cache_dir_files = self.message_list(self.cache_dir)
|
||||||
self.assertEqual(len(cache_dir_files), 0)
|
self.assertEqual(len(cache_dir_files), 0)
|
||||||
create_message(self.chat, self.args)
|
create_message(self.chat, self.args)
|
||||||
cache_dir_files = self.message_list(self.cache_path)
|
cache_dir_files = self.message_list(self.cache_dir)
|
||||||
self.assertEqual(len(cache_dir_files), 1)
|
self.assertEqual(len(cache_dir_files), 1)
|
||||||
message = Message.from_file(cache_dir_files[0])
|
message = Message.from_file(cache_dir_files[0])
|
||||||
self.assertIsInstance(message, Message)
|
self.assertIsInstance(message, Message)
|
||||||
@ -193,3 +199,138 @@ This is embedded source code.
|
|||||||
It is embedded code
|
It is embedded code
|
||||||
```
|
```
|
||||||
"""))
|
"""))
|
||||||
|
|
||||||
|
|
||||||
|
class TestQuestionCmd(unittest.TestCase):
|
||||||
|
|
||||||
|
def setUp(self) -> None:
|
||||||
|
# create DB and cache
|
||||||
|
self.db_dir = tempfile.TemporaryDirectory()
|
||||||
|
self.cache_dir = tempfile.TemporaryDirectory()
|
||||||
|
# create configuration
|
||||||
|
self.config = Config()
|
||||||
|
self.config.cache = self.cache_dir.name
|
||||||
|
self.config.db = self.db_dir.name
|
||||||
|
# create a mock argparse.Namespace
|
||||||
|
self.args = argparse.Namespace(
|
||||||
|
ask=['What is the meaning of life?'],
|
||||||
|
num_answers=1,
|
||||||
|
output_tags=['science'],
|
||||||
|
AI='openai',
|
||||||
|
model='gpt-3.5-turbo',
|
||||||
|
or_tags=None,
|
||||||
|
and_tags=None,
|
||||||
|
exclude_tags=None,
|
||||||
|
source_text=None,
|
||||||
|
source_code=None,
|
||||||
|
create=None,
|
||||||
|
repeat=None,
|
||||||
|
process=None
|
||||||
|
)
|
||||||
|
|
||||||
|
def input_message(self, args: argparse.Namespace) -> Message:
|
||||||
|
"""
|
||||||
|
Create the expected input message for a question using the
|
||||||
|
given arguments.
|
||||||
|
"""
|
||||||
|
# NOTE: we only use the first question from the "ask" list
|
||||||
|
# -> message creation using "question.create_message()" is
|
||||||
|
# tested above
|
||||||
|
# the answer is always empty for the input message
|
||||||
|
return Message(Question(args.ask[0]),
|
||||||
|
tags=args.output_tags,
|
||||||
|
ai=args.AI,
|
||||||
|
model=args.model)
|
||||||
|
|
||||||
|
def mock_request(self,
|
||||||
|
question: Message,
|
||||||
|
chat: Chat,
|
||||||
|
num_answers: int = 1,
|
||||||
|
otags: Optional[set[Tag]] = None) -> AIResponse:
|
||||||
|
"""
|
||||||
|
Mock the 'ai.request()' function
|
||||||
|
"""
|
||||||
|
question.answer = Answer("Answer 0")
|
||||||
|
question.tags = otags
|
||||||
|
question.ai = 'FakeAI'
|
||||||
|
question.model = 'FakeModel'
|
||||||
|
answers: list[Message] = [question]
|
||||||
|
for n in range(1, num_answers):
|
||||||
|
answers.append(Message(question=question.question,
|
||||||
|
answer=Answer(f"Answer {n}"),
|
||||||
|
tags=otags,
|
||||||
|
ai='FakeAI',
|
||||||
|
model='FakeModel'))
|
||||||
|
return AIResponse(answers, Tokens(10, 10, 20))
|
||||||
|
|
||||||
|
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
|
||||||
|
# exclude '.next'
|
||||||
|
return sorted([f for f in Path(tmp_dir.name).glob('*.[ty]*')])
|
||||||
|
|
||||||
|
@mock.patch('chatmastermind.commands.question.create_ai')
|
||||||
|
def test_ask_single_answer(self, mock_create_ai: MagicMock) -> None:
|
||||||
|
"""
|
||||||
|
Test single answer with no errors
|
||||||
|
"""
|
||||||
|
# create a mock AI instance
|
||||||
|
ai = MagicMock(spec=AI)
|
||||||
|
ai.request.side_effect = self.mock_request
|
||||||
|
mock_create_ai.return_value = ai
|
||||||
|
expected_question = self.input_message(self.args)
|
||||||
|
expected_responses = self.mock_request(expected_question,
|
||||||
|
Chat([]),
|
||||||
|
self.args.num_answers,
|
||||||
|
self.args.output_tags).messages
|
||||||
|
|
||||||
|
# execute the command
|
||||||
|
question_cmd(self.args, self.config)
|
||||||
|
|
||||||
|
# check for correct request call
|
||||||
|
ai.request.assert_called_once_with(expected_question,
|
||||||
|
ANY,
|
||||||
|
self.args.num_answers,
|
||||||
|
self.args.output_tags)
|
||||||
|
# check for the expected message files
|
||||||
|
chat = ChatDB.from_dir(Path(self.cache_dir.name),
|
||||||
|
Path(self.db_dir.name))
|
||||||
|
cached_msg = chat.msg_gather(loc='cache')
|
||||||
|
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
|
||||||
|
self.assertSequenceEqual(cached_msg, expected_responses)
|
||||||
|
|
||||||
|
@mock.patch('chatmastermind.commands.question.ChatDB.from_dir')
|
||||||
|
@mock.patch('chatmastermind.commands.question.create_ai')
|
||||||
|
def test_ask_single_answer_mocked(self, mock_create_ai: MagicMock, mock_from_dir: MagicMock) -> None:
|
||||||
|
"""
|
||||||
|
Test single answer with no errors (mocked ChatDB version)
|
||||||
|
"""
|
||||||
|
chat = MagicMock(spec=ChatDB)
|
||||||
|
mock_from_dir.return_value = chat
|
||||||
|
|
||||||
|
# create a mock AI instance
|
||||||
|
ai = MagicMock(spec=AI)
|
||||||
|
ai.request.side_effect = self.mock_request
|
||||||
|
mock_create_ai.return_value = ai
|
||||||
|
expected_question = self.input_message(self.args)
|
||||||
|
expected_responses = self.mock_request(expected_question,
|
||||||
|
Chat([]),
|
||||||
|
self.args.num_answers,
|
||||||
|
self.args.output_tags).messages
|
||||||
|
|
||||||
|
# execute the command
|
||||||
|
question_cmd(self.args, self.config)
|
||||||
|
|
||||||
|
# check for correct request call
|
||||||
|
ai.request.assert_called_once_with(expected_question,
|
||||||
|
chat,
|
||||||
|
self.args.num_answers,
|
||||||
|
self.args.output_tags)
|
||||||
|
|
||||||
|
# check for the correct ChatDB calls:
|
||||||
|
# - initial question has been written (prior to the actual request)
|
||||||
|
# - responses have been written (after the request)
|
||||||
|
chat.cache_write.assert_has_calls([call([expected_question]),
|
||||||
|
call(expected_responses)],
|
||||||
|
any_order=False)
|
||||||
|
|
||||||
|
# check that the messages have not been added to the internal message list
|
||||||
|
chat.cache_add.assert_not_called()
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user