From d080a31a5ca322364f5645aafa3727d47ea32f64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Mon, 30 Sep 2024 11:37:26 -0700
Subject: [PATCH] tests: fix langchanin tests

---
 tests/test_langchain.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/test_langchain.py b/tests/test_langchain.py
index fb222205b..d30d213bd 100644
--- a/tests/test_langchain.py
+++ b/tests/test_langchain.py
@@ -7,9 +7,9 @@
 import unittest
 
 from pipecat.frames.frames import (
+    EndFrame,
     LLMFullResponseEndFrame,
     LLMFullResponseStartFrame,
-    StopTaskFrame,
     TextFrame,
     TranscriptionFrame,
     UserStartedSpeakingFrame,
@@ -32,6 +32,7 @@ from langchain_core.language_models import FakeStreamingListLLM
 class TestLangchain(unittest.IsolatedAsyncioTestCase):
     class MockProcessor(FrameProcessor):
         def __init__(self, name):
+            super().__init__()
             self.name = name
             self.token: list[str] = []
             # Start collecting tokens when we see the start frame
@@ -55,13 +56,13 @@ class TestLangchain(unittest.IsolatedAsyncioTestCase):
     def setUp(self):
         self.expected_response = "Hello dear human"
         self.fake_llm = FakeStreamingListLLM(responses=[self.expected_response])
-        self.mock_proc = self.MockProcessor("token_collector")
 
     async def test_langchain(self):
         messages = [("system", "Say hello to {name}"), ("human", "{input}")]
         prompt = ChatPromptTemplate.from_messages(messages).partial(name="Thomas")
         chain = prompt | self.fake_llm
         proc = LangchainProcessor(chain=chain)
+        self.mock_proc = self.MockProcessor("token_collector")
 
         tma_in = LLMUserResponseAggregator(messages)
         tma_out = LLMAssistantResponseAggregator(messages)
@@ -81,7 +82,7 @@ class TestLangchain(unittest.IsolatedAsyncioTestCase):
                 UserStartedSpeakingFrame(),
                 TranscriptionFrame(text="Hi World", user_id="user", timestamp="now"),
                 UserStoppedSpeakingFrame(),
-                StopTaskFrame(),
+                EndFrame(),
             ]
         )