196 lines
7.4 KiB
Python
196 lines
7.4 KiB
Python
import os
|
|
import unittest
|
|
import argparse
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
from chatmastermind.commands.question import create_message
|
|
from chatmastermind.message import Message, Question, Answer
|
|
from chatmastermind.chat import ChatDB
|
|
|
|
|
|
class TestMessageCreate(unittest.TestCase):
|
|
"""
|
|
Test if messages created by the 'question' command have
|
|
the correct format.
|
|
"""
|
|
def setUp(self) -> None:
|
|
# create ChatDB structure
|
|
self.db_path = tempfile.TemporaryDirectory()
|
|
self.cache_path = tempfile.TemporaryDirectory()
|
|
self.chat = ChatDB.from_dir(cache_path=Path(self.cache_path.name),
|
|
db_path=Path(self.db_path.name))
|
|
# create some messages
|
|
self.message_text = Message(Question("What is this?"),
|
|
Answer("It is pure text"))
|
|
self.message_code = Message(Question("What is this?"),
|
|
Answer("Text\n```\nIt is embedded code\n```\ntext"))
|
|
self.chat.add_to_db([self.message_text, self.message_code])
|
|
# create arguments mock
|
|
self.args = MagicMock(spec=argparse.Namespace)
|
|
self.args.source_text = None
|
|
self.args.source_code = None
|
|
self.args.AI = None
|
|
self.args.model = None
|
|
self.args.output_tags = None
|
|
# File 1 : no source code block, only text
|
|
self.source_file1 = tempfile.NamedTemporaryFile(delete=False)
|
|
self.source_file1_content = """This is just text.
|
|
No source code.
|
|
Nope. Go look elsewhere!"""
|
|
with open(self.source_file1.name, 'w') as f:
|
|
f.write(self.source_file1_content)
|
|
# File 2 : one embedded source code block
|
|
self.source_file2 = tempfile.NamedTemporaryFile(delete=False)
|
|
self.source_file2_content = """This is just text.
|
|
```
|
|
This is embedded source code.
|
|
```
|
|
And some text again."""
|
|
with open(self.source_file2.name, 'w') as f:
|
|
f.write(self.source_file2_content)
|
|
# File 3 : all source code
|
|
self.source_file3 = tempfile.NamedTemporaryFile(delete=False)
|
|
self.source_file3_content = """This is all source code.
|
|
Yes, really.
|
|
Language is called 'brainfart'."""
|
|
with open(self.source_file3.name, 'w') as f:
|
|
f.write(self.source_file3_content)
|
|
# File 4 : two source code blocks
|
|
self.source_file4 = tempfile.NamedTemporaryFile(delete=False)
|
|
self.source_file4_content = """This is just text.
|
|
```
|
|
This is embedded source code.
|
|
```
|
|
And some text again.
|
|
```
|
|
This is embedded source code.
|
|
```
|
|
Aaaand again some text."""
|
|
with open(self.source_file4.name, 'w') as f:
|
|
f.write(self.source_file4_content)
|
|
|
|
def tearDown(self) -> None:
|
|
os.remove(self.source_file1.name)
|
|
os.remove(self.source_file2.name)
|
|
os.remove(self.source_file3.name)
|
|
|
|
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
|
|
# exclude '.next'
|
|
return list(Path(tmp_dir.name).glob('*.[ty]*'))
|
|
|
|
def test_message_file_created(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
cache_dir_files = self.message_list(self.cache_path)
|
|
self.assertEqual(len(cache_dir_files), 0)
|
|
create_message(self.chat, self.args)
|
|
cache_dir_files = self.message_list(self.cache_path)
|
|
self.assertEqual(len(cache_dir_files), 1)
|
|
message = Message.from_file(cache_dir_files[0])
|
|
self.assertIsInstance(message, Message)
|
|
self.assertEqual(message.question, Question("What is this?")) # type: ignore [union-attr]
|
|
|
|
def test_single_question(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
self.assertEqual(message.question, Question("What is this?"))
|
|
self.assertEqual(len(message.question.source_code()), 0)
|
|
|
|
def test_multipart_question(self) -> None:
|
|
self.args.ask = ["What is this", "'bard' thing?", "Is it good?"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
self.assertEqual(message.question, Question("""What is this
|
|
|
|
'bard' thing?
|
|
|
|
Is it good?"""))
|
|
|
|
def test_single_question_with_text_only_file(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
self.args.source_text = [f"{self.source_file1.name}"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
# file contains no source code (only text)
|
|
# -> don't expect any in the question
|
|
self.assertEqual(len(message.question.source_code()), 0)
|
|
self.assertEqual(message.question, Question(f"""What is this?
|
|
|
|
{self.source_file1_content}"""))
|
|
|
|
def test_single_question_with_text_file_and_embedded_code(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
self.args.source_code = [f"{self.source_file2.name}"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
# file contains 1 source code block
|
|
# -> expect it in the question
|
|
self.assertEqual(len(message.question.source_code()), 1)
|
|
self.assertEqual(message.question, Question("""What is this?
|
|
|
|
```
|
|
This is embedded source code.
|
|
```
|
|
"""))
|
|
|
|
def test_single_question_with_code_only_file(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
self.args.source_code = [f"{self.source_file3.name}"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
# file is complete source code
|
|
self.assertEqual(len(message.question.source_code()), 1)
|
|
self.assertEqual(message.question, Question(f"""What is this?
|
|
|
|
```
|
|
{self.source_file3_content}
|
|
```"""))
|
|
|
|
def test_single_question_with_text_file_and_multi_embedded_code(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
self.args.source_code = [f"{self.source_file4.name}"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
# file contains 2 source code blocks
|
|
# -> expect them in the question
|
|
self.assertEqual(len(message.question.source_code()), 2)
|
|
self.assertEqual(message.question, Question("""What is this?
|
|
|
|
```
|
|
This is embedded source code.
|
|
```
|
|
|
|
|
|
```
|
|
This is embedded source code.
|
|
```
|
|
"""))
|
|
|
|
def test_single_question_with_text_only_message(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
self.args.source_text = [f"{self.chat.messages[0].file_path}"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
# file contains no source code (only text)
|
|
# -> don't expect any in the question
|
|
self.assertEqual(len(message.question.source_code()), 0)
|
|
self.assertEqual(message.question, Question(f"""What is this?
|
|
|
|
{self.message_text.answer}"""))
|
|
|
|
def test_single_question_with_message_and_embedded_code(self) -> None:
|
|
self.args.ask = ["What is this?"]
|
|
self.args.source_code = [f"{self.chat.messages[1].file_path}"]
|
|
message = create_message(self.chat, self.args)
|
|
self.assertIsInstance(message, Message)
|
|
# answer contains 1 source code block
|
|
# -> expect it in the question
|
|
self.assertEqual(len(message.question.source_code()), 1)
|
|
self.assertEqual(message.question, Question("""What is this?
|
|
|
|
```
|
|
It is embedded code
|
|
```
|
|
"""))
|