From a31fe03a71df1ab7aaa48c6eab5cfe545143a964 Mon Sep 17 00:00:00 2001 From: Tin-Yin Lai Date: Thu, 9 Jul 2026 12:03:32 -0700 Subject: [PATCH] test: cover per-phase drain timeout in BenchmarkSession The per-phase drain timeout feature (settings.drain, PhaseConfig.drain_timeout) landed via #330 and #337 but had no session-level regression test. Add one: a phase whose samples never complete must exit the drain wait after drain_timeout seconds instead of hanging. --- .../unit/load_generator/test_async_session.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 278589c1..225caa81 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -446,6 +446,30 @@ class _StuckIssuer: ) assert session._current_phase_stopped is True + @pytest.mark.asyncio + async def test_drain_respects_per_phase_timeout(self, caplog): + """PhaseConfig.drain_timeout bounds the drain wait for stuck samples.""" + loop = asyncio.get_running_loop() + publisher = FakePublisher() + + issuer = FakeIssuer() + issuer._loop = loop + issuer._auto_respond = False # samples issue but never complete + + session = BenchmarkSession(issuer, publisher, loop) + phases = [ + PhaseConfig( + "perf", + _make_settings(n_samples=2), + FakeDataset(2), + drain_timeout=0.1, + ), + ] + + result = await asyncio.wait_for(session.run(phases), timeout=10.0) + assert result is not None + assert "Drain timed out after" in caplog.text + @pytest.mark.asyncio async def test_on_sample_complete_callback(self): loop = asyncio.get_running_loop()