CharacterChatBot / tests /test_suggestions.py
cereeenn120's picture
Improve suggestions with adding persona based system
31404a1
Raw History Blame Contribute Delete
2.76 kB
import json
import unittest
from unittest.mock import patch
from engine.suggestions import generate_suggestions, fallback_suggestions, _parse
from test_conversations import Client, answer
class SuggestionTests(unittest.TestCase):
def test_language_and_persona_reach_single_model_call(self):
client = Client()
items = fallback_suggestions('Turkish')
with patch.object(client, 'chat_completion', return_value=answer(json.dumps(items))) as completion:
result = generate_suggestions({'name': 'Hermione'}, [], client, 'test',
language='Turkish', persona={'name': 'Deniz', 'background': 'A shy student'})
self.assertEqual(items, result)
completion.assert_called_once()
system = completion.call_args.args[0][0]['content']
self.assertIn('Turkish', system)
self.assertIn('Deniz', system)
self.assertIn('A shy student', system)
self.assertIn('Write as this user persona', system)
def test_empty_persona_omits_role_data_and_english_is_explicit(self):
client = Client()
with patch.object(client, 'chat_completion', return_value=answer(json.dumps(fallback_suggestions('English')))) as completion:
generate_suggestions({}, [], client, persona={})
system = completion.call_args.args[0][0]['content']
self.assertNotIn('USER PERSONA', system)
self.assertIn('in English', system)
def test_failure_uses_selected_language_without_translation_call(self):
client = Client(fail=True)
result = generate_suggestions({}, [], client, language='Turkish', persona={'name': 'Deniz'})
self.assertEqual(fallback_suggestions('Turkish'), result)
self.assertEqual(1, len(client.calls))
self.assertIn('konuya', result[0]['text'])
def test_each_ui_language_has_four_localized_fallbacks(self):
languages = ['English', 'French', 'German', 'Italian', 'Russian', 'Turkish', 'Chinese']
for language in languages:
items = fallback_suggestions(language)
self.assertEqual(4, len(items))
self.assertEqual(items, _parse(json.dumps(items)))
if language != 'English':
self.assertNotEqual(fallback_suggestions('English')[0]['text'], items[0]['text'])
changed = fallback_suggestions('English')
changed[0]['text'] = 'Modified'
self.assertNotEqual('Modified', fallback_suggestions('English')[0]['text'])
def test_malformed_output_rejected(self):
for invalid in [[], [None] * 4, [{'tone': 'x', 'emoji': 'x', 'text': ''}] * 4]:
with self.subTest(invalid=invalid), self.assertRaises(ValueError):
_parse(json.dumps(invalid))