63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
import unittest
|
|
import argparse
|
|
import tempfile
|
|
import yaml
|
|
from pathlib import Path
|
|
from chatmastermind.message import Message, Question
|
|
from chatmastermind.chat import ChatDB, ChatError
|
|
from chatmastermind.configuration import Config
|
|
from chatmastermind.commands.hist import convert_messages
|
|
|
|
|
|
msg_suffix = Message.file_suffix_write
|
|
|
|
|
|
class TestConvertMessages(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.db_dir = tempfile.TemporaryDirectory()
|
|
self.cache_dir = tempfile.TemporaryDirectory()
|
|
self.db_path = Path(self.db_dir.name)
|
|
self.cache_path = Path(self.cache_dir.name)
|
|
self.args = argparse.Namespace()
|
|
self.config = Config()
|
|
self.config.cache = self.cache_dir.name
|
|
self.config.db = self.db_dir.name
|
|
# Prepare some messages
|
|
self.chat = ChatDB.from_dir(Path(self.cache_path),
|
|
Path(self.db_path))
|
|
self.messages = [Message(Question(f'Question {i}')) for i in range(0, 6)]
|
|
self.chat.db_write(self.messages[0:2])
|
|
self.chat.cache_write(self.messages[2:])
|
|
# Change some of the suffixes
|
|
assert self.messages[0].file_path
|
|
assert self.messages[1].file_path
|
|
self.messages[0].file_path.rename(self.messages[0].file_path.with_suffix('.txt'))
|
|
self.messages[1].file_path.rename(self.messages[1].file_path.with_suffix('.yaml'))
|
|
|
|
def tearDown(self) -> None:
|
|
self.db_dir.cleanup()
|
|
self.cache_dir.cleanup()
|
|
|
|
def test_convert_messages(self) -> None:
|
|
self.args.convert = 'yaml'
|
|
convert_messages(self.args, self.config)
|
|
msgs = self.chat.msg_gather(loc='disk', glob='*.*')
|
|
# Check if the number of messages is the same as before
|
|
self.assertEqual(len(msgs), len(self.messages))
|
|
# Check if all messages have the requested suffix
|
|
for msg in msgs:
|
|
assert msg.file_path
|
|
self.assertEqual(msg.file_path.suffix, msg_suffix)
|
|
# Check if the message IDs are correctly maintained
|
|
for m_new, m_old in zip(msgs, self.messages):
|
|
self.assertEqual(m_new.msg_id(), m_old.msg_id())
|
|
# check if all messages have the new format
|
|
for m in msgs:
|
|
with open(str(m.file_path), "r") as fd:
|
|
yaml.load(fd, Loader=yaml.FullLoader)
|
|
|
|
def test_convert_messages_wrong_format(self) -> None:
|
|
self.args.convert = 'foo'
|
|
with self.assertRaises(ChatError):
|
|
convert_messages(self.args, self.config)
|