from .test_main import CmmTestCase from chatmastermind.message import source_code, MessageError, Question, Answer class SourceCodeTestCase(CmmTestCase): def test_source_code_with_include_delims(self) -> None: text = """ Some text before the code block ```python print("Hello, World!") ``` Some text after the code block ```python x = 10 y = 20 print(x + y) ``` """ expected_result = [ " ```python\n print(\"Hello, World!\")\n ```\n", " ```python\n x = 10\n y = 20\n print(x + y)\n ```\n" ] result = source_code(text, include_delims=True) self.assertEqual(result, expected_result) def test_source_code_without_include_delims(self) -> None: text = """ Some text before the code block ```python print("Hello, World!") ``` Some text after the code block ```python x = 10 y = 20 print(x + y) ``` """ expected_result = [ " print(\"Hello, World!\")\n", " x = 10\n y = 20\n print(x + y)\n" ] result = source_code(text, include_delims=False) self.assertEqual(result, expected_result) def test_source_code_with_single_code_block(self) -> None: text = "```python\nprint(\"Hello, World!\")\n```" expected_result = ["```python\nprint(\"Hello, World!\")\n```\n"] result = source_code(text, include_delims=True) self.assertEqual(result, expected_result) def test_source_code_with_no_code_blocks(self) -> None: text = "Some text without any code blocks" expected_result: list[str] = [] result = source_code(text, include_delims=True) self.assertEqual(result, expected_result) class QuestionTestCase(CmmTestCase): def test_question_with_prefix(self) -> None: with self.assertRaises(MessageError): Question("=== QUESTION === What is your name?") def test_question_without_prefix(self) -> None: question = Question("What is your favorite color?") self.assertIsInstance(question, Question) self.assertEqual(question, "What is your favorite color?") class AnswerTestCase(CmmTestCase): def test_answer_with_prefix(self) -> None: with self.assertRaises(MessageError): Answer("=== ANSWER === Yes") def test_answer_without_prefix(self) -> None: answer = Answer("No") self.assertIsInstance(answer, Answer) self.assertEqual(answer, "No")