tests: added testcases for Message.from_file() and Message.to_file()

chatmastermind/message.py

@@ -170,11 +170,11 @@ class Message():
     and a file path.
     """
     question: Question
-    answer: Optional[Answer]
+    answer: Optional[Answer] = None
-    tags: Optional[set[Tag]]
+    tags: Optional[set[Tag]] = None
-    ai: Optional[str]
+    ai: Optional[str] = None
-    model: Optional[str]
+    model: Optional[str] = None
-    file_path: Optional[pathlib.Path]
+    file_path: Optional[pathlib.Path] = None
     # class variables
     file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml']
     tags_yaml_key: ClassVar[str] = 'tags'
@@ -213,12 +213,31 @@ class Message():
         return tags
 
     @classmethod
-    def from_file(cls: Type[MessageInst], file_path: pathlib.Path,  # noqa: 11
+    def from_file(cls: Type[MessageInst], file_path: pathlib.Path,
                   tags_or: Optional[set[Tag]] = None,
                   tags_and: Optional[set[Tag]] = None,
                   tags_not: Optional[set[Tag]] = None) -> Optional[MessageInst]:
         """
-        Create a Message from the given file. Expects the following file structures:
+        Create a Message from the given file. Returns 'None' if the message does
+        not fulfill the tag requirements.
+        """
+        if not file_path.exists():
+            raise MessageError(f"Message file '{file_path}' does not exist")
+        if file_path.suffix not in cls.file_suffixes:
+            raise MessageError(f"File type '{file_path.suffix}' is not supported")
+
+        if file_path.suffix == '.txt':
+            return cls.__from_file_txt(file_path, tags_or, tags_and, tags_not)
+        else:
+            return cls.__from_file_yaml(file_path, tags_or, tags_and, tags_not)
+
+    @classmethod
+    def __from_file_txt(cls: Type[MessageInst], file_path: pathlib.Path,  # noqa: 11
+                        tags_or: Optional[set[Tag]] = None,
+                        tags_and: Optional[set[Tag]] = None,
+                        tags_not: Optional[set[Tag]] = None) -> Optional[MessageInst]:
+        """
+        Create a Message from the given TXT file. Expects the following file structures:
         For '.txt':
             * TagLine [Optional]
             * AI [Optional]
@@ -226,27 +245,15 @@ class Message():
             * Question.txt_header
             * Question
             * Answer.txt_header [Optional]
-            # Answer [Optional]
+            * Answer [Optional]
-        For '.yaml':
+
-            * Question.yaml_key: single or multiline string
-            * Answer.yaml_key: single or multiline string [Optional]
-            * Message.tags_yaml_key: list of strings [Optional]
-            * Message.ai_yaml_key: str [Optional]
-            * Message.model_yaml_key: str [Optional]
         Returns 'None' if the message does not fulfill the tag requirements.
         """
-        if not file_path.exists():
-            raise MessageError(f"Message file '{file_path}' does not exist")
-        if file_path.suffix not in cls.file_suffixes:
-            raise MessageError(f"File type '{file_path.suffix}' is not supported")
-
         tags: set[Tag] = set()
         question: Question
         answer: Optional[Answer] = None
         ai: Optional[str] = None
         model: Optional[str] = None
-        # TXT
-        if file_path.suffix == '.txt':
         with open(file_path, "r") as fd:
             # TagLine (Optional)
             try:
@@ -280,8 +287,23 @@ class Message():
             except ValueError:
                 question = Question.from_list(text[question_idx:])
             return cls(question, answer, tags, ai, model, file_path)
-        # YAML
+
-        else:
+    @classmethod
+    def __from_file_yaml(cls: Type[MessageInst], file_path: pathlib.Path,
+                         tags_or: Optional[set[Tag]] = None,
+                         tags_and: Optional[set[Tag]] = None,
+                         tags_not: Optional[set[Tag]] = None) -> Optional[MessageInst]:
+        """
+        Create a Message from the given YAML file. Expects the following file structures:
+            * Question.yaml_key: single or multiline string
+            * Answer.yaml_key: single or multiline string [Optional]
+            * Message.tags_yaml_key: list of strings [Optional]
+            * Message.ai_yaml_key: str [Optional]
+            * Message.model_yaml_key: str [Optional]
+
+        Returns 'None' if the message does not fulfill the tag requirements.
+        """
+        tags: set[Tag] = set()
         with open(file_path, "r") as fd:
             data = yaml.load(fd, Loader=yaml.FullLoader)
             if tags_or or tags_and or tags_not:
@@ -295,21 +317,8 @@ class Message():
 
     def to_file(self, file_path: Optional[pathlib.Path]) -> None:  # noqa: 11
         """
-        Write a Message to the given file. Creates the following file structures:
+        Write a Message to the given file. Type is determined based on the suffix.
-        For '.txt':
+        Currently supported suffixes: ['.txt', '.yaml']
-            * TagLine
-            * AI [Optional]
-            * Model [Optional]
-            * Question.txt_header
-            * Question
-            * Answer.txt_header
-            * Answer
-        For '.yaml':
-            * Question.yaml_key: single or multiline string
-            * Answer.yaml_key: single or multiline string
-            * Message.tags_yaml_key: list of strings
-            * Message.ai_yaml_key: str [Optional]
-            * Message.model_yaml_key: str [Optional]
         """
         if file_path:
             self.file_path = file_path
@@ -319,18 +328,44 @@ class Message():
             raise MessageError(f"File type '{self.file_path.suffix}' is not supported")
         # TXT
         if self.file_path.suffix == '.txt':
-            with open(self.file_path, "w") as fd:
+            return self.__to_file_txt(self.file_path)
-                msg_tags = self.tags or set()
+        elif self.file_path.suffix == '.yaml':
-                fd.write(f'{TagLine.from_set(msg_tags)}\n')
+            return self.__to_file_yaml(self.file_path)
+
+    def __to_file_txt(self, file_path: pathlib.Path) -> None:
+        """
+        Write a Message to the given file in TXT format.
+        Creates the following file structures:
+            * TagLine
+            * AI [Optional]
+            * Model [Optional]
+            * Question.txt_header
+            * Question
+            * Answer.txt_header
+            * Answer
+        """
+        with open(file_path, "w") as fd:
+            if self.tags:
+                fd.write(f'{TagLine.from_set(self.tags)}\n')
             if self.ai:
                 fd.write(f'{AILine.from_ai(self.ai)}\n')
             if self.model:
                 fd.write(f'{ModelLine.from_model(self.model)}\n')
             fd.write(f'{Question.txt_header}\n{self.question}\n')
+            if self.answer:
                 fd.write(f'{Answer.txt_header}\n{self.answer}\n')
-        # YAML
+
-        elif self.file_path.suffix == '.yaml':
+    def __to_file_yaml(self, file_path: pathlib.Path) -> None:
-            with open(self.file_path, "w") as fd:
+        """
+        Write a Message to the given file in YAML format.
+        Creates the following file structures:
+            * Question.yaml_key: single or multiline string
+            * Answer.yaml_key: single or multiline string
+            * Message.tags_yaml_key: list of strings
+            * Message.ai_yaml_key: str [Optional]
+            * Message.model_yaml_key: str [Optional]
+        """
+        with open(file_path, "w") as fd:
             data: YamlDict = {Question.yaml_key: str(self.question)}
             if self.answer:
                 data[Answer.yaml_key] = str(self.answer)

tests/test_message.py

@@ -86,19 +86,21 @@ class MessageToFileTxtTestCase(CmmTestCase):
     def setUp(self) -> None:
         self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
         self.file_path = pathlib.Path(self.file.name)
-        self.message = Message(Question('This is a question.'),
+        self.message_complete = Message(Question('This is a question.'),
                                         Answer('This is an answer.'),
                                         {Tag('tag1'), Tag('tag2')},
                                         ai='ChatGPT',
                                         model='gpt-3.5-turbo',
                                         file_path=self.file_path)
+        self.message_min = Message(Question('This is a question.'),
+                                   file_path=self.file_path)
 
     def tearDown(self) -> None:
         self.file.close()
         self.file_path.unlink()
 
-    def test_to_file_txt(self) -> None:
+    def test_to_file_txt_complete(self) -> None:
-        self.message.to_file(self.file_path)
+        self.message_complete.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -109,13 +111,23 @@ class MessageToFileTxtTestCase(CmmTestCase):
 This is a question.
 {Answer.txt_header}
 This is an answer.
+"""
+        self.assertEqual(content, expected_content)
+
+    def test_to_file_txt_min(self) -> None:
+        self.message_min.to_file(self.file_path)
+
+        with open(self.file_path, "r") as fd:
+            content = fd.read()
+        expected_content = f"""{Question.txt_header}
+This is a question.
 """
         self.assertEqual(content, expected_content)
 
     def test_to_file_unsupported_file_type(self) -> None:
         unsupported_file_path = pathlib.Path("example.doc")
         with self.assertRaises(MessageError) as cm:
-            self.message.to_file(unsupported_file_path)
+            self.message_complete.to_file(unsupported_file_path)
         self.assertEqual(str(cm.exception), "File type '.doc' is not supported")
 
     def test_to_file_no_file_path(self) -> None:
@@ -124,18 +136,18 @@ This is an answer.
         """
         with self.assertRaises(MessageError) as cm:
             # clear the internal file_path
-            self.message.file_path = None
+            self.message_complete.file_path = None
-            self.message.to_file(None)
+            self.message_complete.to_file(None)
         self.assertEqual(str(cm.exception), "Got no valid path to write message")
         # reset the internal file_path
-        self.message.file_path = self.file_path
+        self.message_complete.file_path = self.file_path
 
 
 class MessageToFileYamlTestCase(CmmTestCase):
     def setUp(self) -> None:
         self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
         self.file_path = pathlib.Path(self.file.name)
-        self.message = Message(Question('This is a question.'),
+        self.message_complete = Message(Question('This is a question.'),
                                         Answer('This is an answer.'),
                                         {Tag('tag1'), Tag('tag2')},
                                         ai='ChatGPT',
@@ -147,13 +159,15 @@ class MessageToFileYamlTestCase(CmmTestCase):
                                          ai='ChatGPT',
                                          model='gpt-3.5-turbo',
                                          file_path=self.file_path)
+        self.message_min = Message(Question('This is a question.'),
+                                   file_path=self.file_path)
 
     def tearDown(self) -> None:
         self.file.close()
         self.file_path.unlink()
 
-    def test_to_file_yaml(self) -> None:
+    def test_to_file_yaml_complete(self) -> None:
-        self.message.to_file(self.file_path)
+        self.message_complete.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -186,6 +200,14 @@ class MessageToFileYamlTestCase(CmmTestCase):
 """
         self.assertEqual(content, expected_content)
 
+    def test_to_file_yaml_min(self) -> None:
+        self.message_min.to_file(self.file_path)
+
+        with open(self.file_path, "r") as fd:
+            content = fd.read()
+        expected_content = f"{Question.yaml_key}: This is a question.\n"
+        self.assertEqual(content, expected_content)
+
 
 class MessageFromFileTxtTestCase(CmmTestCase):
     def setUp(self) -> None:
@@ -382,6 +404,12 @@ This is an answer.
   - tag2
 """)
 
+    def tearDown(self) -> None:
+        self.file_txt.close()
+        self.file_path_txt.unlink()
+        self.file_yaml.close()
+        self.file_path_yaml.unlink()
+
     def test_tags_from_file_txt(self) -> None:
         tags = Message.tags_from_file(self.file_path_txt)
         self.assertSetEqual(tags, {Tag('tag1'), Tag('tag2')})