From 8bdd7ed0ed2f3ceae354c290dab2d19d6b41ba41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Tue, 11 Feb 2025 22:14:27 -0800
Subject: [PATCH] tests: implement langchain tests with run_test()

---
 tests/test_langchain.py | 48 ++++++++++++++++++-----------------------
 1 file changed, 21 insertions(+), 27 deletions(-)

diff --git a/tests/test_langchain.py b/tests/test_langchain.py
index c94bc1c01..6534f6bb0 100644
--- a/tests/test_langchain.py
+++ b/tests/test_langchain.py
@@ -10,23 +10,22 @@ from langchain.prompts import ChatPromptTemplate
 from langchain_core.language_models import FakeStreamingListLLM
 
 from pipecat.frames.frames import (
-    EndFrame,
     LLMFullResponseEndFrame,
     LLMFullResponseStartFrame,
+    LLMMessagesFrame,
     TextFrame,
     TranscriptionFrame,
     UserStartedSpeakingFrame,
     UserStoppedSpeakingFrame,
 )
 from pipecat.pipeline.pipeline import Pipeline
-from pipecat.pipeline.runner import PipelineRunner
-from pipecat.pipeline.task import PipelineParams, PipelineTask
 from pipecat.processors.aggregators.llm_response import (
     LLMAssistantResponseAggregator,
     LLMUserResponseAggregator,
 )
 from pipecat.processors.frame_processor import FrameProcessor
 from pipecat.processors.frameworks.langchain import LangchainProcessor
+from pipecat.tests.utils import SleepFrame, run_test
 
 
 class TestLangchain(unittest.IsolatedAsyncioTestCase):
@@ -64,31 +63,26 @@ class TestLangchain(unittest.IsolatedAsyncioTestCase):
         self.mock_proc = self.MockProcessor("token_collector")
 
         tma_in = LLMUserResponseAggregator(messages)
-        tma_out = LLMAssistantResponseAggregator(messages)
+        tma_out = LLMAssistantResponseAggregator(messages, expect_stripped_words=False)
 
-        pipeline = Pipeline(
-            [
-                tma_in,
-                proc,
-                self.mock_proc,
-                tma_out,
-            ]
+        pipeline = Pipeline([tma_in, proc, self.mock_proc, tma_out])
+
+        frames_to_send = [
+            UserStartedSpeakingFrame(),
+            TranscriptionFrame(text="Hi World", user_id="user", timestamp="now"),
+            SleepFrame(),
+            UserStoppedSpeakingFrame(),
+        ]
+        expected_down_frames = [
+            UserStartedSpeakingFrame,
+            UserStoppedSpeakingFrame,
+            LLMMessagesFrame,
+        ]
+        await run_test(
+            pipeline,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_down_frames,
         )
 
-        task = PipelineTask(pipeline, PipelineParams(allow_interruptions=False))
-        await task.queue_frames(
-            [
-                UserStartedSpeakingFrame(),
-                TranscriptionFrame(text="Hi World", user_id="user", timestamp="now"),
-                UserStoppedSpeakingFrame(),
-                EndFrame(),
-            ]
-        )
-
-        runner = PipelineRunner()
-        await runner.run(task)
         self.assertEqual("".join(self.mock_proc.token), self.expected_response)
-        # TODO: Address this issue
-        # This next one would fail with:
-        # AssertionError: ' H e l l o   d e a r   h u m a n' != 'Hello dear human'
-        # self.assertEqual(tma_out.messages[-1]["content"], self.expected_response)
+        self.assertEqual(tma_out.messages[-1]["content"], self.expected_response)