From 5633764733cdb64f5e72964d6f2d654eb9b50906 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Mon, 3 Aug 2026 10:48:46 +0530
Subject: [PATCH] test: wait for ledger-wrapped sends in drain assertions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two consumer tests assumed a send completes synchronously within the
handler turn: the continuation-drain test polled on handler-call count
then immediately asserted on adapter.sent, and the split-brain heal
test drained with bare zero-delay yields. With ledger calls hopping to
worker threads around each send, the reply can land microseconds after
those checks. Poll for the actual sends with a bounded 2s window —
same invariants, scheduling-robust.
---
 tests/gateway/test_goal_continuation_drain.py   |  5 ++++-
 tests/gateway/test_session_split_brain_11016.py | 10 +++++++---
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/tests/gateway/test_goal_continuation_drain.py b/tests/gateway/test_goal_continuation_drain.py
index 58c111dde6467..662f541fc6266 100644
--- a/tests/gateway/test_goal_continuation_drain.py
+++ b/tests/gateway/test_goal_continuation_drain.py
@@ -120,7 +120,10 @@ async def test_fifo_enqueued_continuation_is_drained_without_new_user_message():
     await adapter._process_message_background(event, key)
     # The in-band drain hands off to a fresh task (#17758); let it run.
     for _ in range(40):
-        if len(handled) >= 2:
+        # Wait for the SENDS, not just the handler calls: the delivery
+        # ledger hops to worker threads around each send, so the handler
+        # can return while reply-2's send is still in flight.
+        if len(adapter.sent) >= 2:
             break
         await asyncio.sleep(0.05)
 
diff --git a/tests/gateway/test_session_split_brain_11016.py b/tests/gateway/test_session_split_brain_11016.py
index 5a5439f4bd585..b6a850244edbb 100644
--- a/tests/gateway/test_session_split_brain_11016.py
+++ b/tests/gateway/test_session_split_brain_11016.py
@@ -198,9 +198,13 @@ class TestStaleSessionLockSelfHeal:
         # An ordinary message should heal the stale lock, then fall through
         # to normal dispatch.  User gets a reply instead of a busy ack.
         await adapter.handle_message(_make_event("hello"))
-        # Drain any spawned background tasks.
-        for _ in range(5):
-            await asyncio.sleep(0)
+        # Drain any spawned background tasks. Real sleeps, not bare yields:
+        # the delivery ledger hops to worker threads around the send, so a
+        # zero-delay yield loop can finish before the reply lands.
+        for _ in range(40):
+            if any("handled:text" in r for r in adapter.sent_responses):
+                break
+            await asyncio.sleep(0.05)
 
         assert any("handled:text" in r for r in adapter.sent_responses), (
             "stale lock trapped a normal message — split-brain not healed"