Restore erroneously deleted test

Author: ianthomas23

tests/test_subshells.py

@@ -6,12 +6,21 @@
 
 import platform
 import time
+from collections import Counter
 from queue import Empty
 
 import pytest
 from jupyter_client.blocking.client import BlockingKernelClient
 
-from .utils import TIMEOUT, assemble_output, flush_channels, get_replies, get_reply, new_kernel
+from .utils import (
+    TIMEOUT,
+    assemble_output,
+    flush_channels,
+    get_replies,
+    get_reply,
+    new_kernel,
+    wait_for_idle,
+)
 
 # Helpers
 
@@ -54,17 +63,17 @@ def execute_request_subshell_id(
 ):
     msg = execute_request(kc, code, subshell_id)
     msg_id = msg["header"]["msg_id"]
-    stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
+    stdout, _ = assemble_output(kc.get_iopub_msg, None, msg_id)
     return stdout.strip()
 
 
 def execute_thread_count(kc: BlockingKernelClient) -> int:
-    code = "import threading; print(threading.active_count())"
+    code = "print(threading.active_count())"
     return int(execute_request_subshell_id(kc, code, None))
 
 
 def execute_thread_ids(kc: BlockingKernelClient, subshell_id: str | None = None) -> tuple[str, str]:
-    code = "import threading; print(threading.get_ident(), threading.main_thread().ident)"
+    code = "print(threading.get_ident(), threading.main_thread().ident)"
     return execute_request_subshell_id(kc, code, subshell_id).split()
 
 
@@ -263,6 +272,60 @@ def test_execute_stop_on_error(are_subshells):
                 delete_subshell_helper(kc, subshell_id)
 
 
+@pytest.mark.parametrize("are_subshells", [(False, True), (True, False), (True, True)])
+def test_idle_message_parent_headers(are_subshells):
+    with new_kernel() as kc:
+        # import time module on main shell.
+        msg = kc.session.msg("execute_request", {"code": "import time"})
+        kc.shell_channel.send(msg)
+
+        subshell_ids = [
+            create_subshell_helper(kc)["subshell_id"] if is_subshell else None
+            for is_subshell in are_subshells
+        ]
+
+        # Wait for all idle status messages to be received.
+        for _ in range(1 + sum(are_subshells)):
+            wait_for_idle(kc)
+
+        msg_ids = []
+        for subshell_id in subshell_ids:
+            msg = execute_request(kc, "time.sleep(0.5)", subshell_id)
+            msg_ids.append(msg["msg_id"])
+
+        # Expect 4 status messages (2 busy, 2 idle) on iopub channel for the two execute_requests
+        statuses = []
+        timeout = TIMEOUT  # Combined timeout to receive all the status messages
+        t0 = time.time()
+        while True:
+            status = kc.get_iopub_msg(timeout=timeout)
+            if status["msg_type"] != "status" or status["parent_header"]["msg_id"] not in msg_ids:
+                continue
+            statuses.append(status)
+            if len(statuses) == 4:
+                break
+            t1 = time.time()
+            timeout -= t1 - t0
+            t0 = t1
+
+        execution_states = Counter(msg["content"]["execution_state"] for msg in statuses)
+        assert execution_states["busy"] == 2
+        assert execution_states["idle"] == 2
+
+        parent_msg_ids = Counter(msg["parent_header"]["msg_id"] for msg in statuses)
+        assert parent_msg_ids[msg_ids[0]] == 2
+        assert parent_msg_ids[msg_ids[1]] == 2
+
+        parent_subshell_ids = Counter(msg["parent_header"].get("subshell_id") for msg in statuses)
+        assert parent_subshell_ids[subshell_ids[0]] == 2
+        assert parent_subshell_ids[subshell_ids[1]] == 2
+
+        # Cleanup
+        for subshell_id in subshell_ids:
+            if subshell_id:
+                delete_subshell_helper(kc, subshell_id)
+
+
 def test_silent_flag_in_subshells():
     """Verifies that the 'silent' flag suppresses output in main and subshell contexts."""
     with new_kernel() as kc: