diff --git a/tests/test_kernel.py b/tests/test_kernel.py index 07411bd1..70255910 100644 --- a/tests/test_kernel.py +++ b/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"""