test(streaming): drive the Relay finalizer through its collector

The chat_completions finalizer now reads collector-observed chunks, not the
consumer loop, so the test feeds the captured on_chunk before calling it —
the same ordering Relay guarantees.
This commit is contained in:
kshitijk4poor
2026-09-05 10:48:52 +05:30
committed by kshitij
parent 74de0fd4fe
commit f5832f81ed

View File

@@ -368,7 +368,7 @@ class TestStreamingAccumulator:
captured = {}
fake_stream = MagicMock()
fake_stream.final_response = None
fake_stream.__iter__.return_value = iter([
chunks = [
_make_stream_chunk(tool_calls=[
_make_tool_call_delta(
index=0,
@@ -381,10 +381,12 @@ class TestStreamingAccumulator:
_make_tool_call_delta(index=0, arguments='"hello"}')
]),
_make_stream_chunk(finish_reason="tool_calls"),
])
]
fake_stream.__iter__.return_value = iter(chunks)
def relay_stream_impl(*args, **kwargs):
captured["finalizer"] = kwargs["finalizer"]
captured["on_chunk"] = kwargs["on_chunk"]
return fake_stream
mock_relay_stream.side_effect = relay_stream_impl
@@ -404,6 +406,10 @@ class TestStreamingAccumulator:
agent._interruptible_streaming_api_call({})
# Relay's contract: the collector sees every chunk as JSON, then the finalizer runs.
from agent.relay_llm import _jsonable
for chunk in chunks:
captured["on_chunk"](_jsonable(chunk))
payload = captured["finalizer"]()
tool_calls = payload["choices"][0]["message"]["tool_calls"]
assert len(tool_calls) == 1