diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index 1201a156c00..fe8e8d125fb 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -28,6 +28,7 @@ class Reply: chunks: tuple[bytes, ...] | None = None abort_after: int | None = None gate_after_first: threading.Event | None = None + gate_timeout_seconds: float = 5.0 pause_between_chunks: float = 0 headers: Mapping[str, str] = MappingProxyType({}) @@ -104,7 +105,9 @@ def wire_server( self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) self.wfile.flush() if index == 0 and reply.gate_after_first is not None: - assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released" + assert reply.gate_after_first.wait(timeout=reply.gate_timeout_seconds), ( + "Stream barrier was never released" + ) if reply.pause_between_chunks and index + 1 < len(reply.chunks): time.sleep(reply.pause_between_chunks) else: diff --git a/tests/integration/streaming/test_stream_fallback_response_cost.py b/tests/integration/streaming/test_stream_fallback_response_cost.py index e306cf386b6..77f80376fd7 100644 --- a/tests/integration/streaming/test_stream_fallback_response_cost.py +++ b/tests/integration/streaming/test_stream_fallback_response_cost.py @@ -1394,6 +1394,7 @@ def test_c3_surviving_proxy_worker_serves_after_one_worker_is_killed(gateway: Ga content_type="text/event-stream", chunks=_chat_chunks(_request_identity(request_body)), gate_after_first=backup_gate, + gate_timeout_seconds=60, ) with (