From f5832f81ed4ec62c20cd92fd05a68a495bd2e112 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Sat, 5 Sep 2026 10:48:52 +0530
Subject: [PATCH] test(streaming): drive the Relay finalizer through its
 collector
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 tests/run_agent/test_streaming.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/run_agent/test_streaming.py b/tests/run_agent/test_streaming.py
index 654db7e164..b82e9afd85 100644
--- a/tests/run_agent/test_streaming.py
+++ b/tests/run_agent/test_streaming.py
@@ -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