apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-25 13:11:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fixing some tests

Browse source
This commit is contained in:
Cyber MacGeddon 2025-11-26 17:56:15 +00:00
parent 65e297266f
commit 3e9e2aa35f
3 changed files with 12 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

5
tests/integration/test_agent_streaming_integration.py
View file

@ -99,12 +99,11 @@ Args: {
} }
@pytest.fixture @pytest.fixture
def agent_manager(self, mock_flow_context, sample_tools): def agent_manager(self, sample_tools):
"""Create AgentManager instance with streaming support""" """Create AgentManager instance with streaming support"""
return AgentManager( return AgentManager(
prompt_client=mock_flow_context("prompt-request"),
tools=sample_tools, tools=sample_tools,
verbose=True additional_context="You are a helpful AI assistant."
) )
@pytest.mark.asyncio @pytest.mark.asyncio

9
tests/integration/test_document_rag_streaming_integration.py
View file

@ -42,16 +42,17 @@ class TestDocumentRagStreaming:
client = AsyncMock() client = AsyncMock()
async def document_prompt_side_effect(query, documents, timeout=600, streaming=False, chunk_callback=None): 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: if streaming and chunk_callback:
# Simulate streaming chunks # Simulate streaming chunks
full_text = ""
async for chunk in mock_streaming_llm_response(): async for chunk in mock_streaming_llm_response():
full_text += chunk
await chunk_callback(chunk) await chunk_callback(chunk)
return full_text return full_text
else: else:
# Non-streaming response # Non-streaming response - same text
return "Machine learning is a subset of artificial intelligence." return full_text
client.document_prompt.side_effect = document_prompt_side_effect client.document_prompt.side_effect = document_prompt_side_effect
return client return client

9
tests/integration/test_graph_rag_streaming_integration.py
View file

@ -55,16 +55,17 @@ class TestGraphRagStreaming:
client = AsyncMock() client = AsyncMock()
async def kg_prompt_side_effect(query, kg, timeout=600, streaming=False, chunk_callback=None): 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: if streaming and chunk_callback:
# Simulate streaming chunks # Simulate streaming chunks
full_text = ""
async for chunk in mock_streaming_llm_response(): async for chunk in mock_streaming_llm_response():
full_text += chunk
await chunk_callback(chunk) await chunk_callback(chunk)
return full_text return full_text
else: else:
# Non-streaming response # Non-streaming response - same text
return "Machine learning is a subset of artificial intelligence." return full_text
client.kg_prompt.side_effect = kg_prompt_side_effect client.kg_prompt.side_effect = kg_prompt_side_effect
return client return client