From f5255d3f26e2444afcd5ef91422680c9714de697 Mon Sep 17 00:00:00 2001 From: Matthew Printz Date: Tue, 28 Jul 2026 16:45:29 -0600 Subject: [PATCH] Testing fixes --- beaker-vue/package.json | 3 +- beaker-vue/playwright.config.ts | 21 ++- beaker-vue/tests/attachments.spec.ts | 21 ++- beaker-vue/tests/example.spec.ts | 44 ++++- beaker-vue/tests/global-setup.ts | 28 +++ beaker-vue/tests/global-teardown.ts | 72 ++++++++ beaker-vue/tests/helpers/session.ts | 174 ++++++++++++++++++ src/beaker_notebook/kernel.py | 25 ++- .../services/kernel/manager.py | 11 ++ 9 files changed, 380 insertions(+), 19 deletions(-) create mode 100644 beaker-vue/tests/global-setup.ts create mode 100644 beaker-vue/tests/global-teardown.ts create mode 100644 beaker-vue/tests/helpers/session.ts diff --git a/beaker-vue/package.json b/beaker-vue/package.json index 3ef00603..d7baf99f 100644 --- a/beaker-vue/package.json +++ b/beaker-vue/package.json @@ -14,7 +14,8 @@ "build-and-serve": "run-s build serve", "routes": "NODE_ENV=production vite-node utils/extractRoutes.ts", "preview": "vite preview", - "test:unit": "vitest", + "test:unit": "vitest run", + "test:unit:watch": "vitest", "test:e2e": "playwright test", "prepack": "vite-node utils/prepack", "postpack": "vite-node utils/postpack", diff --git a/beaker-vue/playwright.config.ts b/beaker-vue/playwright.config.ts index 9779413a..ac72b5c8 100644 --- a/beaker-vue/playwright.config.ts +++ b/beaker-vue/playwright.config.ts @@ -18,8 +18,12 @@ export default defineConfig({ // be cleanly separated from vitest unit tests (`src/**/__tests__/`) to avoid // the two runners collecting each other's specs. testDir: './tests', + /* Record pre-existing kernels, then reap anything the run leaked. See + * tests/global-setup.ts and tests/global-teardown.ts. */ + globalSetup: './tests/global-setup.ts', + globalTeardown: './tests/global-teardown.ts', /* Maximum time one test can run for. */ - timeout: 60 * 1000, + timeout: 120 * 1000, expect: { /** * Maximum time expect() should wait for the condition to be met. @@ -31,10 +35,17 @@ export default defineConfig({ forbidOnly: !!process.env.CI, /* Retry on CI only */ retries: process.env.CI ? 2 : 0, - /* Opt out of parallel tests on CI. */ - workers: process.env.CI ? 1 : undefined, - /* Reporter to use. See https://playwright.dev/docs/test-reporters */ - reporter: 'html', + /* Each test spins up its own kernel (beaker kernel + subkernel), so concurrency + * is bounded by machine resources rather than by CPU count. Left unbounded, + * every added test widens the startup pile-up and pushes the per-assertion + * timeouts below. */ + workers: 3, + /* Reporter to use. See https://playwright.dev/docs/test-reporters + * `list` gives live progress; the HTML report is still written, but never + * auto-opened -- the html reporter's default (`open: 'on-failure'`) starts a + * report server that blocks the terminal after a failed run. View it on + * demand with `npx playwright show-report`. */ + reporter: [['list'], ['html', {open: 'never'}]], /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */ use: { /* Maximum time each action such as `click()` can take. Defaults to 0 (no limit). */ diff --git a/beaker-vue/tests/attachments.spec.ts b/beaker-vue/tests/attachments.spec.ts index 20b1431d..3cc85083 100644 --- a/beaker-vue/tests/attachments.spec.ts +++ b/beaker-vue/tests/attachments.spec.ts @@ -1,5 +1,21 @@ import {expect, test, type Page} from '@playwright/test'; +import {deleteSessionByName, e2eSessionName, resolveBaseUrl} from './helpers/session'; + +/** + * This spec names its own session rather than going through the `sessionPage` + * fixture in example.spec.ts, because it asserts on session-scoped attachment + * state and needs the id up front. Cleanup is the same idea either way. + * `afterEach` also runs when the test fails, which is when a leaked kernel is + * most likely. + */ +let createdSession: string | null = null; + +test.afterEach(async () => { + if (!createdSession) return; + await deleteSessionByName(createdSession); + createdSession = null; +}); type BrowserFile = { name: string; @@ -22,8 +38,9 @@ async function dropFiles(page: Page, files: BrowserFile[]) { } test('notebook chat attachments upload, remove, extract ZIPs, and send without text', async ({page}) => { - const baseUrl = process.env.BEAKER_E2E_URL ?? 'http://localhost:8888'; - const sessionId = `attachment-e2e-${Date.now()}`; + const baseUrl = resolveBaseUrl(); + const sessionId = e2eSessionName(`attachments-${Date.now()}`); + createdSession = sessionId; await page.goto(`${baseUrl}/?session=${sessionId}`); await expect(page.locator('span.status-label')).toHaveText('Ready', {timeout: 30_000}); diff --git a/beaker-vue/tests/example.spec.ts b/beaker-vue/tests/example.spec.ts index e9610181..0a3c5692 100644 --- a/beaker-vue/tests/example.spec.ts +++ b/beaker-vue/tests/example.spec.ts @@ -1,12 +1,38 @@ import { test as base, expect, type Page } from '@playwright/test'; +import { deleteSessionByName, e2eSessionName, resolveBaseUrl } from './helpers/session'; + +/** + * One session per test, deleted as soon as the test finishes. + * + * Navigating to `/` (as this fixture used to) makes the server redirect to a + * freshly generated `?session=`, so every test already got its own + * session -- it just had no name the suite could recognise afterwards, so + * nothing was ever cleaned up. Naming the session ourselves keeps that + * isolation and makes teardown a direct delete. + * + * Sharing one session per worker was tried and does not work: notebook state + * persists per session, so the second test on a worker inherits the previous + * test's notebook instead of a fresh one, and the `.code-cell` locators find + * nothing. Isolation here is load-bearing, not incidental. + * + * Deleting the session shuts down its kernel and that kernel's subkernel, so + * live kernels track the worker count rather than the test count. + * + * The name must be unique per *run*, not just per test: notebook state is + * snapshotted to disk under the session name and outlives the session itself, + * so a name derived from testId alone would make each run restore the previous + * run's notebook -- which shows up as cells the test didn't expect. + */ const test = base.extend<{sessionPage: Page}>( { - sessionPage: async ({page}, use) => { - await page.goto('http://localhost:8888/'); + sessionPage: async ({page}, use, testInfo) => { + const sessionName = e2eSessionName(`${testInfo.testId}-${Date.now()}`); + await page.goto(`${resolveBaseUrl()}/?session=${encodeURIComponent(sessionName)}`); await expect(page.locator("span.status-label")).toBeVisible(); - await expect(page.locator("span.status-label")).toHaveText("Ready", {timeout: 20_000}); + await expect(page.locator("span.status-label")).toHaveText("Ready", {timeout: 45_000}); await use(page); + await deleteSessionByName(sessionName); } } ); @@ -17,9 +43,9 @@ test('Agent: Say Hello', async ({ sessionPage }) => { await expect(sessionPage.getByPlaceholder("Ask the AI or request an operation.")).toBeVisible(); await sessionPage.getByPlaceholder("Ask the AI or request an operation.").fill("say hello"); await sessionPage.getByLabel("Submit").click(); - await expect(sessionPage.getByText("Beaker Agent")).toBeVisible({timeout: 10_000}); + await expect(sessionPage.getByText("Beaker Agent")).toBeVisible({timeout: 30_000}); await expect(sessionPage.locator(".agent-cell-content")).toBeVisible(); - await expect(sessionPage.locator(".agent-cell-content")).toContainText("hello", {ignoreCase: true}); + await expect(sessionPage.locator(".agent-cell-content")).toContainText("hello", {ignoreCase: true, timeout: 60_000}); }); test('Agent: Run Code', async ({ sessionPage }) => { @@ -28,9 +54,9 @@ test('Agent: Run Code', async ({ sessionPage }) => { .getByPlaceholder("Ask the AI or request an operation.") .fill("use the python tool to compute 21 + 115"); await sessionPage.getByLabel("Submit").click(); - await expect(sessionPage.getByText("Beaker Agent")).toBeVisible({timeout: 10_000}); + await expect(sessionPage.getByText("Beaker Agent")).toBeVisible({timeout: 30_000}); await expect(sessionPage.locator(".code-cell")).toBeVisible(); - await expect(sessionPage.locator(".code-cell-output-box")).toContainText(`${21 + 115}`, {timeout: 10_000}); + await expect(sessionPage.locator(".code-cell-output-box")).toContainText(`${21 + 115}`, {timeout: 60_000}); }); test('Code Cell: Execute Python (control+enter)', async ({ sessionPage }) => { @@ -77,10 +103,10 @@ test('File Browser: Show README.md with metadata', async ({ sessionPage }) => { await filesTabButton.click(); const filePanel = sessionPage.locator('.file-container'); - await expect(filePanel).toBeVisible({timeout: 10_000}); + await expect(filePanel).toBeVisible({timeout: 30_000}); const fileTable = sessionPage.locator('.file-table'); - await expect(fileTable).toBeVisible(); + await expect(fileTable).toBeVisible({timeout: 30_000}); const readmeRow = fileTable.getByRole('row', { name: /README\.md/ }); await expect(readmeRow).toBeVisible(); diff --git a/beaker-vue/tests/global-setup.ts b/beaker-vue/tests/global-setup.ts new file mode 100644 index 00000000..0649187a --- /dev/null +++ b/beaker-vue/tests/global-setup.ts @@ -0,0 +1,28 @@ +/** + * Records which kernels already existed before the suite ran. + * + * `global-teardown.ts` uses this baseline so its kernel sweep only reaps what + * the run created. Without it, a sweep on a shared or local dev server would + * kill kernels a developer is actively using. + */ +import {mkdirSync, writeFileSync} from 'node:fs'; +import {dirname} from 'node:path'; + +import {BASELINE_PATH, createApiContext, listKernelIds} from './helpers/session'; + +export default async function globalSetup(): Promise { + let kernelIds: string[] = []; + const api = await createApiContext(); + try { + kernelIds = await listKernelIds(api); + } catch (error) { + // A server that isn't up yet is the tests' problem to report, not setup's. + console.warn('[e2e setup] could not read pre-existing kernels:', error); + } finally { + await api.dispose(); + } + + mkdirSync(dirname(BASELINE_PATH), {recursive: true}); + writeFileSync(BASELINE_PATH, JSON.stringify({kernelIds}, null, 2)); + console.log(`[e2e setup] ${kernelIds.length} pre-existing kernel(s) recorded; they will be left alone.`); +} diff --git a/beaker-vue/tests/global-teardown.ts b/beaker-vue/tests/global-teardown.ts new file mode 100644 index 00000000..5b28e341 --- /dev/null +++ b/beaker-vue/tests/global-teardown.ts @@ -0,0 +1,72 @@ +/** + * Backstop cleanup for the e2e suite. + * + * Each test already deletes its own session as it finishes; this exists for the + * cases that can't cover -- a crashed worker, a killed run, or a session whose + * kernel outlived its session record. + * + * Two passes, in order: + * 1. Delete every session named with the e2e prefix. Deleting a Beaker + * session shuts down its kernel, which shuts down that kernel's subkernel. + * 2. Delete any kernel that wasn't present before the run and isn't attached + * to a surviving session -- i.e. orphans left by a shutdown that didn't + * cascade. Kernels recorded by global setup are never touched. + */ +import {readFileSync} from 'node:fs'; + +import { + BASELINE_PATH, + createApiContext, + deleteKernels, + deleteSessions, + isE2ESession, + listKernelIds, + listSessions, +} from './helpers/session'; + +function readBaselineKernelIds(): Set | null { + try { + const {kernelIds} = JSON.parse(readFileSync(BASELINE_PATH, 'utf-8')) as {kernelIds: string[]}; + return new Set(kernelIds); + } catch { + // No baseline means we can't tell ours from theirs, so the kernel sweep + // is skipped entirely rather than risking a developer's running kernels. + return null; + } +} + +export default async function globalTeardown(): Promise { + const api = await createApiContext(); + try { + const deletedSessions = await deleteSessions(api, isE2ESession); + if (deletedSessions.length) { + console.log(`[e2e teardown] deleted ${deletedSessions.length} leftover session(s).`); + } + + const baseline = readBaselineKernelIds(); + if (!baseline) { + console.warn(`[e2e teardown] no baseline at ${BASELINE_PATH}; skipping kernel sweep.`); + return; + } + + // Anything still attached to a session is in use by something that isn't us. + const attached = new Set( + (await listSessions(api)) + .map((session) => session.kernel?.id) + .filter((id): id is string => Boolean(id)), + ); + const orphans = (await listKernelIds(api)).filter( + (id) => !baseline.has(id) && !attached.has(id), + ); + + if (orphans.length) { + const deleted = await deleteKernels(api, orphans); + console.log(`[e2e teardown] deleted ${deleted.length} orphaned kernel(s).`); + } + } catch (error) { + // Never fail a green run over cleanup. + console.warn('[e2e teardown] cleanup did not complete:', error); + } finally { + await api.dispose(); + } +} diff --git a/beaker-vue/tests/helpers/session.ts b/beaker-vue/tests/helpers/session.ts new file mode 100644 index 00000000..462af1fe --- /dev/null +++ b/beaker-vue/tests/helpers/session.ts @@ -0,0 +1,174 @@ +/** + * Shared helpers for creating and reaping Beaker sessions from e2e tests. + * + * Every session these tests create is named with `E2E_SESSION_PREFIX` so that + * `global-teardown.ts` can identify and sweep them without touching sessions a + * developer happens to have open against the same server. + * + * Cleanup deliberately goes through a standalone `APIRequestContext` rather + * than `page.request`: teardown most needs to run when a test has timed out or + * a worker is shutting down, and in those cases the page may already be gone. + */ +import {join} from 'node:path'; + +import {request as playwrightRequest, type APIRequestContext} from '@playwright/test'; + +export const E2E_SESSION_PREFIX = 'beaker-e2e'; + +/** + * Where global setup stashes the pre-run kernel list for global teardown to + * read. Both hooks run from the config directory, so a relative path is stable + * across the two. + */ +export const BASELINE_PATH = join('test-results', 'e2e-kernel-baseline.json'); + +export type SessionModel = { + id: string; + path?: string; + name?: string; + kernel?: {id: string} | null; +}; + +/** + * The specs address the server by absolute URL rather than via the config's + * `baseURL`, so resolution lives here to keep the two in agreement. + */ +export function resolveBaseUrl(): string { + return process.env.BEAKER_E2E_URL ?? 'http://localhost:8888'; +} + +export function e2eSessionName(suffix: string): string { + return `${E2E_SESSION_PREFIX}-${suffix}`; +} + +export function isE2ESession(session: SessionModel): boolean { + return [session.path, session.name].some((value) => value?.startsWith(E2E_SESSION_PREFIX)); +} + +/** Set E2E_CLEANUP_DEBUG=1 to trace what cleanup sends and what comes back. */ +const DEBUG = Boolean(process.env.E2E_CLEANUP_DEBUG); + +function debug(message: string): void { + if (DEBUG) console.log(`[e2e cleanup] ${message}`); +} + +/** + * Every `/api/*` call needs `Authorization: token ...`; a cookie-only request + * is 403 even for GET. The app bootstraps that token from an unauthenticated + * `GET /config` (see ConfigHandler in src/beaker_notebook/app/handlers.py), and + * that is where we get it too, so cleanup needs no configuration to work. + * BEAKER_E2E_TOKEN overrides it for servers where /config isn't reachable. + */ +async function resolveToken(): Promise { + if (process.env.BEAKER_E2E_TOKEN) return process.env.BEAKER_E2E_TOKEN; + + const bootstrap = await playwrightRequest.newContext({baseURL: resolveBaseUrl()}); + try { + const response = await bootstrap.get('/config'); + if (!response.ok()) { + console.warn(`[e2e cleanup] GET /config -> ${response.status()}; requests will be unauthenticated.`); + return null; + } + const {token} = (await response.json()) as {token?: string}; + debug(`resolved token from /config: ${token ? `${token.slice(0, 8)}...` : '(empty)'}`); + return token ?? null; + } catch (error) { + console.warn('[e2e cleanup] could not read /config:', error); + return null; + } finally { + await bootstrap.dispose(); + } +} + +export async function createApiContext(): Promise { + const token = await resolveToken(); + return playwrightRequest.newContext({ + baseURL: resolveBaseUrl(), + extraHTTPHeaders: token ? {Authorization: `token ${token}`} : {}, + }); +} + +/** + * Jupyter's Tornado handlers reject non-GET requests without an XSRF token, and + * the cookie alone is not enough -- it has to be echoed back as a header. + * `/config` is what seeds that cookie on a fresh context; `/` would too, but it + * 302s to a generated `?session=` on the way. + */ +export async function xsrfHeaders(api: APIRequestContext): Promise> { + await api.get('/config'); + const {cookies} = await api.storageState(); + const xsrf = cookies.find((cookie) => cookie.name === '_xsrf'); + debug(`xsrf cookie ${xsrf ? 'found' : 'MISSING'}`); + return xsrf ? {'X-XSRFToken': xsrf.value} : {}; +} + +export async function listSessions(api: APIRequestContext): Promise { + const response = await api.get('/api/sessions'); + if (!response.ok()) return []; + return (await response.json()) as SessionModel[]; +} + +export async function listKernelIds(api: APIRequestContext): Promise { + const response = await api.get('/api/kernels'); + if (!response.ok()) return []; + const kernels = (await response.json()) as {id: string}[]; + return kernels.map((kernel) => kernel.id); +} + +/** + * Deleting a Beaker session shuts down its kernel, which in turn shuts down the + * subkernel it spawned -- so sessions are the right granularity to reap at. + * Returns the ids actually deleted. + */ +export async function deleteSessions( + api: APIRequestContext, + matches: (session: SessionModel) => boolean, +): Promise { + const headers = await xsrfHeaders(api); + const deleted: string[] = []; + for (const session of (await listSessions(api)).filter(matches)) { + const response = await api.delete(`/api/sessions/${session.id}`, {headers}); + debug(`DELETE /api/sessions/${session.id} (path=${session.path}) -> ${response.status()}`); + // 404 means something else already reaped it, which is a fine outcome here. + if (response.ok() || response.status() === 404) { + deleted.push(session.id); + } else { + console.warn( + `[e2e cleanup] DELETE session ${session.id} -> ${response.status()} ` + + `${(await response.text()).slice(0, 200)}`, + ); + } + } + return deleted; +} + +export async function deleteKernels(api: APIRequestContext, kernelIds: string[]): Promise { + const headers = await xsrfHeaders(api); + const deleted: string[] = []; + for (const kernelId of kernelIds) { + const response = await api.delete(`/api/kernels/${kernelId}`, {headers}); + debug(`DELETE /api/kernels/${kernelId} -> ${response.status()}`); + if (response.ok() || response.status() === 404) { + deleted.push(kernelId); + } else { + console.warn( + `[e2e cleanup] DELETE kernel ${kernelId} -> ${response.status()} ` + + `${(await response.text()).slice(0, 200)}`, + ); + } + } + return deleted; +} + +/** Reap every session this suite created that still matches `name`. */ +export async function deleteSessionByName(name: string): Promise { + const api = await createApiContext(); + try { + await deleteSessions(api, (session) => session.path === name || session.name === name); + } catch (error) { + // Cleanup failure must never be the reason a run reports red. + console.warn(`[e2e cleanup] failed to delete session '${name}':`, error); + } finally { + await api.dispose(); + } +} diff --git a/src/beaker_notebook/kernel.py b/src/beaker_notebook/kernel.py index 0407043d..9304ac07 100644 --- a/src/beaker_notebook/kernel.py +++ b/src/beaker_notebook/kernel.py @@ -85,8 +85,29 @@ def __init__(self, session_config, kernel_id=None, connection_file=None): # Initialize context (Using the event loop to simulate `await`ing the async func in non-async setup) event_loop = asyncio.get_event_loop() logger.debug(f"About to start default context: {context_args}") + context_task = event_loop.create_task(self.start_default_context(**context_args)) - context_task.add_done_callback(lambda task: None) + context_task.add_done_callback(self.startup_error_callback) + + def startup_error_callback(self, task: asyncio.Task): + try: + exception = task.exception() + except (asyncio.CancelledError, asyncio.InvalidStateError) as err: + exception = err + if exception: + logger.exception("An error occurred while starting the default context.", exc_info=exception) + if self.context: + def log_error(log_task: asyncio.Task): + try: + log_exception = log_task.exception() + if log_exception: + logger.exception("An error occurred during default context cleanup", exc_info=log_exception) + except (asyncio.CancelledError, asyncio.InvalidStateError): + # Intentional swallow + pass + event_loop = asyncio.get_running_loop() + cleanup_task = event_loop.create_task(self.context.cleanup()) + cleanup_task.add_done_callback(log_error) async def start_default_context(self, default_context=None, default_context_payload=None, **options): logger.debug("starting default context!") @@ -936,7 +957,7 @@ def start(connection_file): # Perform shutdown cleanup here try: cleanup_task = partial(cleanup, kernel) - ioloop.IOLoop.current().run_sync(cleanup_task) + loop.run_sync(cleanup_task) except Exception as err: logger.exception(f"Error while shutting down kernel.") finally: diff --git a/src/beaker_notebook/services/kernel/manager.py b/src/beaker_notebook/services/kernel/manager.py index 730a9390..aa69d528 100644 --- a/src/beaker_notebook/services/kernel/manager.py +++ b/src/beaker_notebook/services/kernel/manager.py @@ -1,3 +1,4 @@ +import datetime import os import pwd import shutil @@ -17,6 +18,10 @@ class BeakerKernelManager(AsyncIOLoopKernelManager): beaker_session = Unicode(allow_none=True, help="Beaker session identifier", config=True) + last_activity: Optional[datetime.datetime] + execution_state: str + reason: str + # Longer wait_time for shutdown before killing processed due to potentially needing to shutdown both the subkernel # and the beaker kernel. shutdown_wait_time = Float( @@ -49,6 +54,12 @@ def app(self) -> "BaseBeakerApp": def setup_instance(self, *args, **kwargs): super().setup_instance(*args, **kwargs) + + # Required due to fragile monkey-patch in jupyter_server/services/kernels/kernelmanager.py:265:267 + self.last_activity = datetime.datetime.now(datetime.timezone.utc) + self.execution_state = "starting" + self.reason = "" + self.user = kwargs.get("user", None) async def _async_start_kernel(self, **kw):