apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 20:51: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
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

9
tests/integration/test_document_rag_streaming_integration.py
View file

@ -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

9
tests/integration/test_graph_rag_streaming_integration.py
View file

@ -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