Fixing tests

This commit is contained in:
Cyber MacGeddon 2025-11-26 19:07:00 +00:00
parent 6c0147feda
commit c14116f7a3

View file

@ -28,34 +28,35 @@ class TestAgentStreaming:
client = AsyncMock() client = AsyncMock()
async def agent_react_streaming(variables, timeout=600, streaming=False, chunk_callback=None): async def agent_react_streaming(variables, timeout=600, streaming=False, chunk_callback=None):
if streaming and chunk_callback: # Both modes return the same text for equivalence
# Simulate streaming response with chunks full_text = """Thought: I need to search for information about machine learning.
chunks = [
"Thought: I need to",
" search for",
" information about",
" machine learning.",
"\n",
"Action: knowledge_query\n",
"Args: {\n",
' "question": "What is machine learning?"\n',
"}"
]
full_text = ""
for chunk in chunks:
full_text += chunk
await chunk_callback(chunk)
return full_text
else:
# Non-streaming response
return """Thought: I need to search for information about machine learning.
Action: knowledge_query Action: knowledge_query
Args: { Args: {
"question": "What is machine learning?" "question": "What is machine learning?"
}""" }"""
if streaming and chunk_callback:
# Simulate streaming response with complete lines as chunks
# This mimics how real LLMs stream - complete tokens/words
chunks = [
"Thought: I need to search",
" for information about",
" machine learning.\n",
"Action: knowledge_query\n",
"Args: {\n",
' "question": ',
'"What is machine learning?"',
'\n}'
]
for chunk in chunks:
await chunk_callback(chunk)
return full_text
else:
# Non-streaming response - same text
return full_text
client.agent_react.side_effect = agent_react_streaming client.agent_react.side_effect = agent_react_streaming
return client return client
@ -240,11 +241,19 @@ Args: {
assert result is not None assert result is not None
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_agent_streaming_with_conversation_history(self, agent_manager, mock_flow_context, async def test_agent_streaming_with_conversation_history(self, agent_manager, mock_flow_context):
sample_agent_responses):
"""Test streaming with existing conversation history""" """Test streaming with existing conversation history"""
# Arrange # Arrange
history = sample_agent_responses[:1] # Use first response as history # History should be a list of Action objects
from trustgraph.agent.react.types import Action
history = [
Action(
thought="I need to search for information about machine learning",
name="knowledge_query",
arguments={"question": "What is machine learning?"},
observation="Machine learning is a subset of AI that enables computers to learn from data."
)
]
think = AsyncMock() think = AsyncMock()
# Act # Act