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] = None
-    tags: Optional[set[Tag]] = None
-    ai: Optional[str] = None
-    model: Optional[str] = None
-    file_path: Optional[pathlib.Path] = None
+    answer: Optional[Answer]
+    tags: Optional[set[Tag]]
+    ai: Optional[str]
+    model: Optional[str]
+    file_path: Optional[pathlib.Path]
     # class variables
     file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml']
     tags_yaml_key: ClassVar[str] = 'tags'
@@ -213,31 +213,12 @@ class Message():
         return tags
 
     @classmethod
-    def from_file(cls: Type[MessageInst], file_path: pathlib.Path,
+    def from_file(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 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:
+        Create a Message from the given file. Expects the following file structures:
         For '.txt':
             * TagLine [Optional]
             * AI [Optional]
@@ -245,15 +226,27 @@ 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:
@@ -287,23 +280,8 @@ class Message():
                 except ValueError:
                     question = Question.from_list(text[question_idx:])
                 return cls(question, answer, tags, ai, model, file_path)
-
-    @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()
+        # YAML
+        else:
             with open(file_path, "r") as fd:
                 data = yaml.load(fd, Loader=yaml.FullLoader)
                 if tags_or or tags_and or tags_not:
@@ -317,8 +295,21 @@ class Message():
 
     def to_file(self, file_path: Optional[pathlib.Path]) -> None:  # noqa: 11
         """
-        Write a Message to the given file. Type is determined based on the suffix.
-        Currently supported suffixes: ['.txt', '.yaml']
+        Write a Message to the given file. Creates the following file structures:
+        For '.txt':
+            * 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
@@ -328,44 +319,18 @@ class Message():
             raise MessageError(f"File type '{self.file_path.suffix}' is not supported")
         # TXT
         if self.file_path.suffix == '.txt':
-            return self.__to_file_txt(self.file_path)
-        elif self.file_path.suffix == '.yaml':
-            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')
+            with open(self.file_path, "w") as fd:
+                msg_tags = self.tags or set()
+                fd.write(f'{TagLine.from_set(msg_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')
-
-    def __to_file_yaml(self, file_path: pathlib.Path) -> None:
-        """
-        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:
+        # YAML
+        elif self.file_path.suffix == '.yaml':
+            with open(self.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,21 +86,19 @@ 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_complete = Message(Question('This is a question.'),
+        self.message = 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_complete(self) -> None:
-        self.message_complete.to_file(self.file_path)
+    def test_to_file_txt(self) -> None:
+        self.message.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -111,23 +109,13 @@ 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_complete.to_file(unsupported_file_path)
+            self.message.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:
@@ -136,18 +124,18 @@ This is a question.
         """
         with self.assertRaises(MessageError) as cm:
             # clear the internal file_path
-            self.message_complete.file_path = None
-            self.message_complete.to_file(None)
+            self.message.file_path = None
+            self.message.to_file(None)
         self.assertEqual(str(cm.exception), "Got no valid path to write message")
         # reset the internal file_path
-        self.message_complete.file_path = self.file_path
+        self.message.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_complete = Message(Question('This is a question.'),
+        self.message = Message(Question('This is a question.'),
                                Answer('This is an answer.'),
                                {Tag('tag1'), Tag('tag2')},
                                ai='ChatGPT',
@@ -159,15 +147,13 @@ 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_complete(self) -> None:
-        self.message_complete.to_file(self.file_path)
+    def test_to_file_yaml(self) -> None:
+        self.message.to_file(self.file_path)
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
@@ -200,14 +186,6 @@ 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:
@@ -404,12 +382,6 @@ 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')})