Change storage to use text format per default, instead of yaml, but still support yaml.

This commit is contained in:
Oleksandr Kozachuk
2023-05-07 14:31:17 +02:00
parent 57caba5360
commit c5fd466dda
2 changed files with 57 additions and 40 deletions
+10 -12
View File
@@ -1,12 +1,11 @@
import unittest
import io
import pathlib
import yaml
import argparse
from chatmastermind.utils import terminal_width
from chatmastermind.main import create_parser, handle_question
from chatmastermind.api_client import ai
from chatmastermind.storage import create_chat, save_answers
from chatmastermind.storage import create_chat, save_answers, dump_data
from unittest import mock
from unittest.mock import patch, MagicMock, Mock
@@ -24,8 +23,8 @@ class TestCreateChat(unittest.TestCase):
@patch('os.listdir')
@patch('builtins.open')
def test_create_chat_with_tags(self, open_mock, listdir_mock):
listdir_mock.return_value = ['testfile.yaml']
open_mock.return_value.__enter__.return_value = io.StringIO(yaml.dump(
listdir_mock.return_value = ['testfile.txt']
open_mock.return_value.__enter__.return_value = io.StringIO(dump_data(
{'question': 'test_content', 'answer': 'some answer',
'tags': ['test_tag']}))
@@ -44,8 +43,8 @@ class TestCreateChat(unittest.TestCase):
@patch('os.listdir')
@patch('builtins.open')
def test_create_chat_with_other_tags(self, open_mock, listdir_mock):
listdir_mock.return_value = ['testfile.yaml']
open_mock.return_value.__enter__.return_value = io.StringIO(yaml.dump(
listdir_mock.return_value = ['testfile.txt']
open_mock.return_value.__enter__.return_value = io.StringIO(dump_data(
{'question': 'test_content', 'answer': 'some answer',
'tags': ['other_tag']}))
@@ -60,12 +59,12 @@ class TestCreateChat(unittest.TestCase):
@patch('os.listdir')
@patch('builtins.open')
def test_create_chat_without_tags(self, open_mock, listdir_mock):
listdir_mock.return_value = ['testfile.yaml', 'testfile2.yaml']
listdir_mock.return_value = ['testfile.txt', 'testfile2.txt']
open_mock.side_effect = (
io.StringIO(yaml.dump({'question': 'test_content',
io.StringIO(dump_data({'question': 'test_content',
'answer': 'some answer',
'tags': ['test_tag']})),
io.StringIO(yaml.dump({'question': 'test_content2',
io.StringIO(dump_data({'question': 'test_content2',
'answer': 'some answer2',
'tags': ['test_tag2']})),
)
@@ -109,8 +108,7 @@ class TestHandleQuestion(unittest.TestCase):
@patch("chatmastermind.main.ai", return_value=(["answer1", "answer2", "answer3"], "test_usage"))
@patch("chatmastermind.utils.pp")
@patch("builtins.print")
@patch("chatmastermind.storage.yaml.dump")
def test_handle_question(self, _, mock_print, mock_pp, mock_ai,
def test_handle_question(self, mock_print, mock_pp, mock_ai,
mock_process_tags, mock_create_chat):
open_mock = MagicMock()
with patch("chatmastermind.storage.open", open_mock):
@@ -135,7 +133,7 @@ class TestHandleQuestion(unittest.TestCase):
expected_calls.append((("-" * terminal_width(),),))
expected_calls.append(((f"Usage: {mock_ai.return_value[1]}",),))
self.assertEqual(mock_print.call_args_list, expected_calls)
open_expected_calls = list([mock.call(f"{num:04d}.yaml", "w") for num in range(2, 5)])
open_expected_calls = list([mock.call(f"{num:04d}.txt", "w") for num in range(2, 5)])
open_mock.assert_has_calls(open_expected_calls, any_order=True)