Debug

Author: ianthomas23

tests/test_subshells.py

@@ -10,7 +10,7 @@
 import pytest
 from jupyter_client.blocking.client import BlockingKernelClient
 
-from .utils import TIMEOUT, get_replies, get_reply, new_kernel
+from .utils import TIMEOUT, assemble_output, get_replies, get_reply, new_kernel
 
 # Helpers
 
@@ -20,6 +20,7 @@ def create_subshell_helper(kc: BlockingKernelClient):
     kc.control_channel.send(msg)
     msg_id = msg["header"]["msg_id"]
     reply = get_reply(kc, msg_id, TIMEOUT, channel="control")
+    assemble_output(kc.get_iopub_msg)  # wait for idle
     return reply["content"]
 
 
@@ -28,6 +29,7 @@ def delete_subshell_helper(kc: BlockingKernelClient, subshell_id: str):
     kc.control_channel.send(msg)
     msg_id = msg["header"]["msg_id"]
     reply = get_reply(kc, msg_id, TIMEOUT, channel="control")
+    assemble_output(kc.get_iopub_msg)  # wait for idle
     return reply["content"]
 
 
@@ -36,6 +38,7 @@ def list_subshell_helper(kc: BlockingKernelClient):
     kc.control_channel.send(msg)
     msg_id = msg["header"]["msg_id"]
     reply = get_reply(kc, msg_id, TIMEOUT, channel="control")
+    assemble_output(kc.get_iopub_msg)  # wait for idle
     return reply["content"]
 
 
@@ -49,20 +52,8 @@ def execute_request(kc: BlockingKernelClient, code: str, subshell_id: str | None
 def execute_request_subshell_id(
     kc: BlockingKernelClient, code: str, subshell_id: str | None, terminator: str = "\n"
 ):
-    msg = execute_request(kc, code, subshell_id)
-    msg_id = msg["msg_id"]
-    stdout = ""
-    while True:
-        msg = kc.get_iopub_msg()
-        # Get the stream messages corresponding to msg_id
-        if (
-            msg["msg_type"] == "stream"
-            and msg["parent_header"]["msg_id"] == msg_id
-            and msg["content"]["name"] == "stdout"
-        ):
-            stdout += msg["content"]["text"]
-            if stdout.endswith(terminator):
-                break
+    execute_request(kc, code, subshell_id)
+    stdout, _ = assemble_output(kc.get_iopub_msg)
     return stdout.strip()