Update to add streaming tests (#600)

This commit is contained in:
cybermaggedon 2026-01-06 21:48:05 +00:00 committed by GitHub
parent f0c95a4c5e
commit f79d0603f7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1062 additions and 57 deletions

View file

@ -59,9 +59,16 @@ class TestGraphRagStreaming:
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
# Simulate streaming chunks with end_of_stream flags
chunks = []
async for chunk in mock_streaming_llm_response():
await chunk_callback(chunk)
chunks.append(chunk)
# Send all chunks with end_of_stream=False except the last
for i, chunk in enumerate(chunks):
is_final = (i == len(chunks) - 1)
await chunk_callback(chunk, is_final)
return full_text
else:
# Non-streaming response - same text
@ -102,6 +109,9 @@ class TestGraphRagStreaming:
assert_streaming_chunks_valid(collector.chunks, min_chunks=1)
assert_callback_invoked(AsyncMock(call_count=len(collector.chunks)), min_calls=1)
# Verify streaming protocol compliance
collector.verify_streaming_protocol()
# Verify full response matches concatenated chunks
full_from_chunks = collector.get_full_text()
assert result == full_from_chunks
@ -128,7 +138,7 @@ class TestGraphRagStreaming:
# Act - Streaming
streaming_chunks = []
async def collect(chunk):
async def collect(chunk, end_of_stream):
streaming_chunks.append(chunk)
streaming_result = await graph_rag_streaming.query(