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):