Add tests

tests/test_kernel.py

@@ -58,6 +58,70 @@ def test_simple_print():
         _check_master(kc, expected=True)
 
 
+def test_print_to_correct_cell_from_thread():
+    """should print to the cell that spawned the thread, not a subsequently run cell"""
+    iterations = 5
+    interval = 0.25
+    code = f"""\
+    from threading import Thread
+    from time import sleep
+
+    def thread_target():
+        for i in range({iterations}):
+            print(i, end='', flush=True)
+            sleep({interval})
+
+    Thread(target=thread_target).start()
+    """
+    with kernel() as kc:
+        thread_msg_id = kc.execute(code)
+        _ = kc.execute("pass")
+
+        received = 0
+        while received < iterations:
+            msg = kc.get_iopub_msg(timeout=interval * 2)
+            if msg["msg_type"] != "stream":
+                continue
+            content = msg["content"]
+            assert content["name"] == "stdout"
+            assert content["text"] == str(received)
+            # this is crucial as the parent header decides to which cell the output goes
+            assert msg["parent_header"]["msg_id"] == thread_msg_id
+            received += 1
+
+
+def test_print_to_correct_cell_from_asyncio():
+    """should print to the cell that scheduled the task, not a subsequently run cell"""
+    iterations = 5
+    interval = 0.25
+    code = f"""\
+    import asyncio
+
+    async def async_task():
+        for i in range({iterations}):
+            print(i, end='', flush=True)
+            await asyncio.sleep({interval})
+
+    loop = asyncio.get_event_loop()
+    loop.create_task(async_task());
+    """
+    with kernel() as kc:
+        thread_msg_id = kc.execute(code)
+        _ = kc.execute("pass")
+
+        received = 0
+        while received < iterations:
+            msg = kc.get_iopub_msg(timeout=interval * 2)
+            if msg["msg_type"] != "stream":
+                continue
+            content = msg["content"]
+            assert content["name"] == "stdout"
+            assert content["text"] == str(received)
+            # this is crucial as the parent header decides to which cell the output goes
+            assert msg["parent_header"]["msg_id"] == thread_msg_id
+            received += 1
+
+
 @pytest.mark.skip(reason="Currently don't capture during test as pytest does its own capturing")
 def test_capture_fd():
     """simple print statement in kernel"""