wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tests pass

Browse Source
This commit is contained in:
Moishe Lettvin
2023-12-26 14:13:10 -05:00
parent e724720e76
commit df536b0ad0
3 changed files with 38 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
tests/test_message_handler.py
View File

@@ -4,7 +4,7 @@ import unittest
from unittest.mock import MagicMock, call
from message_handler.message_handler import MessageHandler, IndexingMessageHandler
from services.ai_services import AIService, AIServiceConfig
from services.ai_services import AIService, AIServiceConfig, TTSService, LLMService, ImageGenService
from storage.search import SearchIndexer
@@ -44,7 +44,7 @@ class TestMessageHandler(unittest.TestCase):
message_handler = MessageHandler("System prompt")
message_handler.add_user_message("User message")
message_handler.add_assistant_message("Assistant message")
message_handler.add_user_message("User message plus something else")
message_handler.add_user_message("plus something else")
self.assertEqual(
message_handler.get_llm_messages(),
[
@@ -57,6 +57,7 @@ class TestMessageHandler(unittest.TestCase):
message_handler = MessageHandler("System prompt")
message_handler.add_user_message("User message")
message_handler.add_assistant_message("Assistant message")
message_handler.finalize_user_message()
message_handler.add_user_message("other user message")
self.assertEqual(
message_handler.get_llm_messages(),
@@ -69,25 +70,36 @@ class TestMessageHandler(unittest.TestCase):
)
class MockAIService(AIService):
def __init__(self, **kwargs):
super().__init__(**kwargs)
class MockTTSService(TTSService):
def run_tts(self, sentence):
for word in sentence.split(" "):
time.sleep(0.1)
yield bytes(word, "utf-8")
def run_llm(self, messages, latest_user_message=None, stream=True):
return {"choices": [{"message": {"content": "Parsed user message."}}]}
class MockLLMService(LLMService):
def run_llm(self, messages) -> str:
return "Parsed user message."
class MockImageService(ImageGenService):
def run_image_gen(self, sentence) -> None:
return None
class TestIndexingMessageHandler(unittest.TestCase):
def test_user_message_finalized(self):
mock_ai_service = MockAIService()
mock_tts_service = MockTTSService()
mock_llm_service = MockLLMService()
mock_image_service = MockImageService()
service_config = AIServiceConfig(
mock_ai_service, mock_ai_service, mock_ai_service
tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service
)
mock_indexer = MagicMock(spec=SearchIndexer)
message_handler = IndexingMessageHandler(
"Hello world", "story_id", service_config, mock_indexer
"Hello world", service_config, mock_indexer
)
message_handler.add_user_message("User message")
message_handler.add_assistant_message("Assistant message will be ignored")
@@ -101,11 +113,10 @@ class TestIndexingMessageHandler(unittest.TestCase):
message_handler.write_messages_to_index()
time.sleep(0.5)
self.assertEqual(
mock_indexer.mock_calls,
[
call.index_text("Parsed user message."),
call.index_text('"Parsed user message."'),
call.index_text("New assistant message will not be ignored"),
],
)
@@ -119,7 +130,7 @@ class TestIndexingMessageHandler(unittest.TestCase):
self.assertEqual(
mock_indexer.mock_calls,
[
call.index_text("Parsed user message."),
call.index_text('"Parsed user message."'),
call.index_text("Assistant message second time"),
],
)