From 4f88e589b0a46410b3dbcfbc8fad5b063b67c118 Mon Sep 17 00:00:00 2001
From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com>
Date: Fri, 8 May 2026 12:07:13 +0530
Subject: [PATCH] fix(tests): poll for on_timeout middleware in prefork test

handle_result() flips the DB status to 'dead' before dispatch_outcome
fires on_timeout, so a fast test thread can observe 'dead' and assert on
timeouts_seen before the worker reaches the middleware call. Poll for
the spy with a small budget instead.
---
 tests/worker/test_prefork.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/worker/test_prefork.py b/tests/worker/test_prefork.py
index fa1a366..fa8915e 100644
--- a/tests/worker/test_prefork.py
+++ b/tests/worker/test_prefork.py
@@ -161,7 +161,7 @@ def _wait_for_terminal(job: Any, timeout: float) -> str:
 
 
 @prefork_unix_only
-def test_prefork_kills_hung_task(timeout_app: object) -> None:
+def test_prefork_kills_hung_task(timeout_app: object, poll_until: Any) -> None:
     """A task that hangs past its `timeout=` is SIGKILLed by the watchdog and
     reported as a timeout failure within the timeout + watchdog tick budget."""
     timeouts_seen: list[str] = []
@@ -185,7 +185,14 @@ def on_timeout(self, ctx: JobContext) -> None:
     assert status == "dead", f"expected 'dead', got {status!r} (error={job.error!r})"
     assert "timed out" in (job.error or "").lower()
     assert elapsed < 12, f"hung task took {elapsed:.1f}s to be killed (expected < 12s)"
-    assert job.id in timeouts_seen, "on_timeout middleware did not fire"
+    # The DB status flips to 'dead' inside `handle_result()`, but `on_timeout`
+    # fires a moment later in `dispatch_outcome` on the worker thread — poll
+    # rather than asserting on a single observation to avoid that race.
+    poll_until(
+        lambda: job.id in timeouts_seen,
+        timeout=5,
+        message="on_timeout middleware did not fire",
+    )
 
 
 @prefork_unix_only