diff --git a/docs/02_execution_and_interactive.md b/docs/02_execution_and_interactive.md index ed1d782..061cafb 100644 --- a/docs/02_execution_and_interactive.md +++ b/docs/02_execution_and_interactive.md @@ -5,6 +5,8 @@ log: 2026-05-07: Fixed `print_kitty` (used by `colab exec --output-image` and any image-producing exec) to no-op when `sys.stdout.isatty()` is false. The Kitty Graphics Protocol escape sequence is meaningless when stdout is a file or pipe and was visually corrupting captured output (a multi-KB base64 PNG blob would land in log files, grep targets, or showboat captures). Image bytes are still saved to disk via `handle_image`'s file-write path; only the inline-render attempt is suppressed. 2026-06-04: Bumped the default `--timeout` for `colab exec` from 10s to 30s (and the matching `colab run` default) so brief silent tasks are less likely to hit a premature `TimeoutError`. Explicit `--timeout` overrides are unaffected. + +2026-10-24: `--timeout` for `colab exec` (and `colab run`) is now an idle limit: execution fails only after that many seconds without output, so code that keeps printing is never cut off, and `--timeout 0` means no limit. On 0.8.0/0.9.0 the old deadline was never enforced, so code that is silent for longer than `--timeout` now fails where it used to complete. `colab repl`, which passes no timeout, waits until each cell finishes, or until the kernel dies or restarts. None of them busy-loop while waiting. `colab install`, `colab auth`, and `colab drivemount` are unchanged. See #183. --- # Design: Execution and Interactive Interaction (`repl`, `exec`, `console`) @@ -28,7 +30,7 @@ Execution involves sending Python code (or shell commands) to the Jupyter kernel - If file path is local: Read content, send as code. - If file path is remote: Execute `!python `. - **Multi-Modal Output**: Handle `display_data` messages (e.g., `image/png`, `text/html`). For the CLI, we'll save images to temporary files and print their paths, or if the terminal supports it (e.g., iTerm2), inline them. -- **Timeout Configuration**: Exposes a `--timeout` flag (default 30s) to allow long-running silent tasks (like model compilation or data downloading) to execute without being prematurely killed. +- **Timeout Configuration**: Exposes a `--timeout` flag (default 30s): the longest silence allowed before execution fails. Raise it so long-running silent tasks (like model compilation or data downloading) are not prematurely killed; `0` means no limit. ### 3. Console (`colab console`) - **Implementation**: Connects directly to the backend terminal endpoint (`/colab/tty`) via WebSockets using `websocket-client`. diff --git a/docs/05_run_command.md b/docs/05_run_command.md index 9a1d421..f841461 100644 --- a/docs/05_run_command.md +++ b/docs/05_run_command.md @@ -5,6 +5,7 @@ log: 2026-05-12: Initial design and implementation of `colab run [args...]`. Combines `colab new` + `colab exec` + `colab stop` into a single fire-and-forget invocation so a Python file can use `#!/usr/bin/env -S colab run` as a shebang line and execute on a freshly-allocated Colab VM. Adds `--keep` (skip auto-stop), `--gpu` / `--tpu` (passthrough to session creation), `-s/--session` (name the ephemeral session), and propagates the script's exit status (non-zero on any uncaught exception in the kernel). The script's `sys.argv` is re-set inside the kernel to mirror native `python script.py arg1 arg2` semantics, and `__name__` is set to `"__main__"`. 2026-05-12: Native CPython exit-code semantics for `sys.exit()` / `raise SystemExit(...)` from the script body. The Colab kernel reports a `SystemExit` as `output_type=='error'`, which under the previous logic would have (a) printed the IPython traceback (`An exception has occurred, use %tb...`) and (b) flagged the run as a failure regardless of the integer exit code. Now: `sys.exit()` / `sys.exit(0)` exit 0 silently; `sys.exit(N)` exits N; `sys.exit('msg')` exits 1 (matching CPython). The IPython "To exit: use 'exit', 'quit', or Ctrl-D." UserWarning is filtered via the prelude. Encoded after running `examples/gpu_hello.py` end-to-end and seeing the noisy `SystemExit: 0` traceback at the end of an otherwise-successful GPU run. 2026-06-04: Bumped the default value of the `--timeout` flag from 10.0s to 30.0s so short-but-silent tasks aren't prematurely killed out of the box. Mirrors the same change for `colab exec`. +2026-10-24: `--timeout` is now an idle limit: the run fails only after that many seconds without output, so a script that keeps printing can run as long as it needs, and `--timeout 0` means no limit. The CLI waits for kernel messages itself, in blocking slices of up to one second, instead of using `jupyter-kernel-client`'s `execute_interactive`. That deadline covered the whole run: 1.0.2 enforces it, failing runs after 30s even while they print, and 0.8.0/0.9.0 never enforce it, polling at ~100% CPU once it passes. As a result, on 0.8.0/0.9.0 a script that is silent for longer than `--timeout` now fails (and, without `--keep`, its session is stopped) where it used to complete. See #183. --- # Design: `colab run` — Shebang-Compatible One-Shot Execution @@ -33,7 +34,7 @@ colab run [OPTIONS] SCRIPT [SCRIPT_ARGS]... | `--tpu` | str | None | Same set as `colab new --tpu` (v5e1, v6e1). | | `--high-mem` | bool | False | Same as `colab new --high-mem` — request high-RAM when supported. | | `--keep` | bool | False | Do **not** stop the session after the script finishes. | -| `--timeout` | float | 30.0 | Timeout in seconds for code execution to prevent hanging on silent tasks. | +| `--timeout` | float | 30.0 | Seconds to wait without output before failing, so a hung or silent task doesn't block forever. Raise it for long silent tasks; `0` means no limit. | ### Shebang usage With `--keep` and `--gpu` baked into the shebang line, an entire one-file workload becomes: diff --git a/src/colab_cli/commands/execution.py b/src/colab_cli/commands/execution.py index 336a5ff..4f6ac44 100644 --- a/src/colab_cli/commands/execution.py +++ b/src/colab_cli/commands/execution.py @@ -149,7 +149,10 @@ def exec_command( ] = None, timeout: Annotated[ Optional[float], - typer.Option("--timeout", help="Timeout in seconds for code execution"), + typer.Option( + "--timeout", + help="Seconds to wait without output before failing (0 for no limit)", + ), ] = 30.0, env: Annotated[ Optional[List[str]], diff --git a/src/colab_cli/commands/run.py b/src/colab_cli/commands/run.py index 0b22217..73893cb 100644 --- a/src/colab_cli/commands/run.py +++ b/src/colab_cli/commands/run.py @@ -258,7 +258,10 @@ def run_command( ] = False, timeout: Annotated[ Optional[float], - typer.Option("--timeout", help="Timeout in seconds for code execution"), + typer.Option( + "--timeout", + help="Seconds to wait without output before failing (0 for no limit)", + ), ] = 30.0, env: Annotated[ Optional[List[str]], diff --git a/src/colab_cli/runtime.py b/src/colab_cli/runtime.py index 57e09fd..b2a4146 100644 --- a/src/colab_cli/runtime.py +++ b/src/colab_cli/runtime.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import queue import time from typing import Any, Callable, Dict, List, Optional @@ -20,6 +21,12 @@ import requests +# The longest single wait for a kernel message, so a lost connection is noticed. +_WAIT_SLICE_SEC = 1.0 +# Kernel states that end any execution in progress. +_KERNEL_GONE = frozenset({"dead", "restarting"}) + + class ColabRuntime: def __init__( self, @@ -185,16 +192,17 @@ def execute_code( output_hook: Optional[Callable[[Dict[str, Any]], None]] = None, timeout: Optional[float] = None, ) -> List[Dict[str, Any]]: - # ``jupyter_kernel_client`` defaults ``timeout`` to ``REQUEST_TIMEOUT`` - # (10 seconds) on both ``execute`` and ``execute_interactive``. That - # value is a wall-clock budget that shrinks every time the poll loop - # iterates -- as long as iopub/stdin events arrive back-to-back the - # call survives, but a single >10s quiet stretch (e.g. a kernel - # blocked on ``input_request`` while the user OAuths in the browser) - # will raise ``TimeoutError`` even though the underlying execution is - # still healthy. Callers that know they need a longer ceiling can - # pass ``timeout=`` here; otherwise we forward whatever the upstream - # default is (currently 10s). + # ``timeout`` means one of two things, depending on the path: + # - Streaming without stdin (``output_hook`` set, ``allow_stdin`` + # False; used by run, exec and repl): ``_execute_streaming`` waits + # for messages itself, and ``timeout`` is an idle limit. ``None`` + # or a value <= 0 means no limit. + # - Otherwise the call goes to ``jupyter_kernel_client``, where + # ``timeout`` is a deadline for the whole call, defaulting to + # ``REQUEST_TIMEOUT`` (10 s). 0.8.0 and 0.9.0 never enforce it and + # poll without waiting once it passes; 1.0.2 raises + # ``TimeoutError``. Callers that need a longer ceiling, such as + # the interactive automations, pass ``timeout=`` explicitly. kwargs = {"allow_stdin": allow_stdin} if timeout is not None: kwargs["timeout"] = timeout @@ -218,7 +226,7 @@ def wrapped_stdin_hook(prompt): kwargs["stdin_hook"] = wrapped_stdin_hook if output_hook: - # If we have an output hook, we use execute_interactive and manage buffering ourselves + # With an output hook, we stream messages and manage buffering ourselves outputs = [] def wrapped_output_hook(msg): @@ -234,10 +242,13 @@ def wrapped_output_hook(msg): if idx < len(outputs): output_hook(outputs[idx]) - reply = self.kernel_client.execute_interactive( - code, output_hook=wrapped_output_hook, **kwargs - ) - # execute_interactive returns the raw reply message + if allow_stdin: + reply = self.kernel_client.execute_interactive( + code, output_hook=wrapped_output_hook, **kwargs + ) + else: + reply = self._execute_streaming(code, wrapped_output_hook, timeout) + # Both paths return the raw execute_reply message reply_content = reply["content"] if reply else {"status": "error"} else: reply = self.kernel_client.execute(code, **kwargs) @@ -261,6 +272,68 @@ def wrapped_output_hook(msg): return outputs + def _execute_streaming( + self, + code: str, + output_hook: Callable[[Dict[str, Any]], None], + timeout: Optional[float], + ) -> Dict[str, Any]: + """Run ``code``, pass each output message to ``output_hook``, return the reply. + + ``timeout`` is an idle limit: how long to wait, in seconds, for the next + message from this execution. A run that keeps printing can take as long + as it needs, and ``None`` or a value <= 0 waits without a limit. This + replaces ``jupyter_kernel_client``'s ``execute_interactive``, whose + deadline covers the whole run and which, in 0.8.0 and 0.9.0, polls + without waiting once the deadline passes. + + Like ``_apply_ws_hook``, this uses the websocket client behind the + kernel client. It skips that library's per-client execution lock, + because a ColabRuntime runs one execution at a time. + """ + if timeout is not None and timeout <= 0: + timeout = None + client = self.kernel_client._manager.client + msg_id = client.execute(code, allow_stdin=False) + last_message = time.monotonic() + + def next_message(get_msg, waiting_for): + nonlocal last_message + while True: + if not client.connection_ready.is_set(): + raise RuntimeError("Connection was lost.") + wait = _WAIT_SLICE_SEC + if timeout is not None: + left = last_message + timeout - time.monotonic() + if left <= 0: + raise TimeoutError( + f"Timeout waiting for {waiting_for}: no message from " + f"the kernel for {timeout:g}s. Use --timeout to allow " + "longer silences." + ) + wait = min(wait, left) + try: + msg = get_msg(timeout=wait) + except queue.Empty: + continue + state = msg["content"].get("execution_state") + if msg["header"]["msg_type"] == "status" and state in _KERNEL_GONE: + # The kernel announces these to every client, without a + # parent, and our execution can no longer finish. + raise RuntimeError(f"Kernel {state} before the code finished.") + if msg["parent_header"].get("msg_id") == msg_id: + last_message = time.monotonic() + return msg + + while True: + msg = next_message(client.get_iopub_msg, "output") + output_hook(msg) + if ( + msg["header"]["msg_type"] == "status" + and msg["content"].get("execution_state") == "idle" + ): + return next_message(client.get_shell_msg, "reply") + def stop(self, shutdown_kernel: bool = False): if self._kernel_client: try: diff --git a/tests/mock_kernel.py b/tests/mock_kernel.py new file mode 100644 index 0000000..a157970 --- /dev/null +++ b/tests/mock_kernel.py @@ -0,0 +1,119 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A mock of jupyter-kernel-client's websocket client for streaming tests. + +MockKernelClient replays a script of kernel messages on a simulated clock, so +tests can cover long runs and long silences without waiting. +""" + +import itertools +import queue +import threading +import time +from unittest.mock import MagicMock + +from colab_cli import runtime as runtime_module + +REQUEST = "req-1" +_msg_ids = itertools.count() + + +def kernel_message(msg_type, content, parent=REQUEST): + """A Jupyter message as the websocket client delivers it.""" + return { + "msg_type": msg_type, + "header": {"msg_type": msg_type, "msg_id": f"msg-{next(_msg_ids)}"}, + "parent_header": {"msg_id": parent} if parent else {}, + "content": content, + } + + +def stdout_message(text, parent=REQUEST): + return kernel_message("stream", {"name": "stdout", "text": text}, parent) + + +def status_message(state, parent=REQUEST): + return kernel_message("status", {"execution_state": state}, parent) + + +def reply_message(status="ok", parent=REQUEST, **content): + return kernel_message("execute_reply", {"status": status, **content}, parent) + + +class MockKernelClient: + """The websocket client's methods that ColabRuntime uses, on a simulated clock. + + ``iopub`` and ``shell`` are lists of (seconds after the request, message). + By default the shell channel replies when the last iopub message arrives. A + get_*_msg call that would wait past the next message advances the clock by + ``timeout`` and raises queue.Empty, as a quiet kernel would. The connection + drops when the clock reaches ``disconnect_at``. + """ + + def __init__(self, iopub, shell=None, disconnect_at=None): + self.now = 0.0 + self.iopub = list(iopub) + last = iopub[-1][0] if iopub else 0.0 + self.shell = list(shell if shell is not None else [(last, reply_message())]) + self.disconnect_at = disconnect_at + self.waits = [] + self.connection_ready = threading.Event() + self.connection_ready.set() + + def execute(self, code, **kwargs): + self.code, self.execute_kwargs = code, kwargs + return REQUEST + + def _get(self, messages, timeout): + self.waits.append(timeout) + assert timeout is not None and timeout > 0, "waits must be bounded" + if messages and messages[0][0] <= self.now + timeout: + at, msg = messages.pop(0) + self.now = max(self.now, at) + return msg + self.now += timeout + if self.disconnect_at is not None and self.now >= self.disconnect_at: + self.connection_ready.clear() + raise queue.Empty + + def get_iopub_msg(self, timeout=None): + return self._get(self.iopub, timeout) + + def get_shell_msg(self, timeout=None): + return self._get(self.shell, timeout) + + +class _SimulatedTime: + """The ``time`` module, with ``monotonic`` read from a MockKernelClient.""" + + def __init__(self, mock_kc): + self._mock_kc = mock_kc + + def monotonic(self): + return self._mock_kc.now + + def __getattr__(self, name): + return getattr(time, name) + + +def connect_runtime(runtime, mock_kc, monkeypatch): + """Make ``runtime`` use ``mock_kc``, and read its clock from ``mock_kc``. + + Only ``runtime``'s view of ``time.monotonic`` changes; everything else in + ``time`` stays real. + """ + monkeypatch.setattr(runtime_module, "time", _SimulatedTime(mock_kc)) + runtime._kernel_client = MagicMock() + runtime._kernel_client._manager.client = mock_kc diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 8d439d0..d50c160 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -15,6 +15,13 @@ from unittest.mock import MagicMock, patch import jupyter_kernel_client +import pytest +from mock_kernel import ( + MockKernelClient, + connect_runtime, + status_message, + stdout_message, +) from colab_cli.runtime import ColabRuntime @@ -162,18 +169,41 @@ def test_colab_runtime_execute_code_with_timeout(): def test_colab_runtime_execute_interactive_with_timeout(): """timeout must also be plumbed through the execute_interactive branch - (used when an output_hook is supplied).""" + (used when an output_hook is supplied with allow_stdin).""" runtime = ColabRuntime("http://url", "token123") mock_kc = MagicMock() runtime._kernel_client = mock_kc mock_kc.execute_interactive.return_value = {"content": {"status": "ok"}} - runtime.execute_code("print(1)", output_hook=lambda o: None, timeout=600) + runtime.execute_code( + "print(1)", output_hook=lambda o: None, allow_stdin=True, timeout=600 + ) _, kwargs = mock_kc.execute_interactive.call_args assert kwargs.get("timeout") == 600 +def test_colab_runtime_streaming_timeout_is_an_idle_limit(monkeypatch): + """Without stdin, an output_hook streams through _execute_streaming, where + timeout is the longest silence allowed between messages.""" + for silence, finishes in ((599, True), (601, False)): + runtime = ColabRuntime("http://url", "token123") + mock_kc = MockKernelClient( + [(silence, stdout_message("1\n")), (silence + 0.5, status_message("idle"))] + ) + connect_runtime(runtime, mock_kc, monkeypatch) + if finishes: + outputs = runtime.execute_code( + "print(1)", output_hook=lambda o: None, timeout=600 + ) + assert outputs[0]["text"] == "1\n" + else: + with pytest.raises(TimeoutError): + runtime.execute_code( + "print(1)", output_hook=lambda o: None, timeout=600 + ) + + def test_colab_runtime_stop(): runtime = ColabRuntime("http://url", "token123") mock_kc = MagicMock() diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 50e9f4c..bfe9339 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -12,87 +12,153 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock +import pytest +from mock_kernel import ( + MockKernelClient, + connect_runtime, + reply_message, + status_message, + stdout_message, +) + from colab_cli.runtime import ColabRuntime -def test_runtime_execute_code_streaming(): - runtime = ColabRuntime("http://url", "token") - mock_client = MagicMock() +@pytest.fixture +def run(monkeypatch): + """Run execute_code against a MockKernelClient; returns (outputs, streamed).""" - # Mock the return value of execute_interactive (the raw reply) - mock_client.execute_interactive.return_value = {"content": {"status": "ok"}} + def _run(mock_kc, code="work()", timeout=30.0): + runtime = ColabRuntime("http://url", "token") + connect_runtime(runtime, mock_kc, monkeypatch) + streamed = [] + outputs = runtime.execute_code( + code, output_hook=streamed.append, timeout=timeout + ) + return outputs, streamed - # Inject our mock client - runtime._kernel_client = mock_client + return _run - streamed_outputs = [] - def output_hook(o): - streamed_outputs.append(o) +def test_runtime_execute_code_streaming(run): + mock_kc = MockKernelClient( + [ + (0.1, stdout_message("1\n")), + (0.2, stdout_message("2\n")), + (0.3, status_message("idle")), + ] + ) + outputs, streamed = run(mock_kc, code="print(1); print(2)") - code = "print(1); print(2)" + assert [o["text"] for o in outputs] == ["1\n", "2\n"] + assert [o["text"] for o in streamed] == ["1\n", "2\n"] + assert mock_kc.code == "print(1); print(2)" + assert mock_kc.execute_kwargs == {"allow_stdin": False} - # We need to simulate the execution where execute_interactive is called - # and it calls our wrapped_output_hook. - # Since wrapped_output_hook is defined inside execute_code, we have to - # intercept the call to execute_interactive to get a reference to it. - def side_effect(code, output_hook=None, **kwargs): - # Simulate messages arriving - msg1 = { - "header": {"msg_type": "stream"}, - "content": {"name": "stdout", "text": "1\n"}, - } - msg2 = { - "header": {"msg_type": "stream"}, - "content": {"name": "stdout", "text": "2\n"}, - } - if output_hook: - output_hook(msg1) - output_hook(msg2) - return {"content": {"status": "ok"}} +def test_runtime_execute_code_streaming_error_synthesis(run): + reply = reply_message( + "error", ename="RuntimeError", evalue="something went wrong", traceback=["tb"] + ) + mock_kc = MockKernelClient([(0.1, status_message("idle"))], shell=[(0.1, reply)]) + outputs, streamed = run(mock_kc, code="fail()") - mock_client.execute_interactive.side_effect = side_effect + # The synthesized error is added after the run, so it is returned but not streamed. + assert len(outputs) == 1 + assert outputs[0]["output_type"] == "error" + assert outputs[0]["ename"] == "RuntimeError" + assert streamed == [] - outputs = runtime.execute_code(code, output_hook=output_hook) - assert len(outputs) == 2 - assert outputs[0]["text"] == "1\n" - assert outputs[1]["text"] == "2\n" +def test_a_run_that_keeps_printing_is_not_cut_off(run): + """A 10-minute run that prints every 20 s finishes under a 30 s idle limit.""" + lines = [(t, stdout_message(f"step {t}\n")) for t in range(20, 601, 20)] + mock_kc = MockKernelClient([*lines, (601, status_message("idle"))]) + outputs, _ = run(mock_kc, timeout=30.0) - # Verify streaming hook was called - assert len(streamed_outputs) == 2 - assert streamed_outputs[0]["text"] == "1\n" - assert streamed_outputs[1]["text"] == "2\n" + assert "".join(o["text"] for o in outputs).count("step") == 30 -def test_runtime_execute_code_streaming_error_synthesis(): - runtime = ColabRuntime("http://url", "token") - mock_client = MagicMock() +def test_silence_longer_than_the_timeout_raises(run): + mock_kc = MockKernelClient( + [(5, stdout_message("loading\n")), (100, status_message("idle"))] + ) + with pytest.raises(TimeoutError, match="Timeout waiting for output.*--timeout"): + run(mock_kc, timeout=30.0) + # 30 s after the last message, not 30 s after the start. + assert 35 <= mock_kc.now <= 36 - # Simulate an error reply but NO error output message - mock_client.execute_interactive.return_value = { - "content": { - "status": "error", - "ename": "RuntimeError", - "evalue": "something went wrong", - "traceback": ["tb line 1"], - } - } - runtime._kernel_client = mock_client - streamed_outputs = [] - outputs = runtime.execute_code( - "fail()", output_hook=lambda o: streamed_outputs.append(o) +@pytest.mark.parametrize("timeout", [None, 0]) +def test_no_limit_waits_through_any_silence(run, timeout): + mock_kc = MockKernelClient( + [(3600, stdout_message("done\n")), (3600.5, status_message("idle"))] ) + outputs, _ = run(mock_kc, timeout=timeout) + assert outputs[0]["text"] == "done\n" - # Final outputs should include synthesized error - assert len(outputs) == 1 - assert outputs[0]["output_type"] == "error" - assert outputs[0]["ename"] == "RuntimeError" - # Note: synthesized error is added AFTER execute_interactive returns, - # so it won't be in streamed_outputs unless we specifically add logic for it. - # Currently it's only in the returned list. - assert len(streamed_outputs) == 0 +def test_the_reply_may_arrive_after_the_idle_status(run): + mock_kc = MockKernelClient( + [(5, stdout_message("x\n")), (6, status_message("idle"))], + shell=[(26, reply_message())], + ) + outputs, _ = run(mock_kc, timeout=30.0) + assert outputs[0]["text"] == "x\n" + + +def test_a_reply_that_never_arrives_times_out(run): + mock_kc = MockKernelClient([(5, status_message("idle"))], shell=[]) + with pytest.raises(TimeoutError, match="Timeout waiting for reply"): + run(mock_kc, timeout=30.0) + + +def test_messages_from_other_requests_are_ignored(run): + mock_kc = MockKernelClient( + [ + (1, status_message("idle", parent="req-0")), + (2, stdout_message("ours\n")), + (3, status_message("idle")), + ], + shell=[(1, reply_message("error", parent="req-0")), (3, reply_message())], + ) + outputs, _ = run(mock_kc, timeout=30.0) + assert [o["text"] for o in outputs] == ["ours\n"] + + +def test_output_from_other_requests_does_not_reset_the_timeout(run): + others = [(t, stdout_message("other\n", parent="req-0")) for t in range(5, 100, 5)] + mock_kc = MockKernelClient([*others, (100, status_message("idle"))]) + with pytest.raises(TimeoutError): + run(mock_kc, timeout=30.0) + assert 30 <= mock_kc.now <= 31 + + +def test_waiting_never_spins(run): + """Each wait blocks for up to a second, so a long run costs a few calls a second.""" + lines = [(t, stdout_message("tick\n")) for t in range(25, 301, 25)] + mock_kc = MockKernelClient([*lines, (301, status_message("idle"))]) + run(mock_kc, timeout=30.0) + assert 0 < len(mock_kc.waits) < 330 # about one wait per second over 300 s + + +def test_a_lost_connection_stops_the_wait(run): + mock_kc = MockKernelClient( + [(5, stdout_message("x\n")), (100, status_message("idle"))], disconnect_at=10 + ) + with pytest.raises(RuntimeError, match="Connection was lost"): + run(mock_kc, timeout=None) + assert mock_kc.now < 12 + + +@pytest.mark.parametrize("state", ["dead", "restarting"]) +def test_a_kernel_that_dies_or_restarts_ends_the_wait(run, state): + mock_kc = MockKernelClient( + [ + (5, stdout_message("x\n")), + (8, status_message(state, parent=None)), + (100, status_message("idle")), + ] + ) + with pytest.raises(RuntimeError, match=f"Kernel {state}"): + run(mock_kc, timeout=None)