Skip to content

Commit 88497d2

Browse files
TTTPOBclaude
andcommitted
perf(tests): reuse persistent kernel WebSocket across tests
Add live_kernel (session-scoped), mock_kernel_connection, and mock_execute_code fixtures to conftest. Tests that previously opened a new KernelClient per call now patch the connection layer and reuse one persistent WebSocket, eliminating the ~5s handshake cost per test. exec + notebook_writeback: 98.9s -> 8.7s (11x) variables + vars_cmd: also significantly faster via direct live_kernel use Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 859a6e5 commit 88497d2

5 files changed

Lines changed: 488 additions & 640 deletions

File tree

tests/conftest.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,3 +92,83 @@ def jupyter_server():
9292
except subprocess.TimeoutExpired:
9393
proc.kill()
9494
proc.wait()
95+
96+
97+
@pytest.fixture(scope="session")
98+
def live_session(jupyter_server):
99+
"""A single kernel session shared across the entire test session.
100+
101+
Tests that only run code and inspect results should use this fixture
102+
instead of creating their own session — kernel startup is expensive.
103+
Tests that mutate kernel lifecycle (restart, interrupt) must create
104+
their own private session via _create_session / _kill_session.
105+
"""
106+
import json
107+
from click.testing import CliRunner
108+
from jupyter_jcli.cli import main
109+
110+
runner = CliRunner()
111+
result = runner.invoke(main, [
112+
"-s", jupyter_server["url"], "-t", jupyter_server["token"],
113+
"--json", "session", "create", "--kernel", "python3",
114+
])
115+
data = json.loads(result.output)
116+
sid = data["session_id"]
117+
yield {**jupyter_server, "session_id": sid}
118+
runner.invoke(main, [
119+
"-s", jupyter_server["url"], "-t", jupyter_server["token"],
120+
"session", "kill", sid,
121+
])
122+
123+
124+
@pytest.fixture(scope="session")
125+
def live_kernel(live_session):
126+
"""A persistent WebSocket connection to the shared kernel.
127+
128+
Opened once per test session and reused across all tests. Tests that
129+
want to execute code or inspect variables should use mock_kernel_connection
130+
or mock_execute_code so the CLI path reuses this connection instead of
131+
opening a new one for every call.
132+
"""
133+
from jupyter_jcli.kernel import kernel_connection
134+
from jupyter_jcli.server import get_kernel_id_for_session
135+
136+
kernel_id = get_kernel_id_for_session(
137+
live_session["url"], live_session["session_id"], live_session["token"]
138+
)
139+
with kernel_connection(live_session["url"], live_session["token"], kernel_id) as kernel:
140+
yield kernel
141+
142+
143+
@pytest.fixture
144+
def mock_kernel_connection(live_kernel):
145+
"""Patch kernel_connection so CLI commands reuse live_kernel.
146+
147+
Use this for tests that invoke exec --file or vars through the CLI.
148+
The fixture patches the canonical source (jupyter_jcli.kernel) which
149+
is where both exec_cmd and vars_cmd lazily import from.
150+
"""
151+
from contextlib import contextmanager
152+
from unittest.mock import patch
153+
154+
@contextmanager
155+
def _reuse(*args, **kwargs):
156+
yield live_kernel
157+
158+
with patch("jupyter_jcli.kernel.kernel_connection", _reuse):
159+
yield live_kernel
160+
161+
162+
@pytest.fixture
163+
def mock_execute_code(live_kernel):
164+
"""Patch execute_code so exec --code reuses live_kernel.
165+
166+
Use this for tests that invoke exec --code through the CLI.
167+
"""
168+
from unittest.mock import patch
169+
170+
def _reuse(url, token, kid, code, timeout=300):
171+
return live_kernel.execute(code, timeout=timeout)
172+
173+
with patch("jupyter_jcli.kernel.execute_code", side_effect=_reuse):
174+
yield live_kernel

0 commit comments

Comments
 (0)