From 20ca2d5759defc3795a0d6c75af2639eedc3d2ad Mon Sep 17 00:00:00 2001
From: teknium1 <127238744+teknium1@users.noreply.github.com>
Date: Tue, 30 Jun 2026 16:42:59 -0700
Subject: [PATCH] test(mcp-oauth): yield for done-callback before asserting
 task cleanup

The discard done-callback added via task.add_done_callback runs on a later
event-loop iteration (call_soon) than the one that resolves `pending` and lets
handle_401 return. Both inflight-task tests asserted the live set was empty
immediately after the await returned, racing the callback. Add a single
`await asyncio.sleep(0)` before the cleanup assertions.
---
 tests/tools/test_mcp_oauth_manager.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/tools/test_mcp_oauth_manager.py b/tests/tools/test_mcp_oauth_manager.py
index 2e7d3aa41..448400cad 100644
--- a/tests/tools/test_mcp_oauth_manager.py
+++ b/tests/tools/test_mcp_oauth_manager.py
@@ -174,6 +174,12 @@ async def test_handle_401_tracks_inflight_task_to_prevent_gc(tmp_path, monkeypat
 
     result = await mgr.handle_401("srv", failed_access_token="TOK")
 
+    # The discard done-callback is scheduled via loop.call_soon, so it runs on
+    # a later loop iteration than the one that resolved `pending` and let
+    # handle_401 return. Yield once so the callback fires before we assert the
+    # task was removed from the live set.
+    await asyncio.sleep(0)
+
     # Exactly one handler task was created and tracked.
     assert len(mgr._inflight_tasks.ever_added) == 1
     tracked_task = mgr._inflight_tasks.ever_added[0]
@@ -226,6 +232,8 @@ async def test_handle_401_dedup_survives_even_if_task_reference_dropped(tmp_path
 
     results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0)
     assert results == [False] * 8
+    # Let the shared _do_handle task's discard done-callback (call_soon) run.
+    await asyncio.sleep(0)
     assert len(mgr._inflight_tasks) == 0