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_text = [f"{self.chat.messages[1].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()), 1) self.assertEqual(message.question, Question("""What is this? ``` It is embedded code ```"""))