fix: test

PostHog · Oct 28, 2024 · 067db1f · 067db1f
1 parent e360d47
commit 067db1f
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/ee/hogai/test/test_assistant.py b/ee/hogai/test/test_assistant.py
@@ -1,6 +1,7 @@
 import json
 from unittest.mock import patch
 
+import pytest
 from django.test import override_settings
 from langchain_core.runnables import RunnableLambda
 
@@ -15,6 +16,7 @@
 
 @override_settings(IN_UNIT_TESTING=True)
 class TestAssistant(NonAtomicBaseTest):
+    @pytest.mark.django_db(transaction=True)
     def test_assistant(self):
         mocked_planner_response = """
         Action:
@@ -37,6 +39,7 @@ def test_assistant(self):
             generator = assistant.stream(
                 Conversation(messages=[HumanMessage(content="Launch the chain.")], session_id="id")
             )
+            self.assertEqual(next(generator), "")
             self.assertEqual(
                 json.loads(next(generator)),
                 VisualizationMessage(answer=None, reasoning_steps=[], plan="Plan").model_dump(),