From 3e9e2aa35fffe66148ee735b1742cf04cd461ce6 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Wed, 26 Nov 2025 17:56:15 +0000
Subject: [PATCH] Fixing some tests

---
 tests/integration/test_agent_streaming_integration.py    | 5 ++---
 .../test_document_rag_streaming_integration.py           | 9 +++++----
 .../integration/test_graph_rag_streaming_integration.py  | 9 +++++----
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/tests/integration/test_agent_streaming_integration.py b/tests/integration/test_agent_streaming_integration.py
index b119eeb6..4a7d46e7 100644
--- a/tests/integration/test_agent_streaming_integration.py
+++ b/tests/integration/test_agent_streaming_integration.py
@@ -99,12 +99,11 @@ Args: {
         }
 
     @pytest.fixture
-    def agent_manager(self, mock_flow_context, sample_tools):
+    def agent_manager(self, sample_tools):
         """Create AgentManager instance with streaming support"""
         return AgentManager(
-            prompt_client=mock_flow_context("prompt-request"),
             tools=sample_tools,
-            verbose=True
+            additional_context="You are a helpful AI assistant."
         )
 
     @pytest.mark.asyncio
diff --git a/tests/integration/test_document_rag_streaming_integration.py b/tests/integration/test_document_rag_streaming_integration.py
index 45601b39..4b792443 100644
--- a/tests/integration/test_document_rag_streaming_integration.py
+++ b/tests/integration/test_document_rag_streaming_integration.py
@@ -42,16 +42,17 @@ class TestDocumentRagStreaming:
         client = AsyncMock()
 
         async def document_prompt_side_effect(query, documents, timeout=600, streaming=False, chunk_callback=None):
+            # Both modes return the same text
+            full_text = "Machine learning is a subset of artificial intelligence that focuses on algorithms that learn from data."
+
             if streaming and chunk_callback:
                 # Simulate streaming chunks
-                full_text = ""
                 async for chunk in mock_streaming_llm_response():
-                    full_text += chunk
                     await chunk_callback(chunk)
                 return full_text
             else:
-                # Non-streaming response
-                return "Machine learning is a subset of artificial intelligence."
+                # Non-streaming response - same text
+                return full_text
 
         client.document_prompt.side_effect = document_prompt_side_effect
         return client
diff --git a/tests/integration/test_graph_rag_streaming_integration.py b/tests/integration/test_graph_rag_streaming_integration.py
index d12ec232..92da6527 100644
--- a/tests/integration/test_graph_rag_streaming_integration.py
+++ b/tests/integration/test_graph_rag_streaming_integration.py
@@ -55,16 +55,17 @@ class TestGraphRagStreaming:
         client = AsyncMock()
 
         async def kg_prompt_side_effect(query, kg, timeout=600, streaming=False, chunk_callback=None):
+            # Both modes return the same text
+            full_text = "Machine learning is a subset of artificial intelligence that focuses on algorithms that learn from data."
+
             if streaming and chunk_callback:
                 # Simulate streaming chunks
-                full_text = ""
                 async for chunk in mock_streaming_llm_response():
-                    full_text += chunk
                     await chunk_callback(chunk)
                 return full_text
             else:
-                # Non-streaming response
-                return "Machine learning is a subset of artificial intelligence."
+                # Non-streaming response - same text
+                return full_text
 
         client.kg_prompt.side_effect = kg_prompt_side_effect
         return client