Fix flaky client tests by waiting for first request's reply (#1091)

* Fix flaky client tests by waiting for first request's reply

The tests were sending two requests back-to-back:
1. First call without reply=True (to check msg_id is returned)
2. Second call with reply=True (to check reply content)

This caused flakiness because the reply for the first request would
sit in the channel, and _async_recv_reply would receive it but skip
it since msg_ids didn't match. On slow systems (like MacOS CI), the
second reply might not arrive before timeout.

Fix by using _recv_reply() / _async_recv_reply() to wait for the
first request's reply instead of sending a second request.

* Fix flaky client tests and handle flush race condition

Two fixes:

1. Fix flaky client tests by draining first reply before testing reply=True

   The tests send two requests back-to-back:
   - First call without reply=True (to check msg_id is returned)
   - Second call with reply=True (to check reply content)

   This caused flakiness because the reply for the first request would
   sit in the channel, and _async_recv_reply would skip it since msg_ids
   didn't match. On slow systems (like MacOS CI), the second reply might
   not arrive before timeout.

   Fix by draining the first reply before sending the second request.
   This still tests both paths: returning msg_id and the reply=True
   convenience functionality.

2. Handle stream closed between flush scheduling and execution

   The _flush() callback runs on the ioloop thread after being scheduled
   by flush(). Between scheduling and execution, the stream may be closed
   by stop_channels() during teardown.

   This is an expected race condition during shutdown, not a programming
   error, so we handle it gracefully rather than asserting.

   Uncovered by qtconsole downstream tests after #1089.

Author: rgbkrk

jupyter_client/threaded.py

@@ -214,7 +214,13 @@ def flush(f: Any) -> None:
 
     def _flush(self) -> None:
         """Callback for :method:`self.flush`."""
-        assert self.stream is not None
+        # Race condition: flush() checks stream validity then schedules this
+        # callback on the ioloop thread. Between scheduling and execution,
+        # stop_channels() may close the stream (e.g., during teardown).
+        # Handle gracefully rather than asserting, since this is an expected
+        # edge case during shutdown, not a programming error.
+        if self.stream is None or self.stream.closed():
+            return
         self.stream.flush()
         self._flushed = True
 

tests/test_client.py

@@ -50,34 +50,49 @@ def test_history(self):
         kc = self.kc
         msg_id = kc.history(session=0)
         self.assertIsInstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        kc._recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = kc.history(session=0, reply=True, timeout=TIMEOUT)
         self._check_reply("history", reply)
 
     def test_inspect(self):
         kc = self.kc
         msg_id = kc.inspect("who cares")
         self.assertIsInstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        kc._recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = kc.inspect("code", reply=True, timeout=TIMEOUT)
         self._check_reply("inspect", reply)
 
     def test_complete(self):
         kc = self.kc
         msg_id = kc.complete("who cares")
         self.assertIsInstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        kc._recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = kc.complete("code", reply=True, timeout=TIMEOUT)
         self._check_reply("complete", reply)
 
     def test_kernel_info(self):
         kc = self.kc
         msg_id = kc.kernel_info()
         self.assertIsInstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        kc._recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = kc.kernel_info(reply=True, timeout=TIMEOUT)
         self._check_reply("kernel_info", reply)
 
     def test_comm_info(self):
         kc = self.kc
         msg_id = kc.comm_info()
         self.assertIsInstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        kc._recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = kc.comm_info(reply=True, timeout=TIMEOUT)
         self._check_reply("comm_info", reply)
 
@@ -150,36 +165,54 @@ def output_hook(msg):
     async def test_history(self, kc):
         msg_id = kc.history(session=0)
         assert isinstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        await kc._async_recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = await kc.history(session=0, reply=True, timeout=TIMEOUT)
         self._check_reply("history", reply)
 
     async def test_inspect(self, kc):
         msg_id = kc.inspect("who cares")
         assert isinstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        await kc._async_recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = await kc.inspect("code", reply=True, timeout=TIMEOUT)
         self._check_reply("inspect", reply)
 
     async def test_complete(self, kc):
         msg_id = kc.complete("who cares")
         assert isinstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        await kc._async_recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = await kc.complete("code", reply=True, timeout=TIMEOUT)
         self._check_reply("complete", reply)
 
     async def test_is_complete(self, kc):
         msg_id = kc.is_complete("who cares")
         assert isinstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        await kc._async_recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = await kc.is_complete("code", reply=True, timeout=TIMEOUT)
         self._check_reply("is_complete", reply)
 
     async def test_kernel_info(self, kc):
         msg_id = kc.kernel_info()
         assert isinstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        await kc._async_recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = await kc.kernel_info(reply=True, timeout=TIMEOUT)
         self._check_reply("kernel_info", reply)
 
     async def test_comm_info(self, kc):
         msg_id = kc.comm_info()
         assert isinstance(msg_id, str)
+        # Drain the first reply to avoid race condition
+        await kc._async_recv_reply(msg_id, timeout=TIMEOUT)
+        # Now test the reply=True convenience path
         reply = await kc.comm_info(reply=True, timeout=TIMEOUT)
         self._check_reply("comm_info", reply)