fix threaded tests

jupyter · Sep 19, 2022 · ab93126 · ab93126
1 parent d3a21c4
commit ab93126
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/jupyter_client/threaded.py b/jupyter_client/threaded.py
@@ -117,8 +117,11 @@ def _handle_recv(self, future_msg: asyncio.Future) -> None:
         Unpacks message, and calls handlers with it.
         """
         assert self.ioloop is not None
-        assert future_msg.done()
-        msg_list = future_msg.result()
+        if isinstance(future_msg, asyncio.Future):
+            assert future_msg.done()
+            msg_list = future_msg.result()
+        else:
+            msg_list = future_msg
         assert self.session is not None
         ident, smsg = self.session.feed_identities(msg_list)
         msg = self.session.deserialize(smsg)

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -264,8 +264,8 @@ def test_comm_info(self):
     def test_shutdown(self):
         kc = self.kc
         kc.shutdown()
-        kc.shell_channel.msg_recv.wait()
-        reply = kc.shell_channel.last_msg
+        kc.control_channel.msg_recv.wait()
+        reply = kc.control_channel.last_msg
         self._check_reply("shutdown", reply)
 
     def test_shutdown_id(self):