refactor: Add kernel_session() as context manager, DRY up tests

[mscolnick]

• Add kernel_session() context manager to marimo/_runtime/kernel_lifecycle.py
  and use it from launch_kernel so create+teardown are paired by the type
  system.
• Add tests/_runtime/_helpers/ with canonical MockStream/MockStdout/
  MockStderr/MockStdin, default_app_metadata/default_user_config
  factories, and mocked_kernel_session() returning a TestKernel bundle.
• Rewire tests/conftest.py fixtures (k, strict_kernel, lazy_kernel,
  run_mode_kernel, mocked_kernel, etc.) to delegate to
  mocked_kernel_session(). MockedKernel is preserved as a thin back-compat
  wrapper.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	May 15, 2026 1:58am

[cubic-dev-ai]

No issues found across 12 files

Architecture diagram

sequenceDiagram
    participant ProdRuntime as Production Runtime
    participant KernelLifecycle as kernel_lifecycle.py
    participant Kernel as Kernel
    participant TestFixtures as tests/conftest.py
    participant TestHelpers as tests/_runtime/_helpers/
    participant MockStreams as MockStream / MockStdout / MockStderr
    participant Factories as default_app_metadata / default_user_config
    participant TestSession as mocked_kernel_session()
    participant LoopDriver as LoopDriver
    participant TestKernel as TestKernel Bundle

    Note over ProdRuntime,Kernel: Production flow (unchanged, only refactored)
    ProdRuntime->>KernelLifecycle: launch_kernel()
    KernelLifecycle->>KernelLifecycle: kernel_session() (context manager)
    KernelLifecycle->>Kernel: create_kernel()
    Kernel-->>KernelLifecycle: kernel, ctx tuple
    KernelLifecycle->>KernelLifecycle: yield to caller
    KernelLifecycle-->>ProdRuntime: with kernel, ctx
    ProdRuntime->>KernelLifecycle: listen_messages() loop
    KernelLifecycle->>KernelLifecycle: teardown_kernel() on context exit

    Note over TestFixtures,LoopDriver: Test infrastructure (new / DRY'd)
    TestFixtures->>TestHelpers: fixture: k, strict_kernel, lazy_kernel, mocked_kernel
    TestHelpers->>TestSession: mocked_kernel_session()
    TestSession->>MockStreams: create MockStream / MockStdout / MockStderr / MockStdin
    TestSession->>Factories: default_app_metadata() / default_user_config()
    TestSession->>KernelLifecycle: kernel_session(stream, stdout, stderr, stdin, ...)
    KernelLifecycle->>Kernel: create_kernel() with mock I/O
    Kernel-->>KernelLifecycle: kernel, ctx
    KernelLifecycle-->>TestSession: kernel, ctx
    TestSession->>TestSession: build TestKernel(kernel, ctx, stream, ...)
    TestSession-->>TestHelpers: TestKernel bundle
    TestHelpers-->>TestFixtures: Kernel (or MockedKernel wrapper)
    TestFixtures->>TestHelpers: yield to test
    TestHelpers->>TestSession: exit context manager
    TestSession->>KernelLifecycle: teardown_kernel()
    KernelLifecycle->>Kernel: kernel.teardown()
    Kernel-->>KernelLifecycle: done
    TestSession->>TestSession: restore sys.modules["__main__"] and sys.meta_path

    Note over LoopDriver,TestKernel: LoopDriver for request-by-request control
    TestKernel->>LoopDriver: LoopDriver(kernel, control_queue, ui_queue)
    LoopDriver->>LoopDriver: start() → asyncio.create_task(listen_messages())
    LoopDriver->>TestKernel: enqueue(ExecuteCellsCommand)
    TestKernel->>Kernel: listen_messages processes command
    Kernel-->>TestKernel: execution
    LoopDriver->>TestKernel: settle() → drain control queue
    LoopDriver->>TestKernel: stop() → StopKernelCommand, await task

    Note over MockStreams: Shared canonical mock I/O
    MockStreams->>MockStreams: MockStream.write() captures KernelMessage list
    MockStreams->>MockStreams: MockStdout._write_with_mimetype() captures strings
    MockStreams->>MockStreams: MockStderr._write_with_mimetype() captures strings
    MockStreams->>MockStreams: MockStdin._readline_with_prompt() echoes prompt

    alt Test uses LoopDriver
        LoopDriver->>Kernel: step-wise control
    else Test uses direct kernel.run()
        TestHelpers->>Kernel: await kernel.run([exec_reqs])
    end

Loading

[Copilot]

Pull request overview

This PR refactors runtime test setup by introducing reusable test helpers (mock streams, kernel session context manager, hook recorder, loop driver) and rewiring existing fixtures/tests to use them, reducing boilerplate and pairing kernel creation/teardown more reliably.

Changes:

• Added kernel_session() context manager around create_kernel()/teardown_kernel() and used it in launch_kernel.
• Introduced tests/_runtime/_helpers/ (canonical mock streams, config factories, mocked_kernel_session(), HookRecorder, LoopDriver) and migrated fixtures/tests to use them.
• Added end-to-end tests for driving listen_messages via LoopDriver.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
marimo/_runtime/kernel_lifecycle.py	Adds kernel_session() context manager to ensure teardown pairing.
marimo/_runtime/runtime.py	Updates launch_kernel() to use kernel_session() for safer teardown behavior.
tests/conftest.py	Replaces inline mocked kernel/streams with helpers; keeps back-compat wrapper.
tests/_runtime/test_runtime.py	Uses mocked_kernel_session() and metadata factory to DRY up kernel setup in tests.
tests/_runtime/test_loop_driver.py	New integration-style tests for LoopDriver + real kernel control loop.
tests/_runtime/runner/test_hooks.py	Adds a test covering the new HookRecorder helper.
tests/_runtime/_helpers/__init__.py	Exposes helper APIs via a single import surface.
tests/_runtime/_helpers/streams.py	Centralizes mock stream implementations used by runtime tests.
tests/_runtime/_helpers/session.py	Adds mocked_kernel_session() context manager producing a TestKernel bundle.
tests/_runtime/_helpers/factories.py	Adds default AppMetadata / user-config factories.
tests/_runtime/_helpers/recorder.py	Adds HookRecorder for asserting hook invocation ordering/args.
tests/_runtime/_helpers/loop.py	Adds LoopDriver to drive listen_messages() stepwise in tests.

Comments suppressed due to low confidence (2)

tests/conftest.py:355

• In executing_kernel, setting mocked.k.stdout/stderr/stdin = None means Kernel.teardown() will skip stopping the underlying MockStdout/MockStderr/MockStdin watchers (it only calls _stop() when the attributes are non-None). This can leak watcher threads/file descriptors across tests. Consider either (a) keeping the mocks attached and disabling redirection another way, or (b) explicitly stopping mocked.stdout/_stderr/_stdin (or mocked._tk.*) before exiting the session so teardown remains reliable even after the kernel attributes are set to None.

    with mocked.k._install_execution_context(cell_id="0"):
        yield mocked.k
    mocked.teardown()

tests/_runtime/_helpers/factories.py:35

• default_user_config() uses a shallow DEFAULT_CONFIG.copy(), so nested dicts (e.g. runtime, display) remain shared with DEFAULT_CONFIG. If any test mutates nested config, it can contaminate global defaults and create order-dependent failures. Prefer a deep copy (e.g. marimo._config.utils.deep_copy(DEFAULT_CONFIG) or copy.deepcopy) before applying overrides.

[kirangadhave]

🚀 only nits

[kirangadhave]

can be typing only import

[kirangadhave]

Only applies to VirtualFileStorageType

[mscolnick]

probably. maybe we can enforce with linting?