From 1b38046f5b162daf5dbd4c7f15457cf95ca209d5 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:02:05 +0200 Subject: [PATCH 01/22] feat(appkit): make PluginContext telemetry injectable The testing kit needs to construct a real PluginContext without a live OpenTelemetry pipeline. Add an optional constructor dependency for the telemetry provider, defaulting to the shared "plugin-context" provider so the production path is unchanged. This is the single production edit required to wrap the real class in tests rather than reimplementing it. Signed-off-by: Galymzhan --- packages/appkit/src/core/plugin-context.ts | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/packages/appkit/src/core/plugin-context.ts b/packages/appkit/src/core/plugin-context.ts index 4f08dcd91..752ac53ab 100644 --- a/packages/appkit/src/core/plugin-context.ts +++ b/packages/appkit/src/core/plugin-context.ts @@ -1,7 +1,11 @@ import type express from "express"; import type { BasePlugin, IAppRequest, ToolProvider } from "shared"; import { createLogger } from "../logging/logger"; -import { SpanStatusCode, TelemetryManager } from "../telemetry"; +import { + type ITelemetry, + SpanStatusCode, + TelemetryManager, +} from "../telemetry"; import { forwardAsyncErrors } from "../utils/safe-handler"; const logger = createLogger("plugin-context"); @@ -62,7 +66,20 @@ export class PluginContext { LifecycleEvent, Set<() => void | Promise> >(); - private telemetry = TelemetryManager.getProvider("plugin-context"); + private telemetry: ITelemetry; + + /** + * @param deps.telemetry - Telemetry provider used for `executeTool` spans. + * Defaults to the shared `"plugin-context"` provider — the production + * path. Injectable so the testing kit can pass a mock provider and run + * `executeTool` without a live OpenTelemetry pipeline. This is the only + * seam the mock context needs; route buffering and the tool registry are + * exercised through the existing public API. + */ + constructor(deps: { telemetry?: ITelemetry } = {}) { + this.telemetry = + deps.telemetry ?? TelemetryManager.getProvider("plugin-context"); + } /** * Register a route on the root Express application. From dd503f52c4ad75536b2805440da5d7085b7948b0 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:36:09 +0200 Subject: [PATCH 02/22] feat(appkit): ship @databricks/appkit/testing and migrate first stub Wire the testing kit as a published subpath and prove it against the first of the two hand-rolled context stubs (the design gate): - Add ./testing to both exports maps (dev + publishConfig) following the ./type-generator shape, add src/testing/index.ts to the tsdown entry, and declare vitest as an optional peerDependency. Build passes attw + publint; dist/testing/{index,mock-plugin-context,expect-stream,fixtures}.{js,d.ts} are emitted and vitest stays external to the main entry. - Migrate dispatch-tool-call.test.ts: replace (plugin as any).context = { executeTool } with mockPluginContext. executeTool is now the REAL method, so the forwarded toolCallTimeoutMs is asserted through actual signal composition, the on-behalf-of (asUser) path is verified, and a new test proves the forwarded timeout actually aborts a slow toolkit tool end-to-end. This is the primary win from the plan: executeTool's OBO and timeout paths gain real assertions instead of a stub that proved nothing. Signed-off-by: Galymzhan --- knip.json | 3 + packages/appkit/package.json | 17 + .../agents/tests/dispatch-tool-call.test.ts | 78 +++- packages/appkit/src/testing/expect-stream.ts | 223 +++++++++ packages/appkit/src/testing/fixtures.ts | 437 ++++++++++++++++++ packages/appkit/src/testing/index.ts | 66 +++ .../appkit/src/testing/mock-plugin-context.ts | 291 ++++++++++++ .../src/testing/tests/expect-stream.test.ts | 142 ++++++ .../testing/tests/mock-plugin-context.test.ts | 164 +++++++ packages/appkit/tsdown.config.ts | 2 +- pnpm-lock.yaml | 92 ++++ 11 files changed, 1500 insertions(+), 15 deletions(-) create mode 100644 packages/appkit/src/testing/expect-stream.ts create mode 100644 packages/appkit/src/testing/fixtures.ts create mode 100644 packages/appkit/src/testing/index.ts create mode 100644 packages/appkit/src/testing/mock-plugin-context.ts create mode 100644 packages/appkit/src/testing/tests/expect-stream.test.ts create mode 100644 packages/appkit/src/testing/tests/mock-plugin-context.test.ts diff --git a/knip.json b/knip.json index 0e96b7df5..1ca5a1b7e 100644 --- a/knip.json +++ b/knip.json @@ -9,6 +9,9 @@ "workspaces": { "packages/appkit-ui": { "ignoreDependencies": ["tailwindcss", "tw-animate-css"] + }, + "packages/appkit": { + "ignoreDependencies": ["vitest"] } }, "ignore": [ diff --git a/packages/appkit/package.json b/packages/appkit/package.json index ea70b95d9..dd19f3b52 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -42,6 +42,11 @@ "development": "./src/type-generator/index.ts", "default": "./dist/type-generator/index.js" }, + "./testing": { + "types": "./dist/testing/index.d.ts", + "development": "./src/testing/index.ts", + "default": "./dist/testing/index.js" + }, "./dist/shared/src/plugin": { "types": "./dist/shared/src/plugin.d.ts", "default": "./dist/shared/src/plugin.d.ts" @@ -103,6 +108,14 @@ "@types/ws": "8.18.1", "@vitejs/plugin-react": "5.1.1" }, + "peerDependencies": { + "vitest": ">=1.0.0" + }, + "peerDependenciesMeta": { + "vitest": { + "optional": true + } + }, "overrides": { "vite": "npm:rolldown-vite@7.1.14" }, @@ -113,6 +126,10 @@ "./beta": "./dist/beta.js", "./dist/shared/src/plugin": "./dist/shared/src/plugin.d.ts", "./type-generator": "./dist/type-generator/index.js", + "./testing": { + "types": "./dist/testing/index.d.ts", + "default": "./dist/testing/index.js" + }, "./package.json": "./package.json" } } diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index 2767766c3..d5915feca 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -1,6 +1,7 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; +import { mockPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; /** @@ -289,16 +290,8 @@ describe("dispatchToolCall — toolkit timeout plumbing", () => { * `runState.limits.toolCallTimeoutMs` through to `PluginContext` so the * agents plugin owns the cap and the default (5 minutes) is generous. */ - test("forwards runState.limits.toolCallTimeoutMs to PluginContext.executeTool", async () => { - const plugin = new AgentsPlugin({ dir: false }); - const { runState } = makeRunState(plugin); - runState.limits.toolCallTimeoutMs = 90_000; - - const executeTool = vi.fn().mockResolvedValue("rows"); - // biome-ignore lint/suspicious/noExplicitAny: stub PluginContext shape - (plugin as any).context = { executeTool }; - - const toolIndex = new Map([ + const toolkitToolIndex = () => + new Map([ [ "analytics.query", { @@ -314,19 +307,76 @@ describe("dispatchToolCall — toolkit timeout plumbing", () => { ], ]); - await callDispatch(plugin, { + test("forwards runState.limits.toolCallTimeoutMs to PluginContext.executeTool", async () => { + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + runState.limits.toolCallTimeoutMs = 90_000; + + // Use the real PluginContext via the testing kit rather than a bare + // `{ executeTool }` stub. `executeTool` here is the real method, so the + // forwarded timeout is exercised through actual signal composition — and + // spying on it lets us keep asserting the exact call signature the agents + // plugin passes. + const mock = mockPluginContext({ analytics: { query: "rows" } }); + const executeToolSpy = vi.spyOn(mock.ctx, "executeTool"); + // biome-ignore lint/suspicious/noExplicitAny: attach the real context to the plugin + (plugin as any).context = mock.ctx; + + const result = await callDispatch(plugin, { runState, - toolIndex, + toolIndex: toolkitToolIndex(), name: "analytics.query", args: { sql: "SELECT 1" }, }); - expect(executeTool).toHaveBeenCalledTimes(1); - const call = executeTool.mock.calls[0]; + expect(result).toBe("rows"); + expect(executeToolSpy).toHaveBeenCalledTimes(1); + const call = executeToolSpy.mock.calls[0]; // (req, pluginName, toolName, args, signal, timeoutMs) expect(call[1]).toBe("analytics"); expect(call[2]).toBe("query"); expect(call[5]).toBe(90_000); + + // The stub could never prove this: the real executeTool routed the call + // through the analytics provider's on-behalf-of (asUser) path. + expect(mock.toolCalls).toHaveLength(1); + expect(mock.toolCalls[0]).toMatchObject({ + plugin: "analytics", + tool: "query", + args: { sql: "SELECT 1" }, + asUser: true, + }); + }); + + test("the forwarded timeout actually aborts a slow toolkit tool", async () => { + // End-to-end proof that the timeout value the agents plugin forwards + // reaches real AbortSignal composition inside PluginContext.executeTool — + // a stubbed executeTool would silently ignore the timeout. + const plugin = new AgentsPlugin({ dir: false }); + const { runState } = makeRunState(plugin); + runState.limits.toolCallTimeoutMs = 5; + + const mock = mockPluginContext({ + analytics: { + query: (_args, signal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject(new Error("aborted by toolkit timeout")), + ); + }), + }, + }); + // biome-ignore lint/suspicious/noExplicitAny: attach the real context + (plugin as any).context = mock.ctx; + + await expect( + callDispatch(plugin, { + runState, + toolIndex: toolkitToolIndex(), + name: "analytics.query", + args: { sql: "SELECT 1" }, + }), + ).rejects.toThrow(/aborted by toolkit timeout/); }); test("resolvedLimits exposes the documented 5-minute default", () => { diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts new file mode 100644 index 000000000..cc32de943 --- /dev/null +++ b/packages/appkit/src/testing/expect-stream.ts @@ -0,0 +1,223 @@ +/** + * A single event observed on a stream. AppKit adapters yield objects with a + * `type` discriminator; SSE frames parsed from an HTTP response carry the + * event name under `event`. {@link expectStream} normalizes both to a type + * string, preferring `type` and falling back to `event`. + */ +export interface StreamEvent { + type?: string; + event?: string; + [key: string]: unknown; +} + +/** + * Anything {@link expectStream} can consume: + * - an async event stream (an adapter's `run()`, an SSE reader), + * - an already-collected array of events, + * - an SSE `Response` (or a promise of one) — its body is parsed into events. + */ +export type StreamSource = + | AsyncIterable + | Iterable + | Response + | Promise; + +/** Assertions over the events collected from a {@link StreamSource}. */ +export interface StreamAssertion { + /** + * Assert that `eventTypes` appear, in this order, among the emitted event + * types. Extra events (heartbeats, metadata, deltas) may appear before, + * between, or after — this is an in-order subsequence match, which is what + * you want for streams that interleave bookkeeping events. Resolves to the + * full list of emitted types on success; rejects with a diff otherwise. + */ + toEmit(...eventTypes: string[]): Promise; + /** + * Assert that the emitted event types are exactly `eventTypes`, in order and + * with nothing else. Use when the stream's shape is fully determined. + */ + toEmitExactly(...eventTypes: string[]): Promise; + /** Collect and return the normalized events without asserting. */ + collect(): Promise; + /** Collect and return just the event type strings, in order. */ + collectTypes(): Promise; +} + +function eventType(event: StreamEvent): string { + return event.type ?? event.event ?? ""; +} + +/** + * Parse a finished SSE response body into events. Blocks are delimited by a + * blank line; within a block, `event:` sets the name and `data:` lines are + * joined and JSON-parsed when possible. Comment/heartbeat lines (`:`) and + * blocks without data are ignored. + */ +function parseSSEBody(text: string): StreamEvent[] { + const events: StreamEvent[] = []; + const blocks = text.split(/\n\n/); + + for (const block of blocks) { + let name: string | undefined; + const dataLines: string[] = []; + + for (const rawLine of block.split("\n")) { + const line = rawLine.replace(/\r$/, ""); + if (line.startsWith("event:")) { + name = line.slice("event:".length).trim(); + } else if (line.startsWith("data:")) { + dataLines.push(line.slice("data:".length).replace(/^ /, "")); + } + // `id:` and comment (`:`) lines carry no event type/data we assert on. + } + + if (name === undefined && dataLines.length === 0) continue; + + const data = dataLines.join("\n"); + let parsed: Record = {}; + if (data) { + try { + const json = JSON.parse(data); + if (json && typeof json === "object" && !Array.isArray(json)) { + parsed = json as Record; + } else { + parsed = { data: json }; + } + } catch { + parsed = { data }; + } + } + + events.push({ + type: name ?? (parsed.type as string | undefined), + ...parsed, + }); + } + + return events; +} + +async function collectEvents(source: StreamSource): Promise { + const resolved = await source; + + if (resolved instanceof Response) { + const text = await resolved.text(); + return parseSSEBody(text); + } + + if (resolved && typeof resolved === "object") { + if (Symbol.asyncIterator in resolved) { + const events: StreamEvent[] = []; + for await (const event of resolved as AsyncIterable) { + events.push(event); + } + return events; + } + if (Symbol.iterator in resolved) { + return Array.from(resolved as Iterable); + } + } + + throw new Error( + "expectStream: source must be an async iterable, an iterable, or a Response", + ); +} + +/** Does `expected` appear as an in-order subsequence of `actual`? */ +function isSubsequence(actual: string[], expected: string[]): boolean { + let i = 0; + for (const type of actual) { + if (i < expected.length && type === expected[i]) i++; + } + return i === expected.length; +} + +/** + * Consume a stream and make ordered assertions about the event types it emits. + * + * Deterministic and network-free: pair it with {@link mockPluginContext} to + * exercise a plugin's streaming handler and assert what it emits. + * + * @example Async event stream (adapter output) + * ```ts + * await expectStream(agent.adapter.run(input)).toEmit("tool_call", "message_delta"); + * ``` + * + * @example SSE HTTP response + * ```ts + * const res = await fetch("/api/analytics/query/top_users", { method: "POST" }); + * await expectStream(res).toEmit("warehouse_status", "result"); + * ``` + */ +export function expectStream(source: StreamSource): StreamAssertion { + const events = collectEvents(source); + + return { + async collect() { + return events; + }, + async collectTypes() { + return (await events).map(eventType); + }, + async toEmit(...eventTypes: string[]) { + const types = (await events).map(eventType); + if (!isSubsequence(types, eventTypes)) { + throw new Error( + `expectStream(...).toEmit: expected events ${JSON.stringify( + eventTypes, + )} in order, but stream emitted ${JSON.stringify(types)}`, + ); + } + return types; + }, + async toEmitExactly(...eventTypes: string[]) { + const types = (await events).map(eventType); + const equal = + types.length === eventTypes.length && + types.every((t, i) => t === eventTypes[i]); + if (!equal) { + throw new Error( + `expectStream(...).toEmitExactly: expected exactly ${JSON.stringify( + eventTypes, + )}, but stream emitted ${JSON.stringify(types)}`, + ); + } + return types; + }, + }; +} + +/** + * Parse a single-event SSE `Response` into `{ eventType, ...data }`. + * + * Retained for tests that assert on a one-shot SSE reply; prefer + * {@link expectStream} for multi-event ordering assertions. + */ +export async function parseSSEResponse(response: Response): Promise<{ + eventType: string | null; + [key: string]: unknown; +}> { + const text = await response.text(); + const lines = text.split("\n"); + + let eventType: string | null = null; + let dataLine: string | null = null; + + for (const line of lines) { + if (line.startsWith("event: ")) { + eventType = line.substring(7).trim(); + } else if (line.startsWith("data: ")) { + dataLine = line.substring(6); + } + } + + if (!dataLine) { + throw new Error(`No data found in SSE response: ${text}`); + } + + const parsed = JSON.parse(dataLine); + return { + eventType, + ...parsed, + }; +} diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts new file mode 100644 index 000000000..ddb873f0a --- /dev/null +++ b/packages/appkit/src/testing/fixtures.ts @@ -0,0 +1,437 @@ +import type { Span, SpanOptions } from "@opentelemetry/api"; +import type { IAppRouter } from "shared"; +import { vi } from "vitest"; +import type { ServiceContextState } from "../context/service-context"; +import { ServiceContext } from "../context/service-context"; +import type { UserContext } from "../context/user-context"; +import type { InstrumentConfig, ITelemetry } from "../telemetry/types"; + +// biome-ignore lint/suspicious/noExplicitAny: test fixtures intentionally use loose shapes +type Any = any; + +/** + * Creates a mock telemetry provider for testing. Every span/meter/logger is a + * `vi.fn()` no-op, so plugins that trace, count, or log run without a live + * OpenTelemetry pipeline. Passed into {@link mockPluginContext} as the one + * injectable production seam. + */ +export function createMockTelemetry(): ITelemetry { + const mockSpan: Span = { + addLink: vi.fn(), + addLinks: vi.fn(), + end: vi.fn(), + setAttribute: vi.fn(), + setAttributes: vi.fn(), + setStatus: vi.fn(), + recordException: vi.fn(), + updateName: vi.fn(), + addEvent: vi.fn(), + isRecording: vi.fn().mockReturnValue(false), + spanContext: vi.fn(), + }; + + return { + getTracer: vi.fn().mockReturnValue({ + startActiveSpan: vi.fn().mockImplementation((...args: Any[]) => { + const fn = args[args.length - 1]; + if (typeof fn === "function") { + return fn(mockSpan); + } + return undefined; + }), + }), + getMeter: vi.fn().mockReturnValue({ + createCounter: vi.fn().mockReturnValue({ add: vi.fn() }), + createHistogram: vi.fn().mockReturnValue({ record: vi.fn() }), + }), + getLogger: vi.fn().mockReturnValue({ + emit: vi.fn(), + }), + emit: vi.fn(), + startActiveSpan: vi + .fn() + .mockImplementation( + async ( + _name: string, + _options: SpanOptions, + fn: (span: Span) => Promise, + _tracerOptions?: InstrumentConfig, + ) => { + return await fn(mockSpan); + }, + ), + registerInstrumentations: vi.fn(), + }; +} + +/** + * Creates a mock Express router that captures registered handlers so a test + * can pull a handler back out by method + path and invoke it directly. + */ +export function createMockRouter(): { + router: IAppRouter; + handlers: Record; + getHandler: (method: string, path: string) => Any; +} { + const handlers: Record = {}; + + const mockRouter = { + get: vi.fn((path: string, handler: Any) => { + handlers[`GET:${path}`] = handler; + }), + post: vi.fn((path: string, handler: Any) => { + handlers[`POST:${path}`] = handler; + }), + put: vi.fn((path: string, handler: Any) => { + handlers[`PUT:${path}`] = handler; + }), + delete: vi.fn((path: string, handler: Any) => { + handlers[`DELETE:${path}`] = handler; + }), + patch: vi.fn((path: string, handler: Any) => { + handlers[`PATCH:${path}`] = handler; + }), + } as unknown as IAppRouter; + + return { + router: mockRouter, + handlers, + getHandler: (method: string, path: string) => + handlers[`${method.toUpperCase()}:${path}`], + }; +} + +/** + * Creates a mock Express request object. Carries a default mock + * WorkspaceClient (SQL succeeds, warehouse is RUNNING) on both the user and + * service-principal client slots; override any field via `overrides`. + */ +export function createMockRequest(overrides: Any = {}) { + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }), + }, + // Analytics route now calls `warehouses.get` before issuing SQL to + // ensure the warehouse is RUNNING. Default to RUNNING so existing + // tests that only care about SQL behaviour aren't affected. + warehouses: { + get: vi.fn().mockResolvedValue({ state: "RUNNING" }), + start: vi.fn().mockResolvedValue(undefined), + }, + }; + + const req = { + params: {}, + query: {}, + body: {}, + headers: {}, + userWorkspaceClient: mockWorkspaceClient, + serviceWorkspaceClient: mockWorkspaceClient, + getWarehouseId: vi.fn().mockResolvedValue("test-warehouse-id"), + getWorkspaceId: vi.fn().mockResolvedValue("test-workspace-id"), + header: function (name: string) { + return this.headers[name.toLowerCase()]; + }, + ...overrides, + }; + return req; +} + +/** + * Creates a mock Express response object. `write`/`send`/`setHeader` flip + * `headersSent`, `end` flips `writableEnded` and fires any `close` listener — + * enough for streaming handlers that branch on those flags. + */ +export function createMockResponse() { + const eventListeners: Record void>> = {}; + + const res = { + // Flips to true once headers/body have gone out — mirrors Express so + // streaming handlers can branch between a JSON error (pre-headers) and + // aborting the socket (mid-stream). + headersSent: false, + status: vi.fn().mockReturnThis(), + json: vi.fn().mockReturnThis(), + send: vi.fn(function (this: Any) { + this.headersSent = true; + return this; + }), + sendStatus: vi.fn().mockReturnThis(), + end: vi.fn(function (this: Any) { + this.writableEnded = true; + // Trigger 'close' event when end is called + if (eventListeners.close) { + for (const handler of eventListeners.close) { + handler(); + } + } + return this; + }), + write: vi.fn(function (this: Any) { + this.headersSent = true; + return this; + }), + setHeader: vi.fn(function (this: Any) { + this.headersSent = true; + return this; + }), + flushHeaders: vi.fn().mockReturnThis(), + destroy: vi.fn().mockReturnThis(), + on: vi.fn(function ( + this: Any, + event: string, + handler: (...args: Any[]) => void, + ) { + if (!eventListeners[event]) { + eventListeners[event] = []; + } + eventListeners[event].push(handler); + return this; + }), + off: vi.fn(function ( + this: Any, + event: string, + handler: (...args: Any[]) => void, + ) { + if (eventListeners[event]) { + eventListeners[event] = eventListeners[event].filter( + (h) => h !== handler, + ); + } + return this; + }), + writableEnded: false, + }; + return res; +} + +/** + * Sets up common environment variables for Databricks testing so code that + * reads `DATABRICKS_HOST` / `DATABRICKS_WAREHOUSE_ID` finds test values. + */ +export function setupDatabricksEnv(overrides: Record = {}) { + process.env.DATABRICKS_HOST = "https://test.databricks.com"; + process.env.DATABRICKS_WAREHOUSE_ID = "test-warehouse-id"; + Object.assign(process.env, overrides); +} + +/** + * Context options for running tests with mocked service/user context + */ +export interface TestContextOptions { + /** Mock WorkspaceClient for service principal operations */ + serviceDatabricksClient?: Any; + /** Mock WorkspaceClient for user operations */ + userDatabricksClient?: Any; + /** User ID for user context */ + userId?: string; + /** Service user ID */ + serviceUserId?: string; + /** Warehouse ID */ + warehouseId?: string; + /** Workspace ID */ + workspaceId?: string; +} + +/** + * Creates a default mock WorkspaceClient for testing (SQL succeeds, warehouse + * RUNNING). + */ +export function createMockWorkspaceClient() { + return { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }), + }, + // Analytics route now calls `warehouses.get` before issuing SQL to + // ensure the warehouse is RUNNING. Default to RUNNING so existing + // tests that only care about SQL behaviour aren't affected. + warehouses: { + get: vi.fn().mockResolvedValue({ state: "RUNNING" }), + start: vi.fn().mockResolvedValue(undefined), + }, + }; +} + +/** + * Builds a {@link ServiceContextState} for testing without touching the + * singleton. Use with {@link mockServiceContext} to install it. + */ +export function createMockServiceContext(options: TestContextOptions = {}) { + const mockWorkspaceClient = createMockWorkspaceClient(); + + const serviceContext: ServiceContextState = { + client: (options.serviceDatabricksClient || mockWorkspaceClient) as Any, + serviceUserId: options.serviceUserId || "test-service-user", + warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), + workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), + }; + + return serviceContext; +} + +/** + * Creates a mock UserContext for testing. + */ +export function createMockUserContext( + options: TestContextOptions = {}, +): UserContext { + const mockWorkspaceClient = createMockWorkspaceClient(); + + return { + client: (options.userDatabricksClient || mockWorkspaceClient) as Any, + userId: options.userId || "test-user", + warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), + workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), + isUserContext: true, + }; +} + +/** + * Mocks the `ServiceContext` singleton for testing — spies `get`, + * `initialize`, `isInitialized`, and `createUserContext` so code that resolves + * the service principal or an on-behalf-of user context gets test doubles. + * Call in `beforeEach`; call the returned `restore()` in `afterEach`. + * + * @returns The mock context plus the spies and a `restore()` helper. + */ +export function mockServiceContext(options: TestContextOptions = {}) { + const serviceContext = createMockServiceContext(options); + + const getSpy = vi + .spyOn(ServiceContext, "get") + .mockReturnValue(serviceContext); + + const initSpy = vi + .spyOn(ServiceContext, "initialize") + .mockResolvedValue(serviceContext); + + const isInitializedSpy = vi + .spyOn(ServiceContext, "isInitialized") + .mockReturnValue(true); + + // Mock createUserContext to return a test user context + const createUserContextSpy = vi + .spyOn(ServiceContext, "createUserContext") + .mockImplementation((_token: string, userId: string, userName?: string) => { + const mockWorkspaceClient = createMockWorkspaceClient(); + return { + client: (options.userDatabricksClient || mockWorkspaceClient) as Any, + userId, + userName, + warehouseId: serviceContext.warehouseId, + workspaceId: serviceContext.workspaceId, + isUserContext: true, + }; + }); + + return { + serviceContext, + getSpy, + initSpy, + isInitializedSpy, + createUserContextSpy, + restore: () => { + getSpy.mockRestore(); + initSpy.mockRestore(); + isInitializedSpy.mockRestore(); + createUserContextSpy.mockRestore(); + }, + }; +} + +/** + * Runs a test function within a mocked service context: installs the mock, + * runs `fn`, and restores the singleton afterward. + */ +export async function runWithRequestContext( + fn: () => T | Promise, + context?: TestContextOptions, +): Promise { + const mocks = mockServiceContext(context); + + try { + return await fn(); + } finally { + mocks.restore(); + } +} + +/** + * Builds a SUCCEEDED SQL statement response with a synthetic statement id, + * `data_array` rows, and a manifest schema derived from `columns`. + */ +export function createSuccessfulSQLResponse( + data: Any[][], + columns: Array<{ name: string; type_name?: string }>, +) { + return { + status: { state: "SUCCEEDED" }, + statement_id: `stmt-${Date.now()}`, + result: { + data_array: data, + }, + manifest: { + schema: { + columns: columns.map((col) => ({ + name: col.name, + type_name: col.type_name ?? "STRING", + })), + }, + }, + }; +} + +/** Builds a FAILED SQL statement response carrying `errorMessage`. */ +export function createFailedSQLResponse(errorMessage: string) { + return { + status: { + state: "FAILED", + error: { + message: errorMessage, + }, + }, + statement_id: `stmt-${Date.now()}`, + }; +} + +/** + * A WorkspaceClient whose `executeStatement`/`getStatement` are bare `vi.fn()`s + * (no default resolution) so a test can script exactly what SQL returns. + * `warehouses.get` defaults to RUNNING. + */ +export function createConfigurableMockWorkspaceClient() { + const executeStatement = vi.fn(); + const getStatement = vi.fn(); + // Analytics route now calls `warehouses.get` before issuing SQL; default to + // RUNNING so callers that don't care about warehouse readiness don't have + // to wire it up. + const warehousesGet = vi.fn().mockResolvedValue({ state: "RUNNING" }); + const warehousesStart = vi.fn().mockResolvedValue(undefined); + + const client = { + statementExecution: { + executeStatement, + getStatement, + }, + warehouses: { + get: warehousesGet, + start: warehousesStart, + }, + }; + + return { + client, + mocks: { + executeStatement, + getStatement, + warehousesGet, + warehousesStart, + }, + }; +} diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts new file mode 100644 index 000000000..5c34a2d99 --- /dev/null +++ b/packages/appkit/src/testing/index.ts @@ -0,0 +1,66 @@ +/** + * @packageDocumentation + * + * `@databricks/appkit/testing` — test an AppKit app without a live workspace. + * + * The kit is deterministic and network-free: it wraps the real + * {@link PluginContext} with faked edges (mock telemetry, fake tool providers, + * a stubbed on-behalf-of path) so a plugin's real code paths — route + * buffering, tool dispatch, timeout composition, user scoping — run under test + * with no credentials. + * + * Two entry points: + * - {@link mockPluginContext} — build a real `PluginContext` with faked edges + * and attach it to a plugin. + * - {@link expectStream} — assert the ordered event types a stream emits. + * + * Plus the fixture helpers (`createMockRequest`, `mockServiceContext`, …) for + * wiring up requests, responses, and the service-principal singleton. + * + * @example + * ```ts + * import { mockPluginContext, expectStream } from "@databricks/appkit/testing"; + * + * const mock = mockPluginContext({ analytics: { query: fixtureRows } }); + * await mock.attach(agentsPlugin); + * await expectStream(agentsPlugin._handleStream(req, res)).toEmit( + * "tool_call", + * "message_delta", + * ); + * ``` + * + * @module + */ + +export { + expectStream, + parseSSEResponse, + type StreamAssertion, + type StreamEvent, + type StreamSource, +} from "./expect-stream"; +export { + createConfigurableMockWorkspaceClient, + createFailedSQLResponse, + createMockRequest, + createMockResponse, + createMockRouter, + createMockServiceContext, + createMockTelemetry, + createMockUserContext, + createMockWorkspaceClient, + createSuccessfulSQLResponse, + mockServiceContext, + runWithRequestContext, + setupDatabricksEnv, + type TestContextOptions, +} from "./fixtures"; +export { + type FakeProvider, + type FakeProviders, + type FakeToolResponse, + type MockPluginContext, + mockPluginContext, + type RecordedRoute, + type RecordedToolCall, +} from "./mock-plugin-context"; diff --git a/packages/appkit/src/testing/mock-plugin-context.ts b/packages/appkit/src/testing/mock-plugin-context.ts new file mode 100644 index 000000000..420ac56d9 --- /dev/null +++ b/packages/appkit/src/testing/mock-plugin-context.ts @@ -0,0 +1,291 @@ +import type express from "express"; +import type { + AgentToolDefinition, + BasePlugin, + IAppRequest, + ToolProvider, +} from "shared"; +import { CacheManager } from "../cache"; +import { InMemoryStorage } from "../cache/storage"; +import { PluginContext } from "../core/plugin-context"; +import type { Plugin } from "../plugin"; +import type { ITelemetry } from "../telemetry"; +import { createMockTelemetry } from "./fixtures"; + +/** + * A concrete (non-function) fake tool response — returned as-is. Covers the + * JSON-serializable shapes a tool call yields (rows, objects, primitives, + * nullish). A bare `unknown` is intentionally not used here: unioned with the + * function form below it would collapse to `unknown` and strip contextual + * types from the callback's parameters. + */ +type FakeToolValue = + | Record + | unknown[] + | string + | number + | boolean + | null + | undefined; + +/** + * A canned tool response. Either a static {@link FakeToolValue} returned + * as-is, or a function of the call arguments (and the abort signal + * `PluginContext.executeTool` composes) so a fake can assert on inputs or + * simulate slow/aborting work. Returning a promise is supported (the return + * type is intentionally `unknown`, which also covers `Promise<...>`). + */ +export type FakeToolResponse = + | FakeToolValue + | ((args: unknown, signal?: AbortSignal) => unknown); + +/** + * Fake connector responses, keyed by plugin name and then tool name: + * + * ```ts + * mockPluginContext({ analytics: { query: fixtureRows } }); + * ``` + * + * Each top-level key registers a fake {@link ToolProvider} under that plugin + * name; each inner key becomes a tool that returns the mapped response. + */ +export type FakeProviders = Record>; + +/** A single dispatch observed by a fake provider. */ +export interface RecordedToolCall { + /** Registered plugin name (the key in {@link FakeProviders}). */ + plugin: string; + /** Tool name passed to `executeAgentTool`. */ + tool: string; + /** Arguments the tool received. */ + args: unknown; + /** The abort signal `executeTool` composed (timeout ∘ caller). */ + signal?: AbortSignal; + /** + * Whether the call went through the on-behalf-of (`asUser`) path. `true` + * proves `PluginContext.executeTool` resolved the user scope rather than + * running as the service principal. + */ + asUser: boolean; +} + +/** A single route registered through the context's `addRoute`/`addMiddleware`. */ +export interface RecordedRoute { + method: string; + path: string; + /** + * The raw handlers as passed to `addRoute` — before `PluginContext` wraps + * them with `forwardAsyncErrors`. Recorded here so aliasing assertions + * ("both routes mount the same handler") can compare the original + * references, which the wrapped express-level handlers no longer share. + */ + handlers: express.RequestHandler[]; +} + +/** A fake tool provider registered on a mock context. */ +export interface FakeProvider { + /** Every `asUser(req)` the context resolved for this provider. */ + asUserRequests: express.Request[]; + /** Definitions returned from `getAgentTools()`. */ + tools: AgentToolDefinition[]; +} + +/** + * The result of {@link mockPluginContext}: the real `PluginContext` plus the + * seams a test needs to drive and inspect it. + */ +export interface MockPluginContext { + /** The real {@link PluginContext}, constructed with mock telemetry. */ + ctx: PluginContext; + /** The injected mock telemetry provider — assert on spans here. */ + telemetry: ITelemetry; + /** + * Tool dispatches observed across all fake providers, in call order. Live — + * read it after the action under test runs. + */ + toolCalls: RecordedToolCall[]; + /** + * Routes registered through the context, in registration order. Live — + * populated when the plugin calls `addRoute`/`addMiddleware`. + */ + routes: RecordedRoute[]; + /** Fake providers by plugin name, for direct assertions. */ + providers: Map; + /** + * Register (or replace) a fake tool provider after construction. + * Same shape as one {@link FakeProviders} entry. + */ + registerProvider(name: string, tools: Record): void; + /** + * Attach this context to a plugin the production way: seed an in-memory + * cache (if AppKit hasn't already), then call `plugin.attachContext`, which + * also rebuilds the plugin's telemetry and flips `isReady` to `true`. Await + * it before exercising handlers that read `this.context`, `this.cache`, or + * gate on `isReady`. Returns the same plugin for chaining. + */ + attach

(plugin: P): Promise

; +} + +/** + * Build a real {@link PluginContext} with faked edges for testing — no live + * workspace, no OpenTelemetry pipeline, no network. + * + * The context is the *real* class, so route buffering, the tool registry, + * timeout composition, and the on-behalf-of (`asUser`) path all run for real. + * Only three edges are faked, matching the seams the class actually has: + * + * - **Telemetry** is a mock provider (the one injectable production seam). + * - **Tool providers** are fakes registered through the existing public + * `registerToolProvider`; their `asUser`/`executeAgentTool` are recorded. + * - **Routes** are captured by wrapping the public `addRoute`/`addMiddleware`. + * + * Nothing about `PluginContext` is reimplemented. + * + * @param fakes - Canned tool responses keyed by plugin then tool name. + * + * @example + * ```ts + * const mock = mockPluginContext({ analytics: { query: fixtureRows } }); + * await mock.attach(agentsPlugin); + * // ...exercise a handler that dispatches analytics.query... + * expect(mock.toolCalls[0]).toMatchObject({ plugin: "analytics", asUser: true }); + * ``` + */ +export function mockPluginContext( + fakes: FakeProviders = {}, +): MockPluginContext { + const telemetry = createMockTelemetry(); + const ctx = new PluginContext({ telemetry }); + + const toolCalls: RecordedToolCall[] = []; + const routes: RecordedRoute[] = []; + const providers = new Map(); + + // Wrap the public route API so raw (pre-wrap) handlers are inspectable while + // the real buffering/flush path stays intact. + const realAddRoute = ctx.addRoute.bind(ctx); + ctx.addRoute = ( + method: string, + path: string, + ...handlers: express.RequestHandler[] + ): void => { + routes.push({ method, path, handlers }); + realAddRoute(method, path, ...handlers); + }; + const realAddMiddleware = ctx.addMiddleware.bind(ctx); + ctx.addMiddleware = ( + path: string, + ...handlers: express.RequestHandler[] + ): void => { + routes.push({ method: "use", path, handlers }); + realAddMiddleware(path, ...handlers); + }; + + function registerProvider( + name: string, + tools: Record, + ): void { + const record: FakeProvider = { + asUserRequests: [], + tools: Object.keys(tools).map((toolName) => ({ + name: toolName, + description: `Fake tool ${name}.${toolName}`, + parameters: { type: "object" }, + })), + }; + providers.set(name, record); + + const resolve = async ( + toolName: string, + args: unknown, + signal: AbortSignal | undefined, + asUser: boolean, + ): Promise => { + toolCalls.push({ plugin: name, tool: toolName, args, signal, asUser }); + const response = tools[toolName]; + if (response === undefined) { + throw new Error( + `mockPluginContext: plugin "${name}" has no fake tool "${toolName}". ` + + `Available: ${Object.keys(tools).join(", ") || "(none)"}`, + ); + } + return typeof response === "function" + ? await (response as (a: unknown, s?: AbortSignal) => unknown)( + args, + signal, + ) + : response; + }; + + const base: ToolProvider = { + getAgentTools: () => record.tools, + executeAgentTool: (toolName, args, signal) => + resolve(toolName, args, signal, false), + }; + + // `asUser(req)` returns a user-scoped view whose executeAgentTool records + // that the OBO path ran — this is how executeTool's user scoping becomes + // observable without a real user token. + const asUser = (req: IAppRequest): ToolProvider => { + record.asUserRequests.push(req as express.Request); + return { + getAgentTools: () => record.tools, + executeAgentTool: (toolName, args, signal) => + resolve(toolName, args, signal, true), + }; + }; + + // `registerToolProvider` expects the full ToolProviderPlugin shape + // (BasePlugin & ToolProvider & { asUser }). executeTool only ever calls + // `asUser` and `executeAgentTool`; the remaining BasePlugin surface is + // never touched for a registered provider, so a focused fake plus a cast + // is sufficient and avoids reimplementing a plugin. + const provider = { + name, + setup: async () => {}, + injectRoutes: () => {}, + getEndpoints: () => ({}), + ...base, + asUser, + } as unknown as BasePlugin & + ToolProvider & { + asUser: (req: IAppRequest) => ToolProvider; + }; + + ctx.registerToolProvider(name, provider); + } + + for (const [name, tools] of Object.entries(fakes)) { + registerProvider(name, tools); + } + + async function attach

(plugin: P): Promise

{ + // Seed a real in-memory cache if AppKit hasn't initialized one. Idempotent: + // getInstance returns any existing singleton (e.g. one a suite already set + // up) and ignores the storage argument in that case. + if (!cacheReady()) { + await CacheManager.getInstance({ storage: new InMemoryStorage({}) }); + } + plugin.attachContext({ context: ctx }); + return plugin; + } + + return { + ctx, + telemetry, + toolCalls, + routes, + providers, + registerProvider, + attach, + }; +} + +function cacheReady(): boolean { + try { + CacheManager.getInstanceSync(); + return true; + } catch { + return false; + } +} diff --git a/packages/appkit/src/testing/tests/expect-stream.test.ts b/packages/appkit/src/testing/tests/expect-stream.test.ts new file mode 100644 index 000000000..3d1ba7bf5 --- /dev/null +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -0,0 +1,142 @@ +import { describe, expect, test } from "vitest"; +import { expectStream, parseSSEResponse } from "../expect-stream"; + +async function* asyncEvents(events: T[]): AsyncGenerator { + for (const event of events) { + yield event; + } +} + +/** Build a minimal SSE Response body from event frames. */ +function sseResponse( + frames: Array<{ event: string; data: unknown }>, +): Response { + const body = frames + .map( + (f, i) => + `id: ${i}\nevent: ${f.event}\ndata: ${JSON.stringify(f.data)}\n\n`, + ) + .join(""); + return new Response(body, { + headers: { "Content-Type": "text/event-stream" }, + }); +} + +describe("expectStream — async iterables (adapter output)", () => { + test("toEmit matches an in-order subsequence, ignoring interleaved events", async () => { + const stream = asyncEvents([ + { type: "metadata", data: { threadId: "t" } }, + { type: "tool_call", name: "highlight" }, + { type: "tool_result", output: "ok" }, + { type: "message_delta", content: "done" }, + ]); + + const types = await expectStream(stream).toEmit( + "tool_call", + "message_delta", + ); + expect(types).toEqual([ + "metadata", + "tool_call", + "tool_result", + "message_delta", + ]); + }); + + test("toEmit rejects when an expected type is missing", async () => { + const stream = asyncEvents([{ type: "message_delta" }]); + await expect(expectStream(stream).toEmit("tool_call")).rejects.toThrow( + /expected events.*tool_call.*in order/s, + ); + }); + + test("toEmit rejects when order is wrong", async () => { + const stream = asyncEvents([ + { type: "message_delta" }, + { type: "tool_call" }, + ]); + await expect( + expectStream(stream).toEmit("tool_call", "message_delta"), + ).rejects.toThrow(/in order/); + }); + + test("toEmitExactly requires the precise sequence", async () => { + const events = [{ type: "a" }, { type: "b" }]; + await expect( + expectStream(asyncEvents(events)).toEmitExactly("a", "b"), + ).resolves.toEqual(["a", "b"]); + await expect( + expectStream(asyncEvents(events)).toEmitExactly("a"), + ).rejects.toThrow(/exactly/); + }); + + test("collect and collectTypes return raw events and types", async () => { + const events = [ + { type: "x", n: 1 }, + { type: "y", n: 2 }, + ]; + const assertion = expectStream(events); + expect(await assertion.collectTypes()).toEqual(["x", "y"]); + expect(await assertion.collect()).toEqual(events); + }); +}); + +describe("expectStream — sync iterables", () => { + test("accepts a plain array of events", async () => { + await expect( + expectStream([{ type: "one" }, { type: "two" }]).toEmit("one", "two"), + ).resolves.toBeDefined(); + }); +}); + +describe("expectStream — SSE Response", () => { + test("parses event frames and asserts order", async () => { + const res = sseResponse([ + { event: "warehouse_status", data: { state: "RUNNING" } }, + { event: "result", data: { rows: [] } }, + ]); + await expect( + expectStream(res).toEmit("warehouse_status", "result"), + ).resolves.toEqual(["warehouse_status", "result"]); + }); + + test("accepts a Promise", async () => { + const res = Promise.resolve( + sseResponse([{ event: "result", data: { ok: true } }]), + ); + const events = await expectStream(res).collect(); + expect(events[0]).toMatchObject({ type: "result", ok: true }); + }); + + test("ignores heartbeat/comment lines", async () => { + const body = `: heartbeat\n\nid: 0\nevent: result\ndata: {"ok":true}\n\n`; + const res = new Response(body); + await expect(expectStream(res).toEmitExactly("result")).resolves.toEqual([ + "result", + ]); + }); +}); + +describe("expectStream — invalid source", () => { + test("throws for a non-stream value", async () => { + await expect( + // biome-ignore lint/suspicious/noExplicitAny: intentionally wrong type + expectStream(42 as any).collect(), + ).rejects.toThrow(/async iterable, an iterable, or a Response/); + }); +}); + +describe("parseSSEResponse — single-event helper", () => { + test("returns eventType plus parsed data fields", async () => { + const res = new Response( + `event: result\ndata: ${JSON.stringify({ value: 42 })}\n\n`, + ); + const parsed = await parseSSEResponse(res); + expect(parsed).toEqual({ eventType: "result", value: 42 }); + }); + + test("throws when no data line is present", async () => { + const res = new Response(`event: result\n\n`); + await expect(parseSSEResponse(res)).rejects.toThrow(/No data found/); + }); +}); diff --git a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts b/packages/appkit/src/testing/tests/mock-plugin-context.test.ts new file mode 100644 index 000000000..6b3db5416 --- /dev/null +++ b/packages/appkit/src/testing/tests/mock-plugin-context.test.ts @@ -0,0 +1,164 @@ +import type express from "express"; +import { describe, expect, test } from "vitest"; +import { PluginContext } from "../../core/plugin-context"; +import { mockPluginContext } from "../mock-plugin-context"; + +/** + * Contract for `mockPluginContext`. The point of the kit is that it wraps the + * REAL PluginContext — so these tests drive the real `executeTool`, + * `addRoute`, and `getToolProviders` and assert the observable seams (OBO, + * timeout, route recording) rather than a reimplementation. + */ + +function mockReq(headers: Record = {}): express.Request { + return { + body: {}, + headers, + header: (name: string) => headers[name.toLowerCase()], + } as unknown as express.Request; +} + +describe("mockPluginContext — construction", () => { + test("produces a real PluginContext instance", () => { + const { ctx } = mockPluginContext(); + expect(ctx).toBeInstanceOf(PluginContext); + }); + + test("registers fake providers passed at construction", () => { + const { ctx } = mockPluginContext({ + analytics: { query: [{ id: 1 }] }, + genie: { ask: "hi" }, + }); + const names = ctx.getToolProviders().map((p) => p.name); + expect(names).toContain("analytics"); + expect(names).toContain("genie"); + }); +}); + +describe("mockPluginContext — executeTool runs the REAL user-scoping path", () => { + test("dispatches through asUser and returns the canned static response", async () => { + const rows = [{ user: "alice", n: 3 }]; + const mock = mockPluginContext({ analytics: { top_users: rows } }); + + const result = await mock.ctx.executeTool( + mockReq(), + "analytics", + "top_users", + { limit: 10 }, + ); + + expect(result).toEqual(rows); + // executeTool always resolves the user scope via provider.asUser(req). + expect(mock.toolCalls).toHaveLength(1); + expect(mock.toolCalls[0]).toMatchObject({ + plugin: "analytics", + tool: "top_users", + args: { limit: 10 }, + asUser: true, + }); + // The OBO request object is the one we passed in. + expect(mock.providers.get("analytics")?.asUserRequests).toHaveLength(1); + }); + + test("invokes a function response with the args and the composed signal", async () => { + const mock = mockPluginContext({ + analytics: { + query: (args, signal) => ({ echoed: args, aborted: signal?.aborted }), + }, + }); + + const result = await mock.ctx.executeTool(mockReq(), "analytics", "query", { + sql: "SELECT 1", + }); + + expect(result).toEqual({ echoed: { sql: "SELECT 1" }, aborted: false }); + // executeTool composes a timeout signal even when the caller passes none. + expect(mock.toolCalls[0]?.signal).toBeInstanceOf(AbortSignal); + }); + + test("forwards the caller timeout so a slow tool is aborted", async () => { + const mock = mockPluginContext({ + slow: { + wait: (_args, signal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject(new Error("aborted by timeout")), + ); + }), + }, + }); + + // 5ms timeout — the tool never resolves on its own, so the composed + // timeout signal must fire. This exercises executeTool's real + // AbortSignal.timeout + AbortSignal.any composition. + await expect( + mock.ctx.executeTool(mockReq(), "slow", "wait", {}, undefined, 5), + ).rejects.toThrow(/aborted by timeout/); + }); + + test("throws with a helpful message for an unknown plugin", async () => { + const mock = mockPluginContext({ analytics: { query: [] } }); + await expect( + mock.ctx.executeTool(mockReq(), "nope", "query", {}), + ).rejects.toThrow(/unknown plugin "nope"/); + }); + + test("throws with a helpful message for an unknown tool", async () => { + const mock = mockPluginContext({ analytics: { query: [] } }); + await expect( + mock.ctx.executeTool(mockReq(), "analytics", "missing", {}), + ).rejects.toThrow(/no fake tool "missing"/); + }); +}); + +describe("mockPluginContext — telemetry seam", () => { + test("records a span on the injected mock telemetry for each executeTool", async () => { + const mock = mockPluginContext({ analytics: { query: [] } }); + const tracer = mock.telemetry.getTracer(); + + await mock.ctx.executeTool(mockReq(), "analytics", "query", {}); + + // getTracer() is called inside executeTool; startActiveSpan drives the span. + expect(tracer.startActiveSpan).toHaveBeenCalled(); + }); +}); + +describe("mockPluginContext — route recording", () => { + test("records addRoute calls with raw (pre-wrap) handlers", () => { + const mock = mockPluginContext(); + const handler: express.RequestHandler = (_req, res) => { + res.end(); + }; + + mock.ctx.addRoute("post", "/invocations", handler); + mock.ctx.addRoute("post", "/responses", handler); + + expect(mock.routes).toHaveLength(2); + expect(mock.routes[0]).toMatchObject({ + method: "post", + path: "/invocations", + }); + // Raw handler references are preserved (PluginContext would otherwise wrap + // them with forwardAsyncErrors, losing reference identity). + expect(mock.routes[0]?.handlers[0]).toBe(handler); + expect(mock.routes[1]?.handlers[0]).toBe(handler); + }); + + test("records addMiddleware under the 'use' method", () => { + const mock = mockPluginContext(); + const mw: express.RequestHandler = (_req, _res, next) => next(); + mock.ctx.addMiddleware("/api", mw); + expect(mock.routes).toEqual([ + { method: "use", path: "/api", handlers: [mw] }, + ]); + }); +}); + +describe("mockPluginContext — registerProvider after construction", () => { + test("adds a provider dynamically", async () => { + const mock = mockPluginContext(); + mock.registerProvider("late", { ping: "pong" }); + const result = await mock.ctx.executeTool(mockReq(), "late", "ping", {}); + expect(result).toBe("pong"); + }); +}); diff --git a/packages/appkit/tsdown.config.ts b/packages/appkit/tsdown.config.ts index f5ae00475..679b0c607 100644 --- a/packages/appkit/tsdown.config.ts +++ b/packages/appkit/tsdown.config.ts @@ -9,7 +9,7 @@ export default defineConfig([ excludeEntrypoints: ["./type-generator"], }, name: "@databricks/appkit", - entry: ["src/index.ts", "src/beta.ts"], + entry: ["src/index.ts", "src/beta.ts", "src/testing/index.ts"], outDir: "dist", hash: false, format: "esm", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2e45b03b5..ed122fb5f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -344,6 +344,9 @@ importers: vite: specifier: npm:rolldown-vite@7.1.14 version: rolldown-vite@7.1.14(@types/node@25.2.3)(esbuild@0.25.10)(jiti@2.6.1)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) + vitest: + specifier: '>=1.0.0' + version: 3.2.4(@types/debug@4.1.12)(@types/node@25.2.3)(jiti@2.6.1)(jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6))(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) ws: specifier: 8.21.0 version: 8.21.0(bufferutil@4.0.9) @@ -17608,6 +17611,14 @@ snapshots: optionalDependencies: vite: 7.2.4(@types/node@24.7.2)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) + '@vitest/mocker@3.2.4(vite@7.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2))': + dependencies: + '@vitest/spy': 3.2.4 + estree-walker: 3.0.3 + magic-string: 0.30.21 + optionalDependencies: + vite: 7.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) + '@vitest/pretty-format@3.2.4': dependencies: tinyrainbow: 2.0.0 @@ -24742,6 +24753,27 @@ snapshots: - tsx - yaml + vite-node@3.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2): + dependencies: + cac: 6.7.14 + debug: 4.4.3 + es-module-lexer: 1.7.0 + pathe: 2.0.3 + vite: 7.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) + transitivePeerDependencies: + - '@types/node' + - jiti + - less + - lightningcss + - sass + - sass-embedded + - stylus + - sugarss + - supports-color + - terser + - tsx + - yaml + vite-tsconfig-paths@5.1.4(typescript@5.9.3)(vite@7.2.4(@types/node@24.7.2)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2)): dependencies: debug: 4.4.3 @@ -24787,6 +24819,23 @@ snapshots: tsx: 4.20.6 yaml: 2.8.2 + vite@7.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2): + dependencies: + esbuild: 0.25.10 + fdir: 6.5.0(picomatch@4.0.3) + picomatch: 4.0.3 + postcss: 8.5.6 + rollup: 4.52.4 + tinyglobby: 0.2.15 + optionalDependencies: + '@types/node': 25.2.3 + fsevents: 2.3.3 + jiti: 2.6.1 + lightningcss: 1.30.2 + terser: 5.44.1 + tsx: 4.20.6 + yaml: 2.8.2 + vitest@3.2.4(@types/debug@4.1.12)(@types/node@24.7.2)(jiti@2.6.1)(jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6))(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2): dependencies: '@types/chai': 5.2.2 @@ -24830,6 +24879,49 @@ snapshots: - tsx - yaml + vitest@3.2.4(@types/debug@4.1.12)(@types/node@25.2.3)(jiti@2.6.1)(jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6))(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2): + dependencies: + '@types/chai': 5.2.2 + '@vitest/expect': 3.2.4 + '@vitest/mocker': 3.2.4(vite@7.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2)) + '@vitest/pretty-format': 3.2.4 + '@vitest/runner': 3.2.4 + '@vitest/snapshot': 3.2.4 + '@vitest/spy': 3.2.4 + '@vitest/utils': 3.2.4 + chai: 5.3.3 + debug: 4.4.3 + expect-type: 1.2.2 + magic-string: 0.30.21 + pathe: 2.0.3 + picomatch: 4.0.3 + std-env: 3.10.0 + tinybench: 2.9.0 + tinyexec: 0.3.2 + tinyglobby: 0.2.15 + tinypool: 1.1.1 + tinyrainbow: 2.0.0 + vite: 7.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) + vite-node: 3.2.4(@types/node@25.2.3)(jiti@2.6.1)(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) + why-is-node-running: 2.3.0 + optionalDependencies: + '@types/debug': 4.1.12 + '@types/node': 25.2.3 + jsdom: 27.0.0(bufferutil@4.0.9)(postcss@8.5.6) + transitivePeerDependencies: + - jiti + - less + - lightningcss + - msw + - sass + - sass-embedded + - stylus + - sugarss + - supports-color + - terser + - tsx + - yaml + vscode-jsonrpc@8.2.0: {} vscode-languageserver-protocol@3.17.5: From 55859c802d71d83274bf28b647b6d5190a99a8b0 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:37:43 +0200 Subject: [PATCH 03/22] test(appkit): migrate route-handler-errors context stub to mockPluginContext MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the second and final hand-rolled stub — (plugin as any).context = { addRoute } — with the real PluginContext from mockPluginContext. The kit's route recorder captures raw handlers, so the alias assertion (both /invocations and /responses mount the same handler reference) holds against the real class, where forwardAsyncErrors wrapping would otherwise break reference identity. Both context stubs the plan identified are now migrated. Signed-off-by: Galymzhan --- .../agents/tests/route-handler-errors.test.ts | 28 ++++++++----------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 2fc493ef4..547e512f5 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -1,6 +1,7 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; +import { mockPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; /** @@ -396,27 +397,22 @@ describe("POST /invocations & /responses — successful invoke", () => { describe("/invocations and /responses are aliases", () => { test("both routes are registered and bound to the same handler", () => { const plugin = new AgentsPlugin({ dir: false }); - const addRoute = vi.fn(); - // biome-ignore lint/suspicious/noExplicitAny: inject minimal fake context - (plugin as any).context = { addRoute }; + // Attach the real PluginContext via the testing kit. Its route recorder + // captures the RAW handlers passed to addRoute — the aliasing assertion + // needs the original references, which the context's forwardAsyncErrors + // wrapping would otherwise break. + const mock = mockPluginContext(); + // biome-ignore lint/suspicious/noExplicitAny: attach the real context + (plugin as any).context = mock.ctx; // biome-ignore lint/suspicious/noExplicitAny: invoke private mounter (plugin as any).mountInvokeRoutes(); - expect(addRoute).toHaveBeenCalledTimes(2); - const calls = addRoute.mock.calls.map((c: unknown[]) => ({ - method: c[0], - path: c[1], - handler: c[2], - })); - const invocations = calls.find( - (c: { path: unknown }) => c.path === "/invocations", - ); - const responses = calls.find( - (c: { path: unknown }) => c.path === "/responses", - ); + expect(mock.routes).toHaveLength(2); + const invocations = mock.routes.find((r) => r.path === "/invocations"); + const responses = mock.routes.find((r) => r.path === "/responses"); expect(invocations?.method).toBe("post"); expect(responses?.method).toBe("post"); // The two routes are aliases — same handler reference is mounted on both. - expect(invocations?.handler).toBe(responses?.handler); + expect(invocations?.handlers[0]).toBe(responses?.handlers[0]); }); }); From 4d37d5e95ba70611669a947f0845ce6c8777a9dc Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:42:50 +0200 Subject: [PATCH 04/22] docs(appkit): document the testing kit and ship a template example test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add docs/docs/development/testing.md covering mockPluginContext(), expectStream(), and the fixture helpers, with a full end-to-end example. Cross-links to local-development, custom-plugins, and execution-context. - Add template/server/example.test.ts: a self-contained, plugin-agnostic example that scaffolded apps ship with — it defines a tiny custom plugin and exercises both mockPluginContext (route recording) and expectStream (ordered event assertions), running with no workspace or network. Ships the kit to users, satisfying the plan's acceptance criteria that a docs page exists and the template carries at least one example test. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 151 +++++++++++++++++++++++++++++++ template/server/example.test.ts | 69 ++++++++++++++ 2 files changed, 220 insertions(+) create mode 100644 docs/docs/development/testing.md create mode 100644 template/server/example.test.ts diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md new file mode 100644 index 000000000..be1f3fb5f --- /dev/null +++ b/docs/docs/development/testing.md @@ -0,0 +1,151 @@ +--- +sidebar_position: 7 +--- + +# Testing + +AppKit ships a testing kit at `@databricks/appkit/testing` so you can test a plugin — including its cross-plugin tool calls and streaming responses — without a live Databricks workspace, credentials, or network access. That makes plugin tests fast and lets them run in CI, where no workspace is available. + +## Goal + +Exercise a plugin's real code paths — route registration, cross-plugin tool dispatch, user-scoped (on-behalf-of) execution, and per-call timeouts — against a real `PluginContext` with only its outer edges faked. Nothing about the context is reimplemented, so a test can't drift from production behavior. + +The kit has two entry points plus a set of fixture helpers: + +- **`mockPluginContext()`** — build a real `PluginContext` with faked edges and attach it to a plugin. +- **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. +- **Fixtures** — `createMockRequest`, `createMockResponse`, `mockServiceContext`, and SQL response builders. + +The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so it is declared as an optional peer dependency. Any project that runs Vitest already satisfies it; there is nothing extra to install. + +## `mockPluginContext()` + +`PluginContext` is the mediator AppKit passes to every plugin — it buffers routes, tracks tool providers, and runs cross-plugin tool calls with user scoping and a timeout. `mockPluginContext()` returns the **real** context with three edges faked: + +| Edge | How it's faked | +| --- | --- | +| Telemetry | A no-op mock provider — no OpenTelemetry pipeline needed. | +| Tool providers | Fakes registered through the real `registerToolProvider`, keyed by plugin then tool name. | +| Routes | The real `addRoute`/`addMiddleware` are wrapped to record what a plugin registers. | + +Because the context is real, `executeTool` still resolves the user scope via `asUser(req)` and still composes the abort signal from your timeout — so those paths are genuinely under test. + +### Registering fake tool responses + +Pass canned responses keyed by plugin name, then tool name. A response is either a static value or a function of the call arguments and the composed abort signal: + +```ts +import { mockPluginContext } from "@databricks/appkit/testing"; + +const mock = mockPluginContext({ + analytics: { + // static response + top_users: [{ user: "alice", events: 42 }], + // function response — assert on args, or simulate slow/aborting work + query: (args, signal) => runFakeQuery(args, signal), + }, +}); +``` + +### Attaching to a plugin + +`attach()` wires the context to a plugin the production way: it seeds an in-memory cache (if AppKit hasn't already initialized one), then calls the plugin's `attachContext`, which rebuilds telemetry and flips `isReady` to `true`. Await it before exercising any handler that reads `this.context`, `this.cache`, or gates on `isReady`: + +```ts +const plugin = agents({ dir: false }); +await mock.attach(plugin); +``` + +### Inspecting what happened + +The returned object exposes live views you read after the action under test runs: + +```ts +await someHandler(req, res); + +// Every cross-plugin tool dispatch, in order. +expect(mock.toolCalls[0]).toMatchObject({ + plugin: "analytics", + tool: "query", + asUser: true, // proves the on-behalf-of path ran +}); + +// Every route the plugin registered (raw handlers, before wrapping). +expect(mock.routes).toContainEqual( + expect.objectContaining({ method: "post", path: "/invocations" }), +); + +// The injected telemetry provider, for span assertions. +expect(mock.telemetry.getTracer().startActiveSpan).toHaveBeenCalled(); +``` + +`RecordedToolCall.asUser` is the high-value signal: it confirms the context routed the call through the user's identity rather than the service principal — a distinction that silent stubs cannot verify. + +## `expectStream(...)` + +AppKit plugins stream Server-Sent Events. `expectStream` consumes a stream and asserts the ordered event types it emits. It accepts an async iterable (an agent adapter's `run()`), a plain array of events, or an SSE `Response` (or a promise of one) whose body it parses. + +```ts +import { expectStream } from "@databricks/appkit/testing"; + +// In-order subsequence match — interleaved events (heartbeats, deltas) are ignored. +await expectStream(agent.adapter.run(input)).toEmit("tool_call", "message_delta"); + +// Exact match — the stream's full shape, in order, with nothing else. +await expectStream(events).toEmitExactly("warehouse_status", "result"); + +// Or collect without asserting. +const types = await expectStream(res).collectTypes(); +``` + +`toEmit` checks that the expected types appear **in order** but tolerates other events before, between, or after them — which is what you want for streams that interleave bookkeeping events like heartbeats or metadata. Use `toEmitExactly` when the stream's shape is fully determined. + +## Fixtures + +The kit re-exports the request/response/context fixtures AppKit uses internally: + +- `createMockRequest(overrides?)` / `createMockResponse()` — Express request/response doubles, including the streaming flags (`headersSent`, `writableEnded`) and a mock `WorkspaceClient`. +- `mockServiceContext(options?)` — spy the `ServiceContext` singleton so code that resolves the service principal or a user context gets test doubles. Call in `beforeEach`, and call the returned `restore()` in `afterEach`. +- `createSuccessfulSQLResponse(rows, columns)` / `createFailedSQLResponse(message)` — build SQL Warehouse statement responses. +- `setupDatabricksEnv(overrides?)` — set `DATABRICKS_HOST` / `DATABRICKS_WAREHOUSE_ID` to test values. + +## Full example + +```ts +import { describe, expect, test } from "vitest"; +import { analytics } from "@databricks/appkit"; +import { + createMockRequest, + createMockResponse, + mockPluginContext, +} from "@databricks/appkit/testing"; + +describe("analytics query route", () => { + test("streams warehouse status then the result", async () => { + const mock = mockPluginContext({ + analytics: { top_users: [{ user: "alice", events: 42 }] }, + }); + const plugin = analytics({}); + await mock.attach(plugin); + + const req = createMockRequest({ + params: { query_key: "top_users" }, + body: { format: "JSON_ARRAY" }, + }); + const res = createMockResponse(); + + await plugin._handleQueryRoute( + req as never, + res as never, + ); + + expect(res.status).not.toHaveBeenCalledWith(500); + }); +}); +``` + +## See also + +- [Local development](./local-development.mdx) — run your app with hot reload while iterating. +- [Custom plugins](../plugins/custom-plugins.md) — build the plugins you test with this kit. +- [Execution context](../plugins/execution-context.md) — how `asUser` and the service principal differ at runtime. diff --git a/template/server/example.test.ts b/template/server/example.test.ts new file mode 100644 index 000000000..bae100888 --- /dev/null +++ b/template/server/example.test.ts @@ -0,0 +1,69 @@ +import { Plugin, type PluginManifest, toPlugin } from '@databricks/appkit'; +import { + expectStream, + mockPluginContext, +} from '@databricks/appkit/testing'; +import { describe, expect, test } from 'vitest'; + +/** + * Example test using the AppKit testing kit (`@databricks/appkit/testing`). + * + * The kit lets you test a plugin with NO Databricks workspace, credentials, or + * network — so these tests run anywhere, including CI. Delete this file, or use + * it as a starting point for testing your own plugins. + * + * Two headline helpers are shown below: + * - `mockPluginContext()` — a real PluginContext with faked edges, attachable + * to a plugin so its real code paths (routes, tool dispatch, user scoping) + * run under test. + * - `expectStream(...).toEmit(...)` — assert the ordered event types a + * streaming handler emits. + */ + +// A tiny example plugin: it registers one route and streams two events. +class GreeterPlugin extends Plugin { + static manifest = { + name: 'greeter', + displayName: 'Greeter', + description: 'Example plugin for the testing-kit demo', + resources: { required: [], optional: [] }, + } as PluginManifest<'greeter'>; + + async setup() { + // Routes registered here are captured by mockPluginContext().routes. + this.context?.addRoute('get', '/hello', (_req, res) => { + res.end(); + }); + } + + // A stand-in for a streaming handler: yields SSE-style event objects. + async *greet(name: string) { + yield { type: 'greeting_start', name }; + yield { type: 'greeting_end', message: `Hello, ${name}!` }; + } +} + +const greeter = toPlugin(GreeterPlugin); + +describe('testing kit example', () => { + test('attaches a real PluginContext and records registered routes', async () => { + const mock = mockPluginContext(); + const plugin = greeter(); + + await mock.attach(plugin); + await plugin.setup(); + + expect(mock.routes).toContainEqual( + expect.objectContaining({ method: 'get', path: '/hello' }), + ); + }); + + test('asserts the ordered events a stream emits', async () => { + const plugin = greeter(); + + await expectStream(plugin.greet('world')).toEmit( + 'greeting_start', + 'greeting_end', + ); + }); +}); From 5821f6b7e4a864761ec66dc24df461f4afc82857 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:49:29 +0200 Subject: [PATCH 05/22] docs(appkit): fix testing-kit examples to instantiate the plugin class MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validation by scaffolding a real app with `databricks apps init` surfaced that the examples called the `analytics()`/`toPlugin()` factory and then treated the result as a plugin instance — but a factory returns a { plugin, config, name } descriptor for createApp to construct, so `.attachContext`/handler methods are absent. Rewrite both the template example test and the docs "Full example" to instantiate the plugin class directly (`new GreeterPlugin({})`), matching how the migrated agents suites use the kit. The scaffolded app's `npm test` and `tsc` both pass against the published `@databricks/appkit/testing` subpath with no workspace or network. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 78 ++++++++++++++++++++++---------- template/server/example.test.ts | 13 ++++-- 2 files changed, 63 insertions(+), 28 deletions(-) diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md index be1f3fb5f..03ccb0421 100644 --- a/docs/docs/development/testing.md +++ b/docs/docs/development/testing.md @@ -111,39 +111,71 @@ The kit re-exports the request/response/context fixtures AppKit uses internally: ## Full example +Instantiate the plugin **class** directly with `new`. The `analytics()` / `agents()` factory functions you pass to `createApp` return a descriptor for the app to construct — for a unit test you want the instance itself. + ```ts +import { Plugin, type PluginManifest } from "@databricks/appkit"; +import { expectStream, mockPluginContext } from "@databricks/appkit/testing"; import { describe, expect, test } from "vitest"; -import { analytics } from "@databricks/appkit"; -import { - createMockRequest, - createMockResponse, - mockPluginContext, -} from "@databricks/appkit/testing"; - -describe("analytics query route", () => { - test("streams warehouse status then the result", async () => { - const mock = mockPluginContext({ - analytics: { top_users: [{ user: "alice", events: 42 }] }, - }); - const plugin = analytics({}); - await mock.attach(plugin); - const req = createMockRequest({ - params: { query_key: "top_users" }, - body: { format: "JSON_ARRAY" }, - }); - const res = createMockResponse(); +// A small plugin that registers a route and streams two events. +class GreeterPlugin extends Plugin { + static manifest = { + name: "greeter", + displayName: "Greeter", + description: "Example plugin", + resources: { required: [], optional: [] }, + } as PluginManifest<"greeter">; + + async setup() { + this.context?.addRoute("get", "/hello", (_req, res) => res.end()); + } + + async *greet(name: string) { + yield { type: "greeting_start", name }; + yield { type: "greeting_end", message: `Hello, ${name}!` }; + } +} + +describe("greeter plugin", () => { + test("registers its route through the context", async () => { + const mock = mockPluginContext(); + const plugin = new GreeterPlugin({}); - await plugin._handleQueryRoute( - req as never, - res as never, + await mock.attach(plugin); + await plugin.setup(); + + expect(mock.routes).toContainEqual( + expect.objectContaining({ method: "get", path: "/hello" }), ); + }); - expect(res.status).not.toHaveBeenCalledWith(500); + test("streams events in order", async () => { + const plugin = new GreeterPlugin({}); + await expectStream(plugin.greet("world")).toEmit( + "greeting_start", + "greeting_end", + ); }); }); ``` +To test a plugin that dispatches cross-plugin tool calls, register fake providers and assert on `mock.toolCalls` — including `asUser`, which confirms the on-behalf-of path ran: + +```ts +const mock = mockPluginContext({ analytics: { query: [{ n: 1 }] } }); +const plugin = new MyAgentPlugin({ dir: false }); +await mock.attach(plugin); + +await plugin.runSomethingThatCallsAnalytics(req); + +expect(mock.toolCalls[0]).toMatchObject({ + plugin: "analytics", + tool: "query", + asUser: true, +}); +``` + ## See also - [Local development](./local-development.mdx) — run your app with hot reload while iterating. diff --git a/template/server/example.test.ts b/template/server/example.test.ts index bae100888..149569645 100644 --- a/template/server/example.test.ts +++ b/template/server/example.test.ts @@ -1,4 +1,4 @@ -import { Plugin, type PluginManifest, toPlugin } from '@databricks/appkit'; +import { Plugin, type PluginManifest } from '@databricks/appkit'; import { expectStream, mockPluginContext, @@ -18,6 +18,11 @@ import { describe, expect, test } from 'vitest'; * run under test. * - `expectStream(...).toEmit(...)` — assert the ordered event types a * streaming handler emits. + * + * Note: tests instantiate the plugin CLASS directly (`new GreeterPlugin()`). + * The `analytics()` / `agents()` factory functions you pass to `createApp` + * return a descriptor for the app to construct — for a unit test you want the + * instance itself. */ // A tiny example plugin: it registers one route and streams two events. @@ -43,12 +48,10 @@ class GreeterPlugin extends Plugin { } } -const greeter = toPlugin(GreeterPlugin); - describe('testing kit example', () => { test('attaches a real PluginContext and records registered routes', async () => { const mock = mockPluginContext(); - const plugin = greeter(); + const plugin = new GreeterPlugin({}); await mock.attach(plugin); await plugin.setup(); @@ -59,7 +62,7 @@ describe('testing kit example', () => { }); test('asserts the ordered events a stream emits', async () => { - const plugin = greeter(); + const plugin = new GreeterPlugin({}); await expectStream(plugin.greet('world')).toEmit( 'greeting_start', From cc790d07cfcf4fbb0a392a73fac35e6b2b88c995 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:51:53 +0200 Subject: [PATCH 06/22] refactor(appkit): tighten FakeToolResponse so a missing value is a type error Drop `undefined` from the static FakeToolValue union. `resolve()` treats an undefined map entry as "unregistered tool" and throws, so allowing undefined as a declared response made `{ query: undefined }` a confusing runtime error instead of a compile error. A function returning undefined still works for the rare "returns nothing" case. Add a test pinning that a null response is returned as a value, not misread as a missing tool. Signed-off-by: Galymzhan --- packages/appkit/src/testing/mock-plugin-context.ts | 3 +-- .../src/testing/tests/mock-plugin-context.test.ts | 13 +++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/packages/appkit/src/testing/mock-plugin-context.ts b/packages/appkit/src/testing/mock-plugin-context.ts index 420ac56d9..2f81d9734 100644 --- a/packages/appkit/src/testing/mock-plugin-context.ts +++ b/packages/appkit/src/testing/mock-plugin-context.ts @@ -25,8 +25,7 @@ type FakeToolValue = | string | number | boolean - | null - | undefined; + | null; /** * A canned tool response. Either a static {@link FakeToolValue} returned diff --git a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts b/packages/appkit/src/testing/tests/mock-plugin-context.test.ts index 6b3db5416..3c17b1236 100644 --- a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts +++ b/packages/appkit/src/testing/tests/mock-plugin-context.test.ts @@ -109,6 +109,19 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () mock.ctx.executeTool(mockReq(), "analytics", "missing", {}), ).rejects.toThrow(/no fake tool "missing"/); }); + + test("returns a null response as a value rather than treating it as missing", async () => { + // `resolve` distinguishes a null fake response (valid) from undefined + // (unregistered tool), so a tool can model an empty/absent result. + const mock = mockPluginContext({ analytics: { lookup: null } }); + const result = await mock.ctx.executeTool( + mockReq(), + "analytics", + "lookup", + {}, + ); + expect(result).toBeNull(); + }); }); describe("mockPluginContext — telemetry seam", () => { From bace753d71abc8aef0077b48ad567220f35f1721 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 13:55:50 +0200 Subject: [PATCH 07/22] refactor(appkit): make tools/test-helpers a shim over the shipped testing kit The plan's step 5 was to MOVE the fixtures into the package, not copy them. The shipped kit (src/testing/fixtures.ts) duplicated all 15 exports of tools/test-helpers.ts, which would drift over time. Collapse the original into a thin re-export of @databricks/appkit/testing so src/testing is the single source of truth while the 18 existing @tools/test-helpers importers keep working unchanged. The re-exported mockServiceContext is now synchronous; every call site either awaits it (no-op on a non-promise) or reads it through Awaited>, so all suites pass unchanged (full appkit suite: 3117 passed, 1 pre-existing skip). Signed-off-by: Galymzhan --- tools/test-helpers.ts | 472 +++--------------------------------------- 1 file changed, 27 insertions(+), 445 deletions(-) diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index 9161a0f67..63830f43b 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -1,447 +1,29 @@ -import type { Span, SpanOptions } from "@opentelemetry/api"; -import type { IAppRouter } from "shared"; -import { vi } from "vitest"; -import type { ServiceContextState } from "../packages/appkit/src/context/service-context"; -import type { UserContext } from "../packages/appkit/src/context/user-context"; -import type { - InstrumentConfig, - ITelemetry, -} from "../packages/appkit/src/telemetry/types"; - /** - * Creates a mock telemetry provider for testing - */ -export function createMockTelemetry(): ITelemetry { - const mockSpan: Span = { - addLink: vi.fn(), - addLinks: vi.fn(), - end: vi.fn(), - setAttribute: vi.fn(), - setAttributes: vi.fn(), - setStatus: vi.fn(), - recordException: vi.fn(), - updateName: vi.fn(), - addEvent: vi.fn(), - isRecording: vi.fn().mockReturnValue(false), - spanContext: vi.fn(), - }; - - return { - getTracer: vi.fn().mockReturnValue({ - startActiveSpan: vi.fn().mockImplementation((...args: any[]) => { - const fn = args[args.length - 1]; - if (typeof fn === "function") { - return fn(mockSpan); - } - return undefined; - }), - }), - getMeter: vi.fn().mockReturnValue({ - createCounter: vi.fn().mockReturnValue({ add: vi.fn() }), - createHistogram: vi.fn().mockReturnValue({ record: vi.fn() }), - }), - getLogger: vi.fn().mockReturnValue({ - emit: vi.fn(), - }), - emit: vi.fn(), - startActiveSpan: vi - .fn() - .mockImplementation( - async ( - _name: string, - _options: SpanOptions, - fn: (span: Span) => Promise, - _tracerOptions?: InstrumentConfig, - ) => { - return await fn(mockSpan); - }, - ), - registerInstrumentations: vi.fn(), - }; -} - -/** - * Creates a mock Express router with route handler capturing - */ -export function createMockRouter(): { - router: IAppRouter; - handlers: Record; - getHandler: (method: string, path: string) => any; -} { - const handlers: Record = {}; - - const mockRouter = { - get: vi.fn((path: string, handler: any) => { - handlers[`GET:${path}`] = handler; - }), - post: vi.fn((path: string, handler: any) => { - handlers[`POST:${path}`] = handler; - }), - put: vi.fn((path: string, handler: any) => { - handlers[`PUT:${path}`] = handler; - }), - delete: vi.fn((path: string, handler: any) => { - handlers[`DELETE:${path}`] = handler; - }), - patch: vi.fn((path: string, handler: any) => { - handlers[`PATCH:${path}`] = handler; - }), - } as unknown as IAppRouter; - - return { - router: mockRouter, - handlers, - getHandler: (method: string, path: string) => - handlers[`${method.toUpperCase()}:${path}`], - }; -} - -/** - * Creates a mock Express request object - */ -export function createMockRequest(overrides: any = {}) { - const mockWorkspaceClient = { - statementExecution: { - executeStatement: vi.fn().mockResolvedValue({ - status: { state: "SUCCEEDED" }, - result: { data: [] }, - }), - }, - // Analytics route now calls `warehouses.get` before issuing SQL to - // ensure the warehouse is RUNNING. Default to RUNNING so existing - // tests that only care about SQL behaviour aren't affected. - warehouses: { - get: vi.fn().mockResolvedValue({ state: "RUNNING" }), - start: vi.fn().mockResolvedValue(undefined), - }, - }; - - const req = { - params: {}, - query: {}, - body: {}, - headers: {}, - userWorkspaceClient: mockWorkspaceClient, - serviceWorkspaceClient: mockWorkspaceClient, - getWarehouseId: vi.fn().mockResolvedValue("test-warehouse-id"), - getWorkspaceId: vi.fn().mockResolvedValue("test-workspace-id"), - header: function (name: string) { - return this.headers[name.toLowerCase()]; - }, - ...overrides, - }; - return req; -} - -/** - * Creates a mock Express response object - */ -export function createMockResponse() { - const eventListeners: Record void>> = {}; - - const res = { - // Flips to true once headers/body have gone out — mirrors Express so - // streaming handlers can branch between a JSON error (pre-headers) and - // aborting the socket (mid-stream). - headersSent: false, - status: vi.fn().mockReturnThis(), - json: vi.fn().mockReturnThis(), - send: vi.fn(function (this: any) { - this.headersSent = true; - return this; - }), - sendStatus: vi.fn().mockReturnThis(), - end: vi.fn(function (this: any) { - this.writableEnded = true; - // Trigger 'close' event when end is called - if (eventListeners.close) { - for (const handler of eventListeners.close) { - handler(); - } - } - return this; - }), - write: vi.fn(function (this: any) { - this.headersSent = true; - return this; - }), - setHeader: vi.fn(function (this: any) { - this.headersSent = true; - return this; - }), - flushHeaders: vi.fn().mockReturnThis(), - destroy: vi.fn().mockReturnThis(), - on: vi.fn(function ( - this: any, - event: string, - handler: (...args: any[]) => void, - ) { - if (!eventListeners[event]) { - eventListeners[event] = []; - } - eventListeners[event].push(handler); - return this; - }), - off: vi.fn(function ( - this: any, - event: string, - handler: (...args: any[]) => void, - ) { - if (eventListeners[event]) { - eventListeners[event] = eventListeners[event].filter( - (h) => h !== handler, - ); - } - return this; - }), - writableEnded: false, - }; - return res; -} - -/** - * Sets up common environment variables for Databricks testing - */ -export function setupDatabricksEnv(overrides: Record = {}) { - process.env.DATABRICKS_HOST = "https://test.databricks.com"; - process.env.DATABRICKS_WAREHOUSE_ID = "test-warehouse-id"; - Object.assign(process.env, overrides); -} - -/** - * Context options for running tests with mocked service/user context - */ -export interface TestContextOptions { - /** Mock WorkspaceClient for service principal operations */ - serviceDatabricksClient?: any; - /** Mock WorkspaceClient for user operations */ - userDatabricksClient?: any; - /** User ID for user context */ - userId?: string; - /** Service user ID */ - serviceUserId?: string; - /** Warehouse ID */ - warehouseId?: string; - /** Workspace ID */ - workspaceId?: string; -} - -/** - * Creates a default mock WorkspaceClient for testing - */ -export function createMockWorkspaceClient() { - return { - statementExecution: { - executeStatement: vi.fn().mockResolvedValue({ - status: { state: "SUCCEEDED" }, - result: { data: [] }, - }), - }, - // Analytics route now calls `warehouses.get` before issuing SQL to - // ensure the warehouse is RUNNING. Default to RUNNING so existing - // tests that only care about SQL behaviour aren't affected. - warehouses: { - get: vi.fn().mockResolvedValue({ state: "RUNNING" }), - start: vi.fn().mockResolvedValue(undefined), - }, - }; -} - -/** - * Creates a mock ServiceContext for testing. - * Call this in beforeEach to set up the ServiceContext mock. - */ -export function createMockServiceContext(options: TestContextOptions = {}) { - const mockWorkspaceClient = createMockWorkspaceClient(); - - const serviceContext: ServiceContextState = { - client: (options.serviceDatabricksClient || mockWorkspaceClient) as any, - serviceUserId: options.serviceUserId || "test-service-user", - warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), - workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), - }; - - return serviceContext; -} - -/** - * Creates a mock UserContext for testing. - */ -export function createMockUserContext( - options: TestContextOptions = {}, -): UserContext { - const mockWorkspaceClient = createMockWorkspaceClient(); - - return { - client: (options.userDatabricksClient || mockWorkspaceClient) as any, - userId: options.userId || "test-user", - warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), - workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), - isUserContext: true, - }; -} - -/** - * Mocks the ServiceContext singleton for testing. - * Should be called in beforeEach. + * @deprecated Internal re-export shim. The test helpers now live in the + * shipped testing kit at `packages/appkit/src/testing/` and are published as + * `@databricks/appkit/testing`. This file re-exports them so the existing + * `@tools/test-helpers` importers keep working; new code (inside or outside + * this repo) should import from `@databricks/appkit/testing` instead. * - * @returns Object with spies that can be used to restore the mocks - */ -export async function mockServiceContext(options: TestContextOptions = {}) { - const serviceContext = createMockServiceContext(options); - - const contextModule = await import( - "../packages/appkit/src/context/service-context" - ); - - const getSpy = vi - .spyOn(contextModule.ServiceContext, "get") - .mockReturnValue(serviceContext); - - const initSpy = vi - .spyOn(contextModule.ServiceContext, "initialize") - .mockResolvedValue(serviceContext); - - const isInitializedSpy = vi - .spyOn(contextModule.ServiceContext, "isInitialized") - .mockReturnValue(true); - - // Mock createUserContext to return a test user context - const createUserContextSpy = vi - .spyOn(contextModule.ServiceContext, "createUserContext") - .mockImplementation((_token: string, userId: string, userName?: string) => { - const mockWorkspaceClient = createMockWorkspaceClient(); - return { - client: (options.userDatabricksClient || mockWorkspaceClient) as any, - userId, - userName, - warehouseId: serviceContext.warehouseId, - workspaceId: serviceContext.workspaceId, - isUserContext: true, - }; - }); - - return { - serviceContext, - getSpy, - initSpy, - isInitializedSpy, - createUserContextSpy, - restore: () => { - getSpy.mockRestore(); - initSpy.mockRestore(); - isInitializedSpy.mockRestore(); - createUserContextSpy.mockRestore(); - }, - }; -} - -/** - * Runs a test function within a mocked service context. - * This sets up the ServiceContext mock, runs the function, and restores the mock. - */ -export async function runWithRequestContext( - fn: () => T | Promise, - context?: TestContextOptions, -): Promise { - const mocks = await mockServiceContext(context); - - try { - return await fn(); - } finally { - mocks.restore(); - } -} - -/** - * Parses SSE response. Format: "event: result\ndata: {...}\n\n" - */ -export async function parseSSEResponse(response: Response): Promise { - const text = await response.text(); - const lines = text.split("\n"); - - let eventType: string | null = null; - let dataLine: string | null = null; - - for (const line of lines) { - if (line.startsWith("event: ")) { - eventType = line.substring(7).trim(); - } else if (line.startsWith("data: ")) { - dataLine = line.substring(6); - } - } - - if (!dataLine) { - throw new Error(`No data found in SSE response: ${text}`); - } - - const parsed = JSON.parse(dataLine); - return { - eventType, - ...parsed, - }; -} - -export function createConfigurableMockWorkspaceClient() { - const executeStatement = vi.fn(); - const getStatement = vi.fn(); - // Analytics route now calls `warehouses.get` before issuing SQL; default to - // RUNNING so callers that don't care about warehouse readiness don't have - // to wire it up. - const warehousesGet = vi.fn().mockResolvedValue({ state: "RUNNING" }); - const warehousesStart = vi.fn().mockResolvedValue(undefined); - - const client = { - statementExecution: { - executeStatement, - getStatement, - }, - warehouses: { - get: warehousesGet, - start: warehousesStart, - }, - }; - - return { - client, - mocks: { - executeStatement, - getStatement, - warehousesGet, - warehousesStart, - }, - }; -} - -export function createSuccessfulSQLResponse( - data: any[][], - columns: Array<{ name: string; type_name?: string }>, -) { - return { - status: { state: "SUCCEEDED" }, - statement_id: `stmt-${Date.now()}`, - result: { - data_array: data, - }, - manifest: { - schema: { - columns: columns.map((col) => ({ - name: col.name, - type_name: col.type_name ?? "STRING", - })), - }, - }, - }; -} - -export function createFailedSQLResponse(errorMessage: string) { - return { - status: { - state: "FAILED", - error: { - message: errorMessage, - }, - }, - statement_id: `stmt-${Date.now()}`, - }; -} + * Note: `mockServiceContext` is now synchronous (the previous dynamic + * `import()` became a static one to avoid a circular-init trap once packaged). + * Existing `await mockServiceContext(...)` call sites are unaffected — awaiting + * a non-promise is a no-op, and `Awaited>` unwraps identically. + */ +export { + createConfigurableMockWorkspaceClient, + createFailedSQLResponse, + createMockRequest, + createMockResponse, + createMockRouter, + createMockServiceContext, + createMockTelemetry, + createMockUserContext, + createMockWorkspaceClient, + createSuccessfulSQLResponse, + mockServiceContext, + parseSSEResponse, + runWithRequestContext, + setupDatabricksEnv, + type TestContextOptions, +} from "../packages/appkit/src/testing"; From cfde155f73c70f7f12d134cddd56a1c6e0c7f6f2 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Mon, 10 Aug 2026 14:09:59 +0200 Subject: [PATCH 08/22] fix(appkit): normalize CRLF in expectStream SSE parsing; sharpen testing docs Code review follow-ups: - expectStream's parseSSEBody split frames on \n\n, so a spec-compliant SSE stream delimited by \r\n\r\n (from a real server) collapsed into one event. AppKit's own writer uses \n\n so existing tests were unaffected, but expectStream is public API that accepts any Response. Normalize CRLF to LF before splitting; add a CRLF regression test. - Docs: instantiate the plugin CLASS in the attach() snippet (the factory returns a descriptor, not an instance), and note that the cache attach() seeds is a per-process singleton shared by tests within a file. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 6 +++++- packages/appkit/src/testing/expect-stream.ts | 7 ++++--- .../appkit/src/testing/tests/expect-stream.test.ts | 12 ++++++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md index 03ccb0421..21dbb1948 100644 --- a/docs/docs/development/testing.md +++ b/docs/docs/development/testing.md @@ -52,10 +52,14 @@ const mock = mockPluginContext({ `attach()` wires the context to a plugin the production way: it seeds an in-memory cache (if AppKit hasn't already initialized one), then calls the plugin's `attachContext`, which rebuilds telemetry and flips `isReady` to `true`. Await it before exercising any handler that reads `this.context`, `this.cache`, or gates on `isReady`: ```ts -const plugin = agents({ dir: false }); +const plugin = new MyAgentPlugin({ dir: false }); await mock.attach(plugin); ``` +Instantiate the plugin **class** directly (`new MyAgentPlugin(...)`). The `analytics()` / `agents()` factories you pass to `createApp` return a descriptor for the app to construct — for a unit test you want the instance. + +The cache `attach()` seeds is a process-wide singleton: `CacheManager` is initialized once per test process and reused. Vitest isolates test *files* in separate workers, so caches never leak across files, but tests **within one file** share it. If a test populates the cache and a later test in the same file must not see it, reset between tests (e.g. clear the cache in `beforeEach`). + ### Inspecting what happened The returned object exposes live views you read after the action under test runs: diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts index cc32de943..ca0873d2e 100644 --- a/packages/appkit/src/testing/expect-stream.ts +++ b/packages/appkit/src/testing/expect-stream.ts @@ -55,14 +55,15 @@ function eventType(event: StreamEvent): string { */ function parseSSEBody(text: string): StreamEvent[] { const events: StreamEvent[] = []; - const blocks = text.split(/\n\n/); + // Normalize CRLF to LF first so frames delimited by `\r\n\r\n` (spec-compliant + // SSE from a real server) split the same as AppKit's own `\n\n` writer. + const blocks = text.replace(/\r\n/g, "\n").split("\n\n"); for (const block of blocks) { let name: string | undefined; const dataLines: string[] = []; - for (const rawLine of block.split("\n")) { - const line = rawLine.replace(/\r$/, ""); + for (const line of block.split("\n")) { if (line.startsWith("event:")) { name = line.slice("event:".length).trim(); } else if (line.startsWith("data:")) { diff --git a/packages/appkit/src/testing/tests/expect-stream.test.ts b/packages/appkit/src/testing/tests/expect-stream.test.ts index 3d1ba7bf5..95c001840 100644 --- a/packages/appkit/src/testing/tests/expect-stream.test.ts +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -115,6 +115,18 @@ describe("expectStream — SSE Response", () => { "result", ]); }); + + test("parses CRLF-delimited frames from a spec-compliant SSE stream", async () => { + // A real server may use \r\n\r\n between frames; AppKit's own writer uses + // \n\n. Both must parse to distinct events, not one collapsed block. + const body = + 'event: warehouse_status\r\ndata: {"state":"RUNNING"}\r\n\r\n' + + 'event: result\r\ndata: {"rows":[]}\r\n\r\n'; + const res = new Response(body); + await expect( + expectStream(res).toEmitExactly("warehouse_status", "result"), + ).resolves.toEqual(["warehouse_status", "result"]); + }); }); describe("expectStream — invalid source", () => { From de5b40bf072c305d89c46bdff0a4d818a96aaa30 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Tue, 11 Aug 2026 11:28:32 +0200 Subject: [PATCH 09/22] fix(appkit): resolve repo-wide Biome error blocking CI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI's "Lint & Type Check" job runs `pnpm run check` over the whole repo, so a pre-existing lint error unrelated to this branch failed the build: - remote-tunnel-controller.test.ts had two `afterEach` hooks in one describe (lint/suspicious/noDuplicateTestHooks, error severity). Merge them into one — behavior preserved (env reset + console-spy clear both still run after each test). This file is byte-identical to main; the error predated the branch and only surfaced because CI lints the entire tree. Also drop two dead `biome-ignore lint/suspicious/noExplicitAny` suppressions in the testing kit (fixtures.ts, expect-stream.test.ts): `noExplicitAny` is turned off repo-wide in biome.json, so the comments had no effect (suppressions/unused warnings). The invalid-source test now casts through `unknown as never`. Signed-off-by: Galymzhan --- .../server/remote-tunnel/remote-tunnel-controller.test.ts | 5 +---- packages/appkit/src/testing/fixtures.ts | 3 ++- packages/appkit/src/testing/tests/expect-stream.test.ts | 4 ++-- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/packages/appkit/src/plugins/server/remote-tunnel/remote-tunnel-controller.test.ts b/packages/appkit/src/plugins/server/remote-tunnel/remote-tunnel-controller.test.ts index 01cabb041..afd9fece5 100644 --- a/packages/appkit/src/plugins/server/remote-tunnel/remote-tunnel-controller.test.ts +++ b/packages/appkit/src/plugins/server/remote-tunnel/remote-tunnel-controller.test.ts @@ -38,6 +38,7 @@ describe("RemoteTunnelController", () => { afterEach(() => { process.env = originalEnv; + consoleLogSpy.mockClear(); }); test("middleware hard-blocks in local dev (never initializes manager)", async () => { @@ -168,8 +169,4 @@ describe("RemoteTunnelController", () => { expect(mockManagerInstance.cleanup).toHaveBeenCalledTimes(1); expect(ctrl.isActive()).toBe(false); }); - - afterEach(() => { - consoleLogSpy.mockClear(); - }); }); diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index ddb873f0a..9b98a1930 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -6,7 +6,8 @@ import { ServiceContext } from "../context/service-context"; import type { UserContext } from "../context/user-context"; import type { InstrumentConfig, ITelemetry } from "../telemetry/types"; -// biome-ignore lint/suspicious/noExplicitAny: test fixtures intentionally use loose shapes +// Test fixtures intentionally use loose shapes; `noExplicitAny` is disabled +// repo-wide (see biome.json), so a local alias keeps the intent readable. type Any = any; /** diff --git a/packages/appkit/src/testing/tests/expect-stream.test.ts b/packages/appkit/src/testing/tests/expect-stream.test.ts index 95c001840..e9cb42cb0 100644 --- a/packages/appkit/src/testing/tests/expect-stream.test.ts +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -132,8 +132,8 @@ describe("expectStream — SSE Response", () => { describe("expectStream — invalid source", () => { test("throws for a non-stream value", async () => { await expect( - // biome-ignore lint/suspicious/noExplicitAny: intentionally wrong type - expectStream(42 as any).collect(), + // Intentionally wrong type to exercise the runtime guard. + expectStream(42 as unknown as never).collect(), ).rejects.toThrow(/async iterable, an iterable, or a Response/); }); }); From 8acf7ee8a31157448c5ff1e1d2b159ef0105b2a0 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Tue, 11 Aug 2026 13:43:57 +0200 Subject: [PATCH 10/22] fix(appkit): address cross-model review findings in the testing kit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verified and fixed the findings from an independent code review: - #1 (correctness) expectStream dropped the wire `event:` name when the JSON payload carried its own `type` (spread ran after the assignment). Spread the payload first, then set `type = name ?? parsed.type`, so a frame like `event: error` + `data: {"type":"result"}` reports `error`. Regression test added. - #2 (contract) `@databricks/appkit/testing` eagerly loads vitest via fixtures even for `expectStream`, so vitest is a real requirement. Drop the "optional" peerDependenciesMeta and correct the docs sentence. - #6 (OBO fidelity) the fake `asUser` recorded `asUser: true` unconditionally. Enforce the real `Plugin.asUser` token precondition: a request without `x-forwarded-access-token` throws `missingToken` (missing user id throws too), and the resolved `userId` is recorded on each tool call. Tests now assert both directions (well-formed request vs token-less). - #3 (fidelity) attach() now mirrors AppKit core: registerPlugin plus registerToolProvider for real tool providers, without clobbering injected fakes. getPlugins()/getPluginNames()/hasPlugin() behave as in production. - #12 unknown-tool lookup used `tools[name] === undefined`, so a tool named "constructor"/"toString" hit Object.prototype. Use Object.hasOwn. - #5 drop data-less named SSE frames (real clients ignore them). - #7 re-export the PluginContext type from the testing barrel so MockPluginContext.ctx is nameable through the exports map. - #13 correct the docs: mock.telemetry captures the context's executeTool spans, not plugin-level spans (attachContext rebuilds the plugin's own telemetry). - #4 parseSSEResponse now delegates to the same parseSSEBody as expectStream — one parser, no divergence. All 3 analytics.integration call sites still pass. - #8 reformat template/server/example.test.ts with the template's Prettier so a scaffolded app's `npm run format` passes. - #10 fix the package-doc @example (agentsPlugin._handleStream does not exist). - #11 add kit tests that exercise attach() end-to-end (cache seed, isReady, registration, fake-not-clobbered). Build passes attw + publint; full appkit suite 3125 passed / 1 pre-existing skip. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 9 +- packages/appkit/package.json | 5 - .../agents/tests/dispatch-tool-call.test.ts | 10 +- packages/appkit/src/testing/expect-stream.ts | 48 ++++---- packages/appkit/src/testing/index.ts | 13 ++- .../appkit/src/testing/mock-plugin-context.ts | 79 +++++++++++-- .../src/testing/tests/expect-stream.test.ts | 22 ++++ .../testing/tests/mock-plugin-context.test.ts | 108 +++++++++++++++++- 8 files changed, 245 insertions(+), 49 deletions(-) diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md index 21dbb1948..384843e47 100644 --- a/docs/docs/development/testing.md +++ b/docs/docs/development/testing.md @@ -16,7 +16,7 @@ The kit has two entry points plus a set of fixture helpers: - **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. - **Fixtures** — `createMockRequest`, `createMockResponse`, `mockServiceContext`, and SQL response builders. -The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so it is declared as an optional peer dependency. Any project that runs Vitest already satisfies it; there is nothing extra to install. +The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so `vitest` is a peer dependency and must be installed to import from this subpath. Any project that runs Vitest as its test runner already has it — AppKit apps scaffolded from the template do — so in practice there is nothing extra to add. ## `mockPluginContext()` @@ -79,11 +79,14 @@ expect(mock.routes).toContainEqual( expect.objectContaining({ method: "post", path: "/invocations" }), ); -// The injected telemetry provider, for span assertions. +// The injected telemetry provider records the context's own spans — i.e. the +// span PluginContext.executeTool opens around each cross-plugin tool call. expect(mock.telemetry.getTracer().startActiveSpan).toHaveBeenCalled(); ``` -`RecordedToolCall.asUser` is the high-value signal: it confirms the context routed the call through the user's identity rather than the service principal — a distinction that silent stubs cannot verify. +`mock.telemetry` is injected into the `PluginContext`, so it captures the spans the *context* opens (notably `executeTool`). It is **not** the plugin's own telemetry: `attachContext` rebuilds `this.telemetry` from the real `TelemetryManager`, so spans a plugin opens internally do not land on `mock.telemetry`. + +`RecordedToolCall.asUser` is the high-value signal for cross-plugin calls: because the fake `asUser` enforces the same token precondition as the real `Plugin.asUser`, a dispatch that records `asUser: true` (with `userId` set) genuinely resolved the caller's user scope, and a request missing `x-forwarded-access-token` **rejects** instead — the OBO distinction that silent `{ executeTool }` stubs cannot verify. Assert both directions: a well-formed request records the expected `userId`, and a token-less one throws. ## `expectStream(...)` diff --git a/packages/appkit/package.json b/packages/appkit/package.json index dd19f3b52..1ff6506e7 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -111,11 +111,6 @@ "peerDependencies": { "vitest": ">=1.0.0" }, - "peerDependenciesMeta": { - "vitest": { - "optional": true - } - }, "overrides": { "vite": "npm:rolldown-vite@7.1.14" }, diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index d5915feca..ad4221b69 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -37,10 +37,16 @@ beforeEach(() => { }); function mockReq(): express.Request { + // Carry OBO headers so PluginContext.executeTool's asUser(req) resolves a + // user scope (the mock context enforces the real token precondition). + const headers: Record = { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "alice", + }; return { body: {}, - headers: {}, - header: () => undefined, + headers, + header: (name: string) => headers[name.toLowerCase()], } as unknown as express.Request; } diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts index ca0873d2e..55f3b874c 100644 --- a/packages/appkit/src/testing/expect-stream.ts +++ b/packages/appkit/src/testing/expect-stream.ts @@ -72,7 +72,10 @@ function parseSSEBody(text: string): StreamEvent[] { // `id:` and comment (`:`) lines carry no event type/data we assert on. } - if (name === undefined && dataLines.length === 0) continue; + // A frame with no data line is bookkeeping (a bare `event:`, an `id:`, or a + // `:` comment/heartbeat) that a real SSE client does not surface as an + // event — skip it whether or not it carried an `event:` name. + if (dataLines.length === 0) continue; const data = dataLines.join("\n"); let parsed: Record = {}; @@ -89,9 +92,13 @@ function parseSSEBody(text: string): StreamEvent[] { } } + // The wire `event:` name is authoritative. Spread the payload FIRST, then + // set `type`, so a `data` payload that happens to carry its own `type` + // field (e.g. `event: error` + `data: {"type":"result"}`) cannot override + // the frame's real event name. events.push({ - type: name ?? (parsed.type as string | undefined), ...parsed, + type: name ?? (parsed.type as string | undefined), }); } @@ -189,36 +196,31 @@ export function expectStream(source: StreamSource): StreamAssertion { } /** - * Parse a single-event SSE `Response` into `{ eventType, ...data }`. + * Parse an SSE `Response` and return its **last** event flattened to + * `{ eventType, ...data }`. + * + * A convenience for one-shot assertions on a reply's final event; prefer + * {@link expectStream} for multi-event ordering. It shares {@link parseSSEBody} + * with `expectStream`, so the two never diverge on CRLF handling, comment + * lines, or field parsing. * - * Retained for tests that assert on a one-shot SSE reply; prefer - * {@link expectStream} for multi-event ordering assertions. + * @throws if the response carries no data-bearing event. */ export async function parseSSEResponse(response: Response): Promise<{ eventType: string | null; [key: string]: unknown; }> { const text = await response.text(); - const lines = text.split("\n"); + const events = parseSSEBody(text); + const last = events.at(-1); - let eventType: string | null = null; - let dataLine: string | null = null; - - for (const line of lines) { - if (line.startsWith("event: ")) { - eventType = line.substring(7).trim(); - } else if (line.startsWith("data: ")) { - dataLine = line.substring(6); - } - } - - if (!dataLine) { + if (!last) { throw new Error(`No data found in SSE response: ${text}`); } - const parsed = JSON.parse(dataLine); - return { - eventType, - ...parsed, - }; + // `parseSSEBody` already spread the JSON payload's fields onto the event and + // set `type` from the wire name. Re-key `type` -> `eventType` for this + // helper's historical shape, dropping the internal `type` alias. + const { type, ...rest } = last; + return { eventType: type ?? null, ...rest }; } diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts index 5c34a2d99..61cebf67a 100644 --- a/packages/appkit/src/testing/index.ts +++ b/packages/appkit/src/testing/index.ts @@ -21,9 +21,14 @@ * ```ts * import { mockPluginContext, expectStream } from "@databricks/appkit/testing"; * + * // Attach a real PluginContext (with faked edges) to your plugin instance, + * // then assert on what a streaming source emits. `expectStream` consumes an + * // async event stream, a plain array, or an SSE `Response`. * const mock = mockPluginContext({ analytics: { query: fixtureRows } }); - * await mock.attach(agentsPlugin); - * await expectStream(agentsPlugin._handleStream(req, res)).toEmit( + * const plugin = new MyPlugin({}); + * await mock.attach(plugin); + * + * await expectStream(plugin.streamSomething(input)).toEmit( * "tool_call", * "message_delta", * ); @@ -32,6 +37,10 @@ * @module */ +// Re-export the PluginContext type so `MockPluginContext.ctx` is nameable +// through this entry point — the class is otherwise reachable only via a deep +// path (../core/plugin-context) that is not part of the package's exports map. +export type { PluginContext } from "../core/plugin-context"; export { expectStream, parseSSEResponse, diff --git a/packages/appkit/src/testing/mock-plugin-context.ts b/packages/appkit/src/testing/mock-plugin-context.ts index 2f81d9734..fe6518473 100644 --- a/packages/appkit/src/testing/mock-plugin-context.ts +++ b/packages/appkit/src/testing/mock-plugin-context.ts @@ -7,7 +7,8 @@ import type { } from "shared"; import { CacheManager } from "../cache"; import { InMemoryStorage } from "../cache/storage"; -import { PluginContext } from "../core/plugin-context"; +import { isToolProvider, PluginContext } from "../core/plugin-context"; +import { AuthenticationError } from "../errors"; import type { Plugin } from "../plugin"; import type { ITelemetry } from "../telemetry"; import { createMockTelemetry } from "./fixtures"; @@ -61,11 +62,22 @@ export interface RecordedToolCall { /** The abort signal `executeTool` composed (timeout ∘ caller). */ signal?: AbortSignal; /** - * Whether the call went through the on-behalf-of (`asUser`) path. `true` - * proves `PluginContext.executeTool` resolved the user scope rather than - * running as the service principal. + * Whether the dispatch was resolved through the on-behalf-of (`asUser`) + * path. `PluginContext.executeTool` always calls `provider.asUser(req)`, so + * for a tool reached through `executeTool` this is `true` — and, because the + * fake `asUser` enforces the same token precondition as the real + * {@link Plugin.asUser}, a request with no `x-forwarded-access-token` makes + * that call **throw** rather than record `asUser: true`. The meaningful + * assertions are therefore: a well-formed request records `asUser: true` + * with {@link userId} set, and a token-less request rejects. */ asUser: boolean; + /** + * The user the on-behalf-of scope resolved to (from `x-forwarded-user`), or + * `undefined` for a service-principal call (`asUser: false`). Lets a test + * assert the tool ran as the expected end user, not just that OBO was used. + */ + userId?: string; } /** A single route registered through the context's `addRoute`/`addMiddleware`. */ @@ -199,15 +211,26 @@ export function mockPluginContext( args: unknown, signal: AbortSignal | undefined, asUser: boolean, + userId: string | undefined, ): Promise => { - toolCalls.push({ plugin: name, tool: toolName, args, signal, asUser }); - const response = tools[toolName]; - if (response === undefined) { + toolCalls.push({ + plugin: name, + tool: toolName, + args, + signal, + asUser, + userId, + }); + // `Object.hasOwn`, not `tools[toolName] === undefined`: a tool named + // "constructor"/"toString"/etc. would otherwise resolve to an inherited + // Object.prototype method and be invoked instead of reported missing. + if (!Object.hasOwn(tools, toolName)) { throw new Error( `mockPluginContext: plugin "${name}" has no fake tool "${toolName}". ` + `Available: ${Object.keys(tools).join(", ") || "(none)"}`, ); } + const response = tools[toolName]; return typeof response === "function" ? await (response as (a: unknown, s?: AbortSignal) => unknown)( args, @@ -219,18 +242,36 @@ export function mockPluginContext( const base: ToolProvider = { getAgentTools: () => record.tools, executeAgentTool: (toolName, args, signal) => - resolve(toolName, args, signal, false), + resolve(toolName, args, signal, false, undefined), }; - // `asUser(req)` returns a user-scoped view whose executeAgentTool records - // that the OBO path ran — this is how executeTool's user scoping becomes - // observable without a real user token. + // Mirror the real `Plugin.asUser` token precondition (plugin.ts) so the + // recorded `asUser` flag reflects genuine user-scope resolution rather than + // being unconditionally true: a request with no `x-forwarded-access-token` + // throws `missingToken` (production behavior), except in development where + // the real code skips impersonation. This is edge-faking of asUser's + // *contract*, not a reimplementation of `runInUserContext`/`ServiceContext`. const asUser = (req: IAppRequest): ToolProvider => { record.asUserRequests.push(req as express.Request); + const token = (req as express.Request) + .header?.("x-forwarded-access-token") + ?.trim(); + const userId = (req as express.Request) + .header?.("x-forwarded-user") + ?.trim(); + const isDev = process.env.NODE_ENV === "development"; + + if (!token && !isDev) { + throw AuthenticationError.missingToken("user token"); + } + if (token && !userId && !isDev) { + throw AuthenticationError.missingUserId(); + } + return { getAgentTools: () => record.tools, executeAgentTool: (toolName, args, signal) => - resolve(toolName, args, signal, true), + resolve(toolName, args, signal, true, userId), }; }; @@ -266,6 +307,20 @@ export function mockPluginContext( await CacheManager.getInstance({ storage: new InMemoryStorage({}) }); } plugin.attachContext({ context: ctx }); + + // Mirror what AppKit core does after attachContext (core/appkit.ts): put + // the plugin in the registry so `getPlugins()`/`getPluginNames()`/ + // `hasPlugin()` and any sibling-plugin lookup behave as in production. Only + // register it as a tool provider when it actually is one AND its name does + // not collide with an injected fake — the fakes are the authored test + // doubles and must not be overwritten by the plugin under test. + ctx.registerPlugin(plugin.name, plugin as unknown as BasePlugin); + if (isToolProvider(plugin) && !providers.has(plugin.name)) { + ctx.registerToolProvider( + plugin.name, + plugin as unknown as Parameters[1], + ); + } return plugin; } diff --git a/packages/appkit/src/testing/tests/expect-stream.test.ts b/packages/appkit/src/testing/tests/expect-stream.test.ts index e9cb42cb0..eb87afeed 100644 --- a/packages/appkit/src/testing/tests/expect-stream.test.ts +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -116,6 +116,28 @@ describe("expectStream — SSE Response", () => { ]); }); + test("the wire event: name wins over a type field inside the data payload", async () => { + // Regression: object spread must not let a `data` payload carrying its own + // `type` override the frame's real event name. Here the wire says `error` + // but the payload says `result`; the emitted event must be `error`. + const body = `event: error\ndata: {"type":"result","message":"boom"}\n\n`; + const res = new Response(body); + const events = await expectStream(res).collect(); + expect(events[0]?.type).toBe("error"); + // A stream that actually errored must NOT satisfy an assertion for result. + await expect( + expectStream(new Response(body)).toEmitExactly("result"), + ).rejects.toThrow(/exactly/); + }); + + test("drops a data-less named frame (real clients ignore it)", async () => { + const body = `event: ping\n\nevent: result\ndata: {"ok":true}\n\n`; + const res = new Response(body); + await expect(expectStream(res).toEmitExactly("result")).resolves.toEqual([ + "result", + ]); + }); + test("parses CRLF-delimited frames from a spec-compliant SSE stream", async () => { // A real server may use \r\n\r\n between frames; AppKit's own writer uses // \n\n. Both must parse to distinct events, not one collapsed block. diff --git a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts b/packages/appkit/src/testing/tests/mock-plugin-context.test.ts index 3c17b1236..1da3511ff 100644 --- a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts +++ b/packages/appkit/src/testing/tests/mock-plugin-context.test.ts @@ -1,8 +1,29 @@ import type express from "express"; import { describe, expect, test } from "vitest"; import { PluginContext } from "../../core/plugin-context"; +import { Plugin } from "../../plugin"; +import type { PluginManifest } from "../../registry"; import { mockPluginContext } from "../mock-plugin-context"; +// A minimal real plugin for exercising attach() end-to-end. +class ProbePlugin extends Plugin { + static manifest = { + name: "probe", + displayName: "Probe", + description: "attach() probe", + resources: { required: [], optional: [] }, + } as PluginManifest<"probe">; + + ready() { + // `isReady` is protected; expose it for the attach() assertion. + return (this as unknown as { isReady: boolean }).isReady; + } + + register() { + this.context?.addRoute("get", "/probe", (_req, res) => res.end()); + } +} + /** * Contract for `mockPluginContext`. The point of the kit is that it wraps the * REAL PluginContext — so these tests drive the real `executeTool`, @@ -10,7 +31,14 @@ import { mockPluginContext } from "../mock-plugin-context"; * timeout, route recording) rather than a reimplementation. */ -function mockReq(headers: Record = {}): express.Request { +// Default to a well-formed OBO request (user token + user id) so executeTool's +// asUser path resolves. Pass `{}` explicitly to model a token-less request. +function mockReq( + headers: Record = { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "alice", + }, +): express.Request { return { body: {}, headers, @@ -48,18 +76,61 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () ); expect(result).toEqual(rows); - // executeTool always resolves the user scope via provider.asUser(req). + // executeTool resolves the user scope via provider.asUser(req), and the + // fake resolves the user id from the request headers. expect(mock.toolCalls).toHaveLength(1); expect(mock.toolCalls[0]).toMatchObject({ plugin: "analytics", tool: "top_users", args: { limit: 10 }, asUser: true, + userId: "alice", }); // The OBO request object is the one we passed in. expect(mock.providers.get("analytics")?.asUserRequests).toHaveLength(1); }); + test("rejects a token-less request the way the real asUser does", async () => { + // The fake asUser enforces the same token precondition as Plugin.asUser, + // so a header-less request must reject rather than silently record + // asUser: true — this is what makes the OBO assertion meaningful. + const mock = mockPluginContext({ analytics: { top_users: [] } }); + + await expect( + mock.ctx.executeTool(mockReq({}), "analytics", "top_users", {}), + ).rejects.toThrow(/Missing user token/); + // The dispatch never reached the tool. + expect(mock.toolCalls).toHaveLength(0); + }); + + test("rejects a request with a token but no user id", async () => { + const mock = mockPluginContext({ analytics: { top_users: [] } }); + + await expect( + mock.ctx.executeTool( + mockReq({ "x-forwarded-access-token": "tok" }), + "analytics", + "top_users", + {}, + ), + ).rejects.toThrow(/Missing user id|user id/i); + expect(mock.toolCalls).toHaveLength(0); + }); + + test("records the resolved user id so a test can assert who the tool ran as", async () => { + const mock = mockPluginContext({ analytics: { top_users: [] } }); + await mock.ctx.executeTool( + mockReq({ + "x-forwarded-access-token": "tok", + "x-forwarded-user": "bob", + }), + "analytics", + "top_users", + {}, + ); + expect(mock.toolCalls[0]).toMatchObject({ asUser: true, userId: "bob" }); + }); + test("invokes a function response with the args and the composed signal", async () => { const mock = mockPluginContext({ analytics: { @@ -175,3 +246,36 @@ describe("mockPluginContext — registerProvider after construction", () => { expect(result).toBe("pong"); }); }); + +describe("mockPluginContext — attach()", () => { + test("seeds the cache, flips isReady, and registers the plugin", async () => { + const mock = mockPluginContext(); + const plugin = new ProbePlugin({}); + + // Before attach the plugin may not be ready (no cache seeded yet in a + // fresh process); after attach it is, and it is in the context registry. + const returned = await mock.attach(plugin); + + expect(returned).toBe(plugin); + expect(plugin.ready()).toBe(true); + expect(mock.ctx.getPluginNames()).toContain("probe"); + expect(mock.ctx.hasPlugin("probe")).toBe(true); + + // A route the plugin registers post-attach is captured through the context. + plugin.register(); + expect(mock.routes).toContainEqual( + expect.objectContaining({ method: "get", path: "/probe" }), + ); + }); + + test("does not overwrite an injected fake provider of the same name", async () => { + // If the plugin under test shares a name with an injected fake, the fake + // (the authored double) must win — attach must not clobber it. + const mock = mockPluginContext({ probe: { canned: "fake" } }); + const plugin = new ProbePlugin({}); + await mock.attach(plugin); + + const result = await mock.ctx.executeTool(mockReq(), "probe", "canned", {}); + expect(result).toBe("fake"); + }); +}); From 9297c73045a852a941971fe8139762457119186d Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Tue, 11 Aug 2026 13:45:48 +0200 Subject: [PATCH 11/22] chore(appkit): drop knip vitest-ignore now that vitest is a real peer dep With vitest declared as a (non-optional) peerDependency, knip recognizes it as used, so the earlier ignoreDependencies entry is unnecessary. This reverts knip.json to its original state. Signed-off-by: Galymzhan --- knip.json | 3 --- 1 file changed, 3 deletions(-) diff --git a/knip.json b/knip.json index 1ca5a1b7e..0e96b7df5 100644 --- a/knip.json +++ b/knip.json @@ -9,9 +9,6 @@ "workspaces": { "packages/appkit-ui": { "ignoreDependencies": ["tailwindcss", "tw-animate-css"] - }, - "packages/appkit": { - "ignoreDependencies": ["vitest"] } }, "ignore": [ From 3f11ea3c6f4d3bdc1bb52c1a47d2b7d2572e76d8 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Tue, 11 Aug 2026 14:18:18 +0200 Subject: [PATCH 12/22] fix(appkit): make vitest a normal dependency, not a package-wide peer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A required peerDependency has no per-subpath scope: it applied to the whole @databricks/appkit package, so every production consumer that never imports the testing kit got an unsatisfied peer (npm 7+ auto-installs vitest into their tree; pnpm warns) — a wider blast radius than the eager-import bug it was meant to fix. Follow appkit's own precedent instead: `vite` backs the ./type-generator subpath as a normal `dependency`, installed for everyone but loaded only by importers of that subpath. Do the same for `vitest` and ./testing. vitest is referenced solely by dist/testing/fixtures.js, never by the main/plugin/core entry, so a consumer importing createApp never loads it. Verified end-to-end: scaffolded an app whose own vitest (4.1.9) differs in major from appkit's dependency (3.2.4), forcing a nested second copy. The testing kit's vi.fn()/vi.spyOn() mocks and expect(...).toHaveBeenCalled() assertions work across the two instances (vi spies carry their own call state), and npm install emits no peer-dep warning. Build passes attw + publint. Also fold in the template example's Prettier formatting (template uses Prettier, not Biome) so a scaffolded app's `npm run format` passes. Signed-off-by: Galymzhan --- packages/appkit/package.json | 4 +--- pnpm-lock.yaml | 4 ++-- template/server/example.test.ts | 14 +++----------- 3 files changed, 6 insertions(+), 16 deletions(-) diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 1ff6506e7..7be28e242 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -96,6 +96,7 @@ "semver": "7.7.3", "shared": "workspace:*", "vite": "npm:rolldown-vite@7.1.14", + "vitest": "3.2.4", "ws": "8.21.0", "zod": "4.3.6" }, @@ -108,9 +109,6 @@ "@types/ws": "8.18.1", "@vitejs/plugin-react": "5.1.1" }, - "peerDependencies": { - "vitest": ">=1.0.0" - }, "overrides": { "vite": "npm:rolldown-vite@7.1.14" }, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index ed122fb5f..ecbcc2a50 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -345,7 +345,7 @@ importers: specifier: npm:rolldown-vite@7.1.14 version: rolldown-vite@7.1.14(@types/node@25.2.3)(esbuild@0.25.10)(jiti@2.6.1)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) vitest: - specifier: '>=1.0.0' + specifier: 3.2.4 version: 3.2.4(@types/debug@4.1.12)(@types/node@25.2.3)(jiti@2.6.1)(jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6))(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) ws: specifier: 8.21.0 @@ -7382,7 +7382,7 @@ packages: git-raw-commits@4.0.0: resolution: {integrity: sha512-ICsMM1Wk8xSGMowkOmPrzo2Fgmfo4bMHLNX6ytHjajRJUqvHOw/TFapQ+QG75c3X/tTDDhOSRPGC52dDbNM8FQ==} engines: {node: '>=16'} - deprecated: This package is no longer maintained. For the JavaScript API, please use @conventional-changelog/git-client instead. + deprecated: Deprecated and no longer maintained. Use @conventional-changelog/git-client instead. hasBin: true git-up@8.1.1: diff --git a/template/server/example.test.ts b/template/server/example.test.ts index 149569645..fb9ad9253 100644 --- a/template/server/example.test.ts +++ b/template/server/example.test.ts @@ -1,8 +1,5 @@ import { Plugin, type PluginManifest } from '@databricks/appkit'; -import { - expectStream, - mockPluginContext, -} from '@databricks/appkit/testing'; +import { expectStream, mockPluginContext } from '@databricks/appkit/testing'; import { describe, expect, test } from 'vitest'; /** @@ -56,17 +53,12 @@ describe('testing kit example', () => { await mock.attach(plugin); await plugin.setup(); - expect(mock.routes).toContainEqual( - expect.objectContaining({ method: 'get', path: '/hello' }), - ); + expect(mock.routes).toContainEqual(expect.objectContaining({ method: 'get', path: '/hello' })); }); test('asserts the ordered events a stream emits', async () => { const plugin = new GreeterPlugin({}); - await expectStream(plugin.greet('world')).toEmit( - 'greeting_start', - 'greeting_end', - ); + await expectStream(plugin.greet('world')).toEmit('greeting_start', 'greeting_end'); }); }); From c2c1aa96acc652e175434b9ee33e08bd2c1d67be Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Wed, 12 Aug 2026 12:08:42 +0200 Subject: [PATCH 13/22] refactor(appkit): rename mockPluginContext to createTestPluginContext MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The helper builds the REAL PluginContext with faked edges — it does not mock the context — so the name was misleading. Rename to createTestPluginContext (and the MockPluginContext type to TestPluginContext), matching the create*-for-tests convention, and rename the files to test-plugin-context.ts. Pre-merge and unreleased, so no external consumers are affected. Also finish the #13 doc-accuracy fix in the shipped JSDoc (not just the docs page): the telemetry field comment now states it captures the context's spans (executeTool), not plugin-internal spans — attachContext rebuilds the plugin's this.telemetry from the real TelemetryManager. These comments ship in dist/testing/*.d.ts, so IntelliSense previously showed the unqualified claim. Build passes attw + publint; full appkit suite 3125 passed / 1 pre-existing skip. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 16 +++--- .../agents/tests/dispatch-tool-call.test.ts | 6 +-- .../agents/tests/route-handler-errors.test.ts | 4 +- packages/appkit/src/testing/expect-stream.ts | 2 +- packages/appkit/src/testing/fixtures.ts | 2 +- packages/appkit/src/testing/index.ts | 14 +++--- ...ugin-context.ts => test-plugin-context.ts} | 25 ++++++---- ...xt.test.ts => test-plugin-context.test.ts} | 50 +++++++++---------- template/server/example.test.ts | 8 +-- 9 files changed, 67 insertions(+), 60 deletions(-) rename packages/appkit/src/testing/{mock-plugin-context.ts => test-plugin-context.ts} (92%) rename packages/appkit/src/testing/tests/{mock-plugin-context.test.ts => test-plugin-context.test.ts} (84%) diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md index 384843e47..baa51865a 100644 --- a/docs/docs/development/testing.md +++ b/docs/docs/development/testing.md @@ -12,15 +12,15 @@ Exercise a plugin's real code paths — route registration, cross-plugin tool di The kit has two entry points plus a set of fixture helpers: -- **`mockPluginContext()`** — build a real `PluginContext` with faked edges and attach it to a plugin. +- **`createTestPluginContext()`** — build a real `PluginContext` with faked edges and attach it to a plugin. - **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. - **Fixtures** — `createMockRequest`, `createMockResponse`, `mockServiceContext`, and SQL response builders. The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so `vitest` is a peer dependency and must be installed to import from this subpath. Any project that runs Vitest as its test runner already has it — AppKit apps scaffolded from the template do — so in practice there is nothing extra to add. -## `mockPluginContext()` +## `createTestPluginContext()` -`PluginContext` is the mediator AppKit passes to every plugin — it buffers routes, tracks tool providers, and runs cross-plugin tool calls with user scoping and a timeout. `mockPluginContext()` returns the **real** context with three edges faked: +`PluginContext` is the mediator AppKit passes to every plugin — it buffers routes, tracks tool providers, and runs cross-plugin tool calls with user scoping and a timeout. `createTestPluginContext()` returns the **real** context with three edges faked: | Edge | How it's faked | | --- | --- | @@ -35,9 +35,9 @@ Because the context is real, `executeTool` still resolves the user scope via `as Pass canned responses keyed by plugin name, then tool name. A response is either a static value or a function of the call arguments and the composed abort signal: ```ts -import { mockPluginContext } from "@databricks/appkit/testing"; +import { createTestPluginContext } from "@databricks/appkit/testing"; -const mock = mockPluginContext({ +const mock = createTestPluginContext({ analytics: { // static response top_users: [{ user: "alice", events: 42 }], @@ -122,7 +122,7 @@ Instantiate the plugin **class** directly with `new`. The `analytics()` / `agent ```ts import { Plugin, type PluginManifest } from "@databricks/appkit"; -import { expectStream, mockPluginContext } from "@databricks/appkit/testing"; +import { expectStream, createTestPluginContext } from "@databricks/appkit/testing"; import { describe, expect, test } from "vitest"; // A small plugin that registers a route and streams two events. @@ -146,7 +146,7 @@ class GreeterPlugin extends Plugin { describe("greeter plugin", () => { test("registers its route through the context", async () => { - const mock = mockPluginContext(); + const mock = createTestPluginContext(); const plugin = new GreeterPlugin({}); await mock.attach(plugin); @@ -170,7 +170,7 @@ describe("greeter plugin", () => { To test a plugin that dispatches cross-plugin tool calls, register fake providers and assert on `mock.toolCalls` — including `asUser`, which confirms the on-behalf-of path ran: ```ts -const mock = mockPluginContext({ analytics: { query: [{ n: 1 }] } }); +const mock = createTestPluginContext({ analytics: { query: [{ n: 1 }] } }); const plugin = new MyAgentPlugin({ dir: false }); await mock.attach(plugin); diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index ad4221b69..d335a933c 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -1,7 +1,7 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; -import { mockPluginContext } from "../../../testing"; +import { createTestPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; /** @@ -323,7 +323,7 @@ describe("dispatchToolCall — toolkit timeout plumbing", () => { // forwarded timeout is exercised through actual signal composition — and // spying on it lets us keep asserting the exact call signature the agents // plugin passes. - const mock = mockPluginContext({ analytics: { query: "rows" } }); + const mock = createTestPluginContext({ analytics: { query: "rows" } }); const executeToolSpy = vi.spyOn(mock.ctx, "executeTool"); // biome-ignore lint/suspicious/noExplicitAny: attach the real context to the plugin (plugin as any).context = mock.ctx; @@ -362,7 +362,7 @@ describe("dispatchToolCall — toolkit timeout plumbing", () => { const { runState } = makeRunState(plugin); runState.limits.toolCallTimeoutMs = 5; - const mock = mockPluginContext({ + const mock = createTestPluginContext({ analytics: { query: (_args, signal) => new Promise((_resolve, reject) => { diff --git a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 547e512f5..3a31bfdfb 100644 --- a/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts +++ b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts @@ -1,7 +1,7 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; -import { mockPluginContext } from "../../../testing"; +import { createTestPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; /** @@ -401,7 +401,7 @@ describe("/invocations and /responses are aliases", () => { // captures the RAW handlers passed to addRoute — the aliasing assertion // needs the original references, which the context's forwardAsyncErrors // wrapping would otherwise break. - const mock = mockPluginContext(); + const mock = createTestPluginContext(); // biome-ignore lint/suspicious/noExplicitAny: attach the real context (plugin as any).context = mock.ctx; // biome-ignore lint/suspicious/noExplicitAny: invoke private mounter diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts index 55f3b874c..1d71e3c55 100644 --- a/packages/appkit/src/testing/expect-stream.ts +++ b/packages/appkit/src/testing/expect-stream.ts @@ -143,7 +143,7 @@ function isSubsequence(actual: string[], expected: string[]): boolean { /** * Consume a stream and make ordered assertions about the event types it emits. * - * Deterministic and network-free: pair it with {@link mockPluginContext} to + * Deterministic and network-free: pair it with {@link createTestPluginContext} to * exercise a plugin's streaming handler and assert what it emits. * * @example Async event stream (adapter output) diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 9b98a1930..0536484bd 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -13,7 +13,7 @@ type Any = any; /** * Creates a mock telemetry provider for testing. Every span/meter/logger is a * `vi.fn()` no-op, so plugins that trace, count, or log run without a live - * OpenTelemetry pipeline. Passed into {@link mockPluginContext} as the one + * OpenTelemetry pipeline. Passed into {@link createTestPluginContext} as the one * injectable production seam. */ export function createMockTelemetry(): ITelemetry { diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts index 61cebf67a..fccb64a2a 100644 --- a/packages/appkit/src/testing/index.ts +++ b/packages/appkit/src/testing/index.ts @@ -10,7 +10,7 @@ * with no credentials. * * Two entry points: - * - {@link mockPluginContext} — build a real `PluginContext` with faked edges + * - {@link createTestPluginContext} — build a real `PluginContext` with faked edges * and attach it to a plugin. * - {@link expectStream} — assert the ordered event types a stream emits. * @@ -19,12 +19,12 @@ * * @example * ```ts - * import { mockPluginContext, expectStream } from "@databricks/appkit/testing"; + * import { createTestPluginContext, expectStream } from "@databricks/appkit/testing"; * * // Attach a real PluginContext (with faked edges) to your plugin instance, * // then assert on what a streaming source emits. `expectStream` consumes an * // async event stream, a plain array, or an SSE `Response`. - * const mock = mockPluginContext({ analytics: { query: fixtureRows } }); + * const mock = createTestPluginContext({ analytics: { query: fixtureRows } }); * const plugin = new MyPlugin({}); * await mock.attach(plugin); * @@ -37,7 +37,7 @@ * @module */ -// Re-export the PluginContext type so `MockPluginContext.ctx` is nameable +// Re-export the PluginContext type so `TestPluginContext.ctx` is nameable // through this entry point — the class is otherwise reachable only via a deep // path (../core/plugin-context) that is not part of the package's exports map. export type { PluginContext } from "../core/plugin-context"; @@ -65,11 +65,11 @@ export { type TestContextOptions, } from "./fixtures"; export { + createTestPluginContext, type FakeProvider, type FakeProviders, type FakeToolResponse, - type MockPluginContext, - mockPluginContext, type RecordedRoute, type RecordedToolCall, -} from "./mock-plugin-context"; + type TestPluginContext, +} from "./test-plugin-context"; diff --git a/packages/appkit/src/testing/mock-plugin-context.ts b/packages/appkit/src/testing/test-plugin-context.ts similarity index 92% rename from packages/appkit/src/testing/mock-plugin-context.ts rename to packages/appkit/src/testing/test-plugin-context.ts index fe6518473..f14a5543d 100644 --- a/packages/appkit/src/testing/mock-plugin-context.ts +++ b/packages/appkit/src/testing/test-plugin-context.ts @@ -43,7 +43,7 @@ export type FakeToolResponse = * Fake connector responses, keyed by plugin name and then tool name: * * ```ts - * mockPluginContext({ analytics: { query: fixtureRows } }); + * createTestPluginContext({ analytics: { query: fixtureRows } }); * ``` * * Each top-level key registers a fake {@link ToolProvider} under that plugin @@ -102,13 +102,18 @@ export interface FakeProvider { } /** - * The result of {@link mockPluginContext}: the real `PluginContext` plus the + * The result of {@link createTestPluginContext}: the real `PluginContext` plus the * seams a test needs to drive and inspect it. */ -export interface MockPluginContext { +export interface TestPluginContext { /** The real {@link PluginContext}, constructed with mock telemetry. */ ctx: PluginContext; - /** The injected mock telemetry provider — assert on spans here. */ + /** + * The mock telemetry provider injected into the {@link PluginContext}. + * Captures the spans the *context* opens (notably `executeTool`) — not the + * plugin's own spans: `attachContext` rebuilds the plugin's `this.telemetry` + * from the real `TelemetryManager`, so plugin-internal spans do not land here. + */ telemetry: ITelemetry; /** * Tool dispatches observed across all fake providers, in call order. Live — @@ -145,7 +150,9 @@ export interface MockPluginContext { * timeout composition, and the on-behalf-of (`asUser`) path all run for real. * Only three edges are faked, matching the seams the class actually has: * - * - **Telemetry** is a mock provider (the one injectable production seam). + * - **Telemetry** is a mock provider injected into the context (the one + * injectable production seam); it records the context's own spans, not the + * plugin's. * - **Tool providers** are fakes registered through the existing public * `registerToolProvider`; their `asUser`/`executeAgentTool` are recorded. * - **Routes** are captured by wrapping the public `addRoute`/`addMiddleware`. @@ -156,15 +163,15 @@ export interface MockPluginContext { * * @example * ```ts - * const mock = mockPluginContext({ analytics: { query: fixtureRows } }); + * const mock = createTestPluginContext({ analytics: { query: fixtureRows } }); * await mock.attach(agentsPlugin); * // ...exercise a handler that dispatches analytics.query... * expect(mock.toolCalls[0]).toMatchObject({ plugin: "analytics", asUser: true }); * ``` */ -export function mockPluginContext( +export function createTestPluginContext( fakes: FakeProviders = {}, -): MockPluginContext { +): TestPluginContext { const telemetry = createMockTelemetry(); const ctx = new PluginContext({ telemetry }); @@ -226,7 +233,7 @@ export function mockPluginContext( // Object.prototype method and be invoked instead of reported missing. if (!Object.hasOwn(tools, toolName)) { throw new Error( - `mockPluginContext: plugin "${name}" has no fake tool "${toolName}". ` + + `createTestPluginContext: plugin "${name}" has no fake tool "${toolName}". ` + `Available: ${Object.keys(tools).join(", ") || "(none)"}`, ); } diff --git a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts b/packages/appkit/src/testing/tests/test-plugin-context.test.ts similarity index 84% rename from packages/appkit/src/testing/tests/mock-plugin-context.test.ts rename to packages/appkit/src/testing/tests/test-plugin-context.test.ts index 1da3511ff..9c845454b 100644 --- a/packages/appkit/src/testing/tests/mock-plugin-context.test.ts +++ b/packages/appkit/src/testing/tests/test-plugin-context.test.ts @@ -3,7 +3,7 @@ import { describe, expect, test } from "vitest"; import { PluginContext } from "../../core/plugin-context"; import { Plugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; -import { mockPluginContext } from "../mock-plugin-context"; +import { createTestPluginContext } from "../test-plugin-context"; // A minimal real plugin for exercising attach() end-to-end. class ProbePlugin extends Plugin { @@ -25,7 +25,7 @@ class ProbePlugin extends Plugin { } /** - * Contract for `mockPluginContext`. The point of the kit is that it wraps the + * Contract for `createTestPluginContext`. The point of the kit is that it wraps the * REAL PluginContext — so these tests drive the real `executeTool`, * `addRoute`, and `getToolProviders` and assert the observable seams (OBO, * timeout, route recording) rather than a reimplementation. @@ -46,14 +46,14 @@ function mockReq( } as unknown as express.Request; } -describe("mockPluginContext — construction", () => { +describe("createTestPluginContext — construction", () => { test("produces a real PluginContext instance", () => { - const { ctx } = mockPluginContext(); + const { ctx } = createTestPluginContext(); expect(ctx).toBeInstanceOf(PluginContext); }); test("registers fake providers passed at construction", () => { - const { ctx } = mockPluginContext({ + const { ctx } = createTestPluginContext({ analytics: { query: [{ id: 1 }] }, genie: { ask: "hi" }, }); @@ -63,10 +63,10 @@ describe("mockPluginContext — construction", () => { }); }); -describe("mockPluginContext — executeTool runs the REAL user-scoping path", () => { +describe("createTestPluginContext — executeTool runs the REAL user-scoping path", () => { test("dispatches through asUser and returns the canned static response", async () => { const rows = [{ user: "alice", n: 3 }]; - const mock = mockPluginContext({ analytics: { top_users: rows } }); + const mock = createTestPluginContext({ analytics: { top_users: rows } }); const result = await mock.ctx.executeTool( mockReq(), @@ -94,7 +94,7 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () // The fake asUser enforces the same token precondition as Plugin.asUser, // so a header-less request must reject rather than silently record // asUser: true — this is what makes the OBO assertion meaningful. - const mock = mockPluginContext({ analytics: { top_users: [] } }); + const mock = createTestPluginContext({ analytics: { top_users: [] } }); await expect( mock.ctx.executeTool(mockReq({}), "analytics", "top_users", {}), @@ -104,7 +104,7 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () }); test("rejects a request with a token but no user id", async () => { - const mock = mockPluginContext({ analytics: { top_users: [] } }); + const mock = createTestPluginContext({ analytics: { top_users: [] } }); await expect( mock.ctx.executeTool( @@ -118,7 +118,7 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () }); test("records the resolved user id so a test can assert who the tool ran as", async () => { - const mock = mockPluginContext({ analytics: { top_users: [] } }); + const mock = createTestPluginContext({ analytics: { top_users: [] } }); await mock.ctx.executeTool( mockReq({ "x-forwarded-access-token": "tok", @@ -132,7 +132,7 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () }); test("invokes a function response with the args and the composed signal", async () => { - const mock = mockPluginContext({ + const mock = createTestPluginContext({ analytics: { query: (args, signal) => ({ echoed: args, aborted: signal?.aborted }), }, @@ -148,7 +148,7 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () }); test("forwards the caller timeout so a slow tool is aborted", async () => { - const mock = mockPluginContext({ + const mock = createTestPluginContext({ slow: { wait: (_args, signal) => new Promise((_resolve, reject) => { @@ -168,14 +168,14 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () }); test("throws with a helpful message for an unknown plugin", async () => { - const mock = mockPluginContext({ analytics: { query: [] } }); + const mock = createTestPluginContext({ analytics: { query: [] } }); await expect( mock.ctx.executeTool(mockReq(), "nope", "query", {}), ).rejects.toThrow(/unknown plugin "nope"/); }); test("throws with a helpful message for an unknown tool", async () => { - const mock = mockPluginContext({ analytics: { query: [] } }); + const mock = createTestPluginContext({ analytics: { query: [] } }); await expect( mock.ctx.executeTool(mockReq(), "analytics", "missing", {}), ).rejects.toThrow(/no fake tool "missing"/); @@ -184,7 +184,7 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () test("returns a null response as a value rather than treating it as missing", async () => { // `resolve` distinguishes a null fake response (valid) from undefined // (unregistered tool), so a tool can model an empty/absent result. - const mock = mockPluginContext({ analytics: { lookup: null } }); + const mock = createTestPluginContext({ analytics: { lookup: null } }); const result = await mock.ctx.executeTool( mockReq(), "analytics", @@ -195,9 +195,9 @@ describe("mockPluginContext — executeTool runs the REAL user-scoping path", () }); }); -describe("mockPluginContext — telemetry seam", () => { +describe("createTestPluginContext — telemetry seam", () => { test("records a span on the injected mock telemetry for each executeTool", async () => { - const mock = mockPluginContext({ analytics: { query: [] } }); + const mock = createTestPluginContext({ analytics: { query: [] } }); const tracer = mock.telemetry.getTracer(); await mock.ctx.executeTool(mockReq(), "analytics", "query", {}); @@ -207,9 +207,9 @@ describe("mockPluginContext — telemetry seam", () => { }); }); -describe("mockPluginContext — route recording", () => { +describe("createTestPluginContext — route recording", () => { test("records addRoute calls with raw (pre-wrap) handlers", () => { - const mock = mockPluginContext(); + const mock = createTestPluginContext(); const handler: express.RequestHandler = (_req, res) => { res.end(); }; @@ -229,7 +229,7 @@ describe("mockPluginContext — route recording", () => { }); test("records addMiddleware under the 'use' method", () => { - const mock = mockPluginContext(); + const mock = createTestPluginContext(); const mw: express.RequestHandler = (_req, _res, next) => next(); mock.ctx.addMiddleware("/api", mw); expect(mock.routes).toEqual([ @@ -238,18 +238,18 @@ describe("mockPluginContext — route recording", () => { }); }); -describe("mockPluginContext — registerProvider after construction", () => { +describe("createTestPluginContext — registerProvider after construction", () => { test("adds a provider dynamically", async () => { - const mock = mockPluginContext(); + const mock = createTestPluginContext(); mock.registerProvider("late", { ping: "pong" }); const result = await mock.ctx.executeTool(mockReq(), "late", "ping", {}); expect(result).toBe("pong"); }); }); -describe("mockPluginContext — attach()", () => { +describe("createTestPluginContext — attach()", () => { test("seeds the cache, flips isReady, and registers the plugin", async () => { - const mock = mockPluginContext(); + const mock = createTestPluginContext(); const plugin = new ProbePlugin({}); // Before attach the plugin may not be ready (no cache seeded yet in a @@ -271,7 +271,7 @@ describe("mockPluginContext — attach()", () => { test("does not overwrite an injected fake provider of the same name", async () => { // If the plugin under test shares a name with an injected fake, the fake // (the authored double) must win — attach must not clobber it. - const mock = mockPluginContext({ probe: { canned: "fake" } }); + const mock = createTestPluginContext({ probe: { canned: "fake" } }); const plugin = new ProbePlugin({}); await mock.attach(plugin); diff --git a/template/server/example.test.ts b/template/server/example.test.ts index fb9ad9253..9140c2e24 100644 --- a/template/server/example.test.ts +++ b/template/server/example.test.ts @@ -1,5 +1,5 @@ import { Plugin, type PluginManifest } from '@databricks/appkit'; -import { expectStream, mockPluginContext } from '@databricks/appkit/testing'; +import { expectStream, createTestPluginContext } from '@databricks/appkit/testing'; import { describe, expect, test } from 'vitest'; /** @@ -10,7 +10,7 @@ import { describe, expect, test } from 'vitest'; * it as a starting point for testing your own plugins. * * Two headline helpers are shown below: - * - `mockPluginContext()` — a real PluginContext with faked edges, attachable + * - `createTestPluginContext()` — a real PluginContext with faked edges, attachable * to a plugin so its real code paths (routes, tool dispatch, user scoping) * run under test. * - `expectStream(...).toEmit(...)` — assert the ordered event types a @@ -32,7 +32,7 @@ class GreeterPlugin extends Plugin { } as PluginManifest<'greeter'>; async setup() { - // Routes registered here are captured by mockPluginContext().routes. + // Routes registered here are captured by createTestPluginContext().routes. this.context?.addRoute('get', '/hello', (_req, res) => { res.end(); }); @@ -47,7 +47,7 @@ class GreeterPlugin extends Plugin { describe('testing kit example', () => { test('attaches a real PluginContext and records registered routes', async () => { - const mock = mockPluginContext(); + const mock = createTestPluginContext(); const plugin = new GreeterPlugin({}); await mock.attach(plugin); From 25649711bec33d49d077bbdcf8fe9c6e3909fb26 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Wed, 12 Aug 2026 13:32:26 +0200 Subject: [PATCH 14/22] refactor(appkit): dedupe testing fixtures and tidy test-plugin-context Behavior-preserving cleanups in the testing kit: - createMockRequest reuses createMockWorkspaceClient() instead of an inline copy of the same mock client (verified identical). - createMockServiceContext / createMockUserContext / mockServiceContext inline the createMockWorkspaceClient() call into the `||` fallback, so the mock client is built only when the caller did not supply one. - The fake asUser view spreads `...base` and overrides executeAgentTool rather than re-declaring getAgentTools. - expectStream's isSubsequence breaks once the expected sequence is fully matched. No semantic change; typecheck clean and all kit + migrated tests pass. Signed-off-by: Galymzhan --- packages/appkit/src/testing/expect-stream.ts | 1 + packages/appkit/src/testing/fixtures.ts | 32 ++++--------------- .../appkit/src/testing/test-plugin-context.ts | 2 +- 3 files changed, 9 insertions(+), 26 deletions(-) diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts index 1d71e3c55..1b182939c 100644 --- a/packages/appkit/src/testing/expect-stream.ts +++ b/packages/appkit/src/testing/expect-stream.ts @@ -136,6 +136,7 @@ function isSubsequence(actual: string[], expected: string[]): boolean { let i = 0; for (const type of actual) { if (i < expected.length && type === expected[i]) i++; + if (i === expected.length) break; } return i === expected.length; } diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 0536484bd..240fdbc6e 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -108,21 +108,7 @@ export function createMockRouter(): { * service-principal client slots; override any field via `overrides`. */ export function createMockRequest(overrides: Any = {}) { - const mockWorkspaceClient = { - statementExecution: { - executeStatement: vi.fn().mockResolvedValue({ - status: { state: "SUCCEEDED" }, - result: { data: [] }, - }), - }, - // Analytics route now calls `warehouses.get` before issuing SQL to - // ensure the warehouse is RUNNING. Default to RUNNING so existing - // tests that only care about SQL behaviour aren't affected. - warehouses: { - get: vi.fn().mockResolvedValue({ state: "RUNNING" }), - start: vi.fn().mockResolvedValue(undefined), - }, - }; + const mockWorkspaceClient = createMockWorkspaceClient(); const req = { params: {}, @@ -163,7 +149,6 @@ export function createMockResponse() { sendStatus: vi.fn().mockReturnThis(), end: vi.fn(function (this: Any) { this.writableEnded = true; - // Trigger 'close' event when end is called if (eventListeners.close) { for (const handler of eventListeners.close) { handler(); @@ -264,10 +249,9 @@ export function createMockWorkspaceClient() { * singleton. Use with {@link mockServiceContext} to install it. */ export function createMockServiceContext(options: TestContextOptions = {}) { - const mockWorkspaceClient = createMockWorkspaceClient(); - const serviceContext: ServiceContextState = { - client: (options.serviceDatabricksClient || mockWorkspaceClient) as Any, + client: (options.serviceDatabricksClient || + createMockWorkspaceClient()) as Any, serviceUserId: options.serviceUserId || "test-service-user", warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), @@ -282,10 +266,9 @@ export function createMockServiceContext(options: TestContextOptions = {}) { export function createMockUserContext( options: TestContextOptions = {}, ): UserContext { - const mockWorkspaceClient = createMockWorkspaceClient(); - return { - client: (options.userDatabricksClient || mockWorkspaceClient) as Any, + client: (options.userDatabricksClient || + createMockWorkspaceClient()) as Any, userId: options.userId || "test-user", warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), @@ -316,13 +299,12 @@ export function mockServiceContext(options: TestContextOptions = {}) { .spyOn(ServiceContext, "isInitialized") .mockReturnValue(true); - // Mock createUserContext to return a test user context const createUserContextSpy = vi .spyOn(ServiceContext, "createUserContext") .mockImplementation((_token: string, userId: string, userName?: string) => { - const mockWorkspaceClient = createMockWorkspaceClient(); return { - client: (options.userDatabricksClient || mockWorkspaceClient) as Any, + client: (options.userDatabricksClient || + createMockWorkspaceClient()) as Any, userId, userName, warehouseId: serviceContext.warehouseId, diff --git a/packages/appkit/src/testing/test-plugin-context.ts b/packages/appkit/src/testing/test-plugin-context.ts index f14a5543d..26770b804 100644 --- a/packages/appkit/src/testing/test-plugin-context.ts +++ b/packages/appkit/src/testing/test-plugin-context.ts @@ -276,7 +276,7 @@ export function createTestPluginContext( } return { - getAgentTools: () => record.tools, + ...base, executeAgentTool: (toolName, args, signal) => resolve(toolName, args, signal, true, userId), }; From 1e28a96a720b0792af512aac87c474df6b6fa0e6 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Wed, 12 Aug 2026 15:20:33 +0200 Subject: [PATCH 15/22] fix(appkit): resolve third-review findings in the testing kit - #1 (P1) The docs called vitest a peer dependency, but the manifest ships it under `dependencies` (the decision we landed on, matching how appkit ships `vite` for ./type-generator). Correct the docs to match: appkit installs vitest for you, and it loads only when you import ./testing. Manifest and docs now agree. - #2 (P2) expectStream buffered the source eagerly with no bound, so a non-terminating stream hung until the runner's own timeout. Add an optional `{ timeout }` that fails fast with a clear, kit-specific error; document it and cover both directions with tests. - #3 (P2) The fake asUser replicates asUser's token precondition but not the real dev-mode `DEV_OBO_FALLBACK_KEY` OTel marker (a module-private telemetry detail). Narrow the docs and JSDoc to say so and point users at the recorded asUser/userId fields instead of isDevOboFallback(). Build passes attw + publint; full appkit suite 3141 passed / 1 pre-existing skip. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 10 ++- packages/appkit/src/testing/expect-stream.ts | 62 ++++++++++++++++++- packages/appkit/src/testing/index.ts | 1 + .../appkit/src/testing/test-plugin-context.ts | 8 +++ .../src/testing/tests/expect-stream.test.ts | 22 +++++++ 5 files changed, 99 insertions(+), 4 deletions(-) diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md index baa51865a..a4c9bdeae 100644 --- a/docs/docs/development/testing.md +++ b/docs/docs/development/testing.md @@ -16,7 +16,7 @@ The kit has two entry points plus a set of fixture helpers: - **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. - **Fixtures** — `createMockRequest`, `createMockResponse`, `mockServiceContext`, and SQL response builders. -The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so `vitest` is a peer dependency and must be installed to import from this subpath. Any project that runs Vitest as its test runner already has it — AppKit apps scaffolded from the template do — so in practice there is nothing extra to add. +The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so AppKit lists `vitest` as a dependency and installs it for you — there is nothing extra to add. It loads only when you import `@databricks/appkit/testing`; apps that never import the testing subpath never pull it into their runtime. (This mirrors how AppKit ships `vite` for the `@databricks/appkit/type-generator` subpath.) ## `createTestPluginContext()` @@ -88,6 +88,8 @@ expect(mock.telemetry.getTracer().startActiveSpan).toHaveBeenCalled(); `RecordedToolCall.asUser` is the high-value signal for cross-plugin calls: because the fake `asUser` enforces the same token precondition as the real `Plugin.asUser`, a dispatch that records `asUser: true` (with `userId` set) genuinely resolved the caller's user scope, and a request missing `x-forwarded-access-token` **rejects** instead — the OBO distinction that silent `{ executeTool }` stubs cannot verify. Assert both directions: a well-formed request records the expected `userId`, and a token-less one throws. +The fake replicates `asUser`'s **token precondition**, not its internal dev-mode telemetry marker: in `NODE_ENV=development` the real `Plugin.asUser` skips impersonation and sets an OTel `isDevOboFallback()` flag, which the fake does not reproduce. Assert OBO through the recorded `asUser`/`userId` fields rather than `isDevOboFallback()`. + ## `expectStream(...)` AppKit plugins stream Server-Sent Events. `expectStream` consumes a stream and asserts the ordered event types it emits. It accepts an async iterable (an agent adapter's `run()`), a plain array of events, or an SSE `Response` (or a promise of one) whose body it parses. @@ -107,6 +109,12 @@ const types = await expectStream(res).collectTypes(); `toEmit` checks that the expected types appear **in order** but tolerates other events before, between, or after them — which is what you want for streams that interleave bookkeeping events like heartbeats or metadata. Use `toEmitExactly` when the stream's shape is fully determined. +`expectStream` buffers the whole source before asserting, so a stream that never terminates would otherwise hang until the test runner's own timeout. Pass `{ timeout }` to fail fast with a clear error instead: + +```ts +await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); +``` + ## Fixtures The kit re-exports the request/response/context fixtures AppKit uses internally: diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts index 1b182939c..e196e64fb 100644 --- a/packages/appkit/src/testing/expect-stream.ts +++ b/packages/appkit/src/testing/expect-stream.ts @@ -105,7 +105,9 @@ function parseSSEBody(text: string): StreamEvent[] { return events; } -async function collectEvents(source: StreamSource): Promise { +async function collectEventsInner( + source: StreamSource, +): Promise { const resolved = await source; if (resolved instanceof Response) { @@ -131,6 +133,48 @@ async function collectEvents(source: StreamSource): Promise { ); } +async function collectEvents( + source: StreamSource, + timeoutMs?: number, +): Promise { + // `expectStream` buffers the whole source before asserting. Without a bound, + // a stream that never terminates hangs until Vitest's per-test timeout — + // a poor signal. When a timeout is given, surface a clear, kit-specific + // error instead. The pending collection is abandoned (it cannot be force + // -cancelled), so callers should pair this with an aborting source. + if (timeoutMs === undefined) return collectEventsInner(source); + + let timer: ReturnType | undefined; + const timeout = new Promise((_resolve, reject) => { + timer = setTimeout( + () => + reject( + new Error( + `expectStream: stream did not terminate within ${timeoutMs}ms. ` + + "Ensure the source ends, or raise the { timeout } option.", + ), + ), + timeoutMs, + ); + }); + + try { + return await Promise.race([collectEventsInner(source), timeout]); + } finally { + if (timer) clearTimeout(timer); + } +} + +/** Options for {@link expectStream}. */ +export interface ExpectStreamOptions { + /** + * Fail with a clear error if the source has not finished within this many + * milliseconds, instead of hanging until the test runner's own timeout. + * Omit to buffer the source with no bound (the default). + */ + timeout?: number; +} + /** Does `expected` appear as an in-order subsequence of `actual`? */ function isSubsequence(actual: string[], expected: string[]): boolean { let i = 0; @@ -157,9 +201,21 @@ function isSubsequence(actual: string[], expected: string[]): boolean { * const res = await fetch("/api/analytics/query/top_users", { method: "POST" }); * await expectStream(res).toEmit("warehouse_status", "result"); * ``` + * + * @example Guard against a non-terminating stream + * ```ts + * await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); + * ``` + * + * @param source - The stream, iterable, or SSE `Response` to consume. + * @param options - See {@link ExpectStreamOptions}; pass `{ timeout }` to fail + * fast on a stream that never ends. */ -export function expectStream(source: StreamSource): StreamAssertion { - const events = collectEvents(source); +export function expectStream( + source: StreamSource, + options: ExpectStreamOptions = {}, +): StreamAssertion { + const events = collectEvents(source, options.timeout); return { async collect() { diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts index fccb64a2a..83f2f371c 100644 --- a/packages/appkit/src/testing/index.ts +++ b/packages/appkit/src/testing/index.ts @@ -42,6 +42,7 @@ // path (../core/plugin-context) that is not part of the package's exports map. export type { PluginContext } from "../core/plugin-context"; export { + type ExpectStreamOptions, expectStream, parseSSEResponse, type StreamAssertion, diff --git a/packages/appkit/src/testing/test-plugin-context.ts b/packages/appkit/src/testing/test-plugin-context.ts index 26770b804..a70054a42 100644 --- a/packages/appkit/src/testing/test-plugin-context.ts +++ b/packages/appkit/src/testing/test-plugin-context.ts @@ -70,6 +70,9 @@ export interface RecordedToolCall { * that call **throw** rather than record `asUser: true`. The meaningful * assertions are therefore: a well-formed request records `asUser: true` * with {@link userId} set, and a token-less request rejects. + * + * The fake replicates the token precondition only, not the real dev-mode + * OTel `isDevOboFallback()` marker — assert OBO here, not via that flag. */ asUser: boolean; /** @@ -258,6 +261,11 @@ export function createTestPluginContext( // throws `missingToken` (production behavior), except in development where // the real code skips impersonation. This is edge-faking of asUser's // *contract*, not a reimplementation of `runInUserContext`/`ServiceContext`. + // + // Deliberately NOT reproduced: the real dev-mode path sets an OTel + // `DEV_OBO_FALLBACK_KEY` marker (read by `isDevOboFallback()`). That key is + // module-private telemetry plumbing; assert OBO via the recorded + // `asUser`/`userId` fields, not `isDevOboFallback()`. const asUser = (req: IAppRequest): ToolProvider => { record.asUserRequests.push(req as express.Request); const token = (req as express.Request) diff --git a/packages/appkit/src/testing/tests/expect-stream.test.ts b/packages/appkit/src/testing/tests/expect-stream.test.ts index eb87afeed..c91cd176f 100644 --- a/packages/appkit/src/testing/tests/expect-stream.test.ts +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -160,6 +160,28 @@ describe("expectStream — invalid source", () => { }); }); +describe("expectStream — timeout", () => { + test("fails with a clear error when a stream never terminates", async () => { + // A generator that yields once then hangs forever. + async function* neverEnds(): AsyncGenerator<{ type: string }> { + yield { type: "start" }; + await new Promise(() => {}); // never resolves + } + + await expect( + expectStream(neverEnds(), { timeout: 20 }).toEmit("start"), + ).rejects.toThrow(/did not terminate within 20ms/); + }); + + test("a terminating stream resolves normally under a generous timeout", async () => { + await expect( + expectStream(asyncEvents([{ type: "a" }, { type: "b" }]), { + timeout: 1000, + }).toEmit("a", "b"), + ).resolves.toEqual(["a", "b"]); + }); +}); + describe("parseSSEResponse — single-event helper", () => { test("returns eventType plus parsed data fields", async () => { const res = new Response( From 01dd25c4473bab1d322b42d926d144206f4f77e4 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Wed, 12 Aug 2026 17:16:38 +0200 Subject: [PATCH 16/22] test(appkit): dogfood the testing kit on analytics and genie plugins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Exercise @databricks/appkit/testing against real core plugins to validate it beyond the two agent proof sites and produce usage references: - analytics.kit.test.ts: cross-plugin executeTool via createTestPluginContext — OBO identity (asUser/userId), token-precondition rejection, and per-call timeout abort. Needs only the kit (no workspace/ServiceContext). - genie.kit.test.ts: drives the real _handleSendMessage SSE stream and asserts event order with expectStream(...).toEmit(...). Both add genuinely new coverage (streamed SSE order + OBO dispatch identity were untested). Full appkit suite 3145 passed / 1 pre-existing skip. Developer-experience notes (kit wins + friction, e.g. createMockResponse doesn't compose with expectStream) captured in internal/ for the milestone review. Signed-off-by: Galymzhan --- .../analytics/tests/analytics.kit.test.ts | 102 +++++++++++ .../src/plugins/genie/tests/genie.kit.test.ts | 170 ++++++++++++++++++ 2 files changed, 272 insertions(+) create mode 100644 packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts create mode 100644 packages/appkit/src/plugins/genie/tests/genie.kit.test.ts diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts new file mode 100644 index 000000000..7b04e51be --- /dev/null +++ b/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts @@ -0,0 +1,102 @@ +import type express from "express"; +import { describe, expect, test } from "vitest"; +import { createTestPluginContext } from "../../../testing"; + +/** + * Dogfooding `@databricks/appkit/testing` for the cross-plugin tool-call (OBO) + * scenario — the case `createTestPluginContext` is purpose-built for. + * + * A plugin that consumes another plugin's tools calls + * `this.context.executeTool(pluginName, toolName, ...)`. Here we register a + * fake `analytics` provider and drive that dispatch, asserting both the result + * and that it resolved through the user's identity (the on-behalf-of path that + * silent `{ executeTool }` stubs could never verify). + * + * This needs ONLY the kit — no workspace, no ServiceContext, no network — which + * is the sweet spot noted in internal/testing-kit-dogfooding.md. (Driving + * analytics' own SQL handlers, by contrast, still needs the ServiceContext / + * workspace-client fixtures because that work lives behind those seams.) + */ + +function mockReq(headers: Record): express.Request { + return { + body: {}, + headers, + header: (name: string) => headers[name.toLowerCase()], + } as unknown as express.Request; +} + +describe("analytics as a cross-plugin tool provider — dogfooding the kit", () => { + test("a consumer dispatches analytics.query on-behalf-of the user", async () => { + const rows = [{ customer: "Acme", revenue: 1_000_000 }]; + const mock = createTestPluginContext({ + analytics: { + // A canned result for the analytics `query` tool. + query: (args) => ({ rows, echoedArgs: args }), + }, + }); + + // Simulate what a consumer plugin (e.g. agents) does internally: resolve a + // sibling plugin's tool through the shared PluginContext. + const req = mockReq({ + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "analyst@example.com", + }); + const result = await mock.ctx.executeTool(req, "analytics", "query", { + sql: "SELECT * FROM top_customers", + }); + + expect(result).toEqual({ + rows, + echoedArgs: { sql: "SELECT * FROM top_customers" }, + }); + + // The kit proves the dispatch ran as the end user, not the service + // principal — and records who. + expect(mock.toolCalls).toHaveLength(1); + expect(mock.toolCalls[0]).toMatchObject({ + plugin: "analytics", + tool: "query", + asUser: true, + userId: "analyst@example.com", + }); + }); + + test("a token-less request is rejected before the tool runs", async () => { + const mock = createTestPluginContext({ + analytics: { query: () => ({ rows: [] }) }, + }); + + await expect( + mock.ctx.executeTool(mockReq({}), "analytics", "query", {}), + ).rejects.toThrow(/Missing user token/); + expect(mock.toolCalls).toHaveLength(0); + }); + + test("the per-call timeout the caller forwards actually aborts a slow tool", async () => { + const mock = createTestPluginContext({ + analytics: { + query: (_args, signal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject(new Error("aborted by timeout")), + ); + }), + }, + }); + + await expect( + mock.ctx.executeTool( + mockReq({ + "x-forwarded-access-token": "t", + "x-forwarded-user": "u", + }), + "analytics", + "query", + {}, + undefined, + 5, // 5ms timeout + ), + ).rejects.toThrow(/aborted by timeout/); + }); +}); diff --git a/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts b/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts new file mode 100644 index 000000000..ca6aee5a0 --- /dev/null +++ b/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts @@ -0,0 +1,170 @@ +import type express from "express"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { ServiceContext } from "../../../context"; +import { createTestPluginContext, expectStream } from "../../../testing"; +import { type GeniePlugin, genie } from "../genie"; + +/** + * Dogfooding `@databricks/appkit/testing` on a real core plugin (genie). + * + * Genie's `_handleSendMessage` streams SSE via the base `executeStream`. This + * suite drives that real handler and asserts the emitted event ORDER with + * `expectStream` — the streaming-assertion path the kit is meant to make easy. + * + * See internal/testing-kit-dogfooding.md for the developer-experience notes + * this exercise produced (notably: the kit's `createMockResponse` does not + * capture written SSE bytes, so a small capturing response is needed to bridge + * a `res.write`-based handler into `expectStream`). + */ + +// The base Plugin reads the cache singleton on attach; a tiny in-memory stub +// keeps this unit-level (mirrors the pattern in the sibling genie.test.ts). +const { mockCacheInstance } = vi.hoisted(() => ({ + mockCacheInstance: { + get: vi.fn(), + set: vi.fn(), + delete: vi.fn(), + getOrExecute: vi.fn( + async (_k: unknown[], fn: (s?: AbortSignal) => Promise) => fn(), + ), + generateKey: vi.fn((...a: unknown[]) => JSON.stringify(a)), + }, +})); + +vi.mock("../../../cache", () => ({ + CacheManager: { getInstanceSync: vi.fn(() => mockCacheInstance) }, +})); + +/** + * Collects the SSE bytes a handler writes and exposes them as a `Response`, + * so a `res.write`-based handler can be asserted with `expectStream`. + * + * NOTE: this bridge is exactly the friction the dogfooding writeup flags — the + * kit's own `createMockResponse` throws written chunks away, so streaming + * handler tests need this until the kit ships a capturing response. + */ +function createCapturingResponse() { + const chunks: string[] = []; + const listeners: Record void>> = {}; + const res = { + headersSent: false, + writableEnded: false, + statusCode: 200, + status: vi.fn().mockReturnThis(), + json: vi.fn().mockReturnThis(), + setHeader: vi.fn().mockReturnThis(), + flushHeaders: vi.fn().mockReturnThis(), + write: vi.fn((chunk: unknown) => { + chunks.push(String(chunk)); + return true; + }), + end: vi.fn(function (this: { writableEnded: boolean }) { + this.writableEnded = true; + for (const fn of listeners.close ?? []) fn(); + return this; + }), + on: vi.fn((event: string, fn: () => void) => { + listeners[event] ??= []; + listeners[event].push(fn); + return res; + }), + off: vi.fn().mockReturnThis(), + destroy: vi.fn().mockReturnThis(), + }; + return { + res: res as unknown as express.Response, + toResponse: () => new Response(chunks.join("")), + }; +} + +function mockReq(body: unknown): express.Request { + const headers: Record = { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "alice", + }; + return { + params: { alias: "myspace" }, + query: {}, + body, + headers, + header: (name: string) => headers[name.toLowerCase()], + on: vi.fn(), + off: vi.fn(), + } as unknown as express.Request; +} + +describe("genie plugin — dogfooding the testing kit", () => { + let plugin: GeniePlugin; + let serviceContextMock: ReturnType; + + // A minimal ServiceContext stand-in (the kit's mockServiceContext fixture + // covers this, but genie's streaming path only needs a resolvable context). + function mockServiceContextLite() { + const state = { + client: {} as never, + serviceUserId: "sp", + warehouseId: Promise.resolve("wh"), + workspaceId: Promise.resolve("ws"), + }; + const get = vi.spyOn(ServiceContext, "get").mockReturnValue(state); + const isInit = vi + .spyOn(ServiceContext, "isInitialized") + .mockReturnValue(true); + return { + restore: () => { + get.mockRestore(); + isInit.mockRestore(); + }, + }; + } + + beforeEach(async () => { + process.env.DATABRICKS_HOST = "https://test.databricks.com"; + ServiceContext.reset(); + serviceContextMock = mockServiceContextLite(); + + // `genie(...)` returns a { plugin: ctor, config } descriptor for createApp; + // for a unit test, instantiate the class and attach a real PluginContext + // through the kit (seeds cache + flips isReady, the production path). + const GenieCtor = genie({}).plugin as unknown as new ( + c: unknown, + ) => GeniePlugin; + plugin = new GenieCtor({ spaces: { myspace: "space-1" }, timeout: 5000 }); + await createTestPluginContext().attach(plugin); + + // Fake the one real edge — the network connector — to yield a known event + // sequence. Everything else (executeStream, SSE writing) runs for real. + ( + plugin as unknown as { + genieConnector: { + streamSendMessage: (...a: unknown[]) => AsyncGenerator; + }; + } + ).genieConnector.streamSendMessage = async function* () { + yield { type: "status", status: "ASKING_AI" }; + yield { type: "message", content: "Here are your results" }; + yield { type: "complete" }; + }; + }); + + afterEach(() => { + serviceContextMock.restore(); + vi.restoreAllMocks(); + }); + + test("_handleSendMessage streams status -> message -> complete in order", async () => { + const { res, toResponse } = createCapturingResponse(); + + await ( + plugin as unknown as { + _handleSendMessage: ( + r: express.Request, + w: express.Response, + ) => Promise; + } + )._handleSendMessage(mockReq({ content: "top customers?" }), res); + + // The kit's expectStream parses the real SSE the handler wrote. + await expectStream(toResponse()).toEmit("status", "message", "complete"); + }); +}); From f09ca5de335dcfe036ee777e16807310fcd90647 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Thu, 13 Aug 2026 13:21:47 +0200 Subject: [PATCH 17/22] refactor(appkit): address testing-kit review feedback Resolve the eight review comments on the testing kit: - createMockResponse now captures written SSE bytes and exposes sseResponse(); expectStream reads a captured mock response directly, so streaming-route tests no longer need a hand-rolled bridge. - Ship vitest as an optional peer dependency (+ devDependency) instead of a plain runtime dependency, keeping the test framework out of production installs and deduping to the app's own copy. Ignore it in knip. - Add an obo option to createMockRequest so on-behalf-of tests set the forwarded identity headers with one flag. - Add resetTestCache() to clear the shared cache singleton between tests. - Use the documented attach() instead of an any-cast in the agents dispatch tests. - Drop the unused createMockServiceContext/createMockUserContext builders from the public surface; keep the service-context builder internal. - Pin the previously untested edges: the Object.hasOwn tool-lookup guard, the dev-mode asUser branch, and parseSSEBody's non-object data values. - Add useServiceContextMock() to register the mock lifecycle in one line, returning a live accessor. Dogfood the new helpers in the analytics, genie, and serving suites, and document them in the testing guide. Signed-off-by: Galymzhan --- docs/docs/development/testing.md | 50 +++- knip.json | 3 + packages/appkit/package.json | 12 +- .../agents/tests/dispatch-tool-call.test.ts | 6 +- .../analytics/tests/analytics.kit.test.ts | 30 +-- .../src/plugins/genie/tests/genie.kit.test.ts | 68 ++---- .../src/plugins/serving/tests/serving.test.ts | 12 +- packages/appkit/src/testing/expect-stream.ts | 52 ++++- packages/appkit/src/testing/fixtures.ts | 214 +++++++++++++++--- packages/appkit/src/testing/index.ts | 7 +- .../src/testing/tests/expect-stream.test.ts | 86 +++++++ .../appkit/src/testing/tests/fixtures.test.ts | 125 ++++++++++ .../testing/tests/test-plugin-context.test.ts | 47 ++++ pnpm-lock.yaml | 6 +- tools/test-helpers.ts | 3 +- 15 files changed, 591 insertions(+), 130 deletions(-) create mode 100644 packages/appkit/src/testing/tests/fixtures.test.ts diff --git a/docs/docs/development/testing.md b/docs/docs/development/testing.md index a4c9bdeae..874abb263 100644 --- a/docs/docs/development/testing.md +++ b/docs/docs/development/testing.md @@ -16,7 +16,7 @@ The kit has two entry points plus a set of fixture helpers: - **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. - **Fixtures** — `createMockRequest`, `createMockResponse`, `mockServiceContext`, and SQL response builders. -The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so AppKit lists `vitest` as a dependency and installs it for you — there is nothing extra to add. It loads only when you import `@databricks/appkit/testing`; apps that never import the testing subpath never pull it into their runtime. (This mirrors how AppKit ships `vite` for the `@databricks/appkit/type-generator` subpath.) +The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so `vitest` is an **optional peer dependency**: you already have it (you're writing Vitest tests), and the kit resolves to your copy rather than bundling a second one. Because it's optional, it is not installed into apps that never import `@databricks/appkit/testing` — production installs stay free of the test framework. Any Vitest v3 or v4 works. ## `createTestPluginContext()` @@ -58,7 +58,17 @@ await mock.attach(plugin); Instantiate the plugin **class** directly (`new MyAgentPlugin(...)`). The `analytics()` / `agents()` factories you pass to `createApp` return a descriptor for the app to construct — for a unit test you want the instance. -The cache `attach()` seeds is a process-wide singleton: `CacheManager` is initialized once per test process and reused. Vitest isolates test *files* in separate workers, so caches never leak across files, but tests **within one file** share it. If a test populates the cache and a later test in the same file must not see it, reset between tests (e.g. clear the cache in `beforeEach`). +The cache `attach()` seeds is a process-wide singleton: `CacheManager` is initialized once per test process and reused. Vitest isolates test *files* in separate workers, so caches never leak across files, but tests **within one file** share it. If a test populates the cache and a later test in the same file must not see it, clear it between tests with `resetTestCache()`: + +```ts +import { resetTestCache } from "@databricks/appkit/testing"; + +beforeEach(async () => { + await resetTestCache(); // no-op if the cache isn't initialized yet +}); +``` + +It also helps *within* a single test — clear the cache to force a miss, then assert the following call is a hit. ### Inspecting what happened @@ -92,7 +102,7 @@ The fake replicates `asUser`'s **token precondition**, not its internal dev-mode ## `expectStream(...)` -AppKit plugins stream Server-Sent Events. `expectStream` consumes a stream and asserts the ordered event types it emits. It accepts an async iterable (an agent adapter's `run()`), a plain array of events, or an SSE `Response` (or a promise of one) whose body it parses. +AppKit plugins stream Server-Sent Events. `expectStream` consumes a stream and asserts the ordered event types it emits. It accepts an async iterable (an agent adapter's `run()`), a plain array of events, an SSE `Response` (or a promise of one) whose body it parses, or a `createMockResponse()` whose captured writes it replays. ```ts import { expectStream } from "@databricks/appkit/testing"; @@ -107,6 +117,22 @@ await expectStream(events).toEmitExactly("warehouse_status", "result"); const types = await expectStream(res).collectTypes(); ``` +### Asserting a plugin's streaming route + +Most plugins stream SSE from a **route handler** (`res.write(...)`), not a bare generator. `createMockResponse()` captures those writes, and `expectStream` reads them straight back — drive the real handler, then assert: + +```ts +import { createMockRequest, createMockResponse, expectStream } from "@databricks/appkit/testing"; + +const res = createMockResponse(); +await plugin._handleStream(createMockRequest({ obo: true }), res); + +// The mock captured the SSE the handler wrote; expectStream parses it. +await expectStream(res).toEmit("status", "result"); +``` + +`expectStream(res)` and `expectStream(res.sseResponse())` are equivalent — the latter hands you the raw `Response` if you want it. Do **not** pass the SSE body as a string: a string is an iterable of characters, so `expectStream` rejects it with a pointer to `sseResponse()` rather than emitting one "event" per character. + `toEmit` checks that the expected types appear **in order** but tolerates other events before, between, or after them — which is what you want for streams that interleave bookkeeping events like heartbeats or metadata. Use `toEmitExactly` when the stream's shape is fully determined. `expectStream` buffers the whole source before asserting, so a stream that never terminates would otherwise hang until the test runner's own timeout. Pass `{ timeout }` to fail fast with a clear error instead: @@ -119,10 +145,21 @@ await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); The kit re-exports the request/response/context fixtures AppKit uses internally: -- `createMockRequest(overrides?)` / `createMockResponse()` — Express request/response doubles, including the streaming flags (`headersSent`, `writableEnded`) and a mock `WorkspaceClient`. +- `createMockRequest(overrides?)` / `createMockResponse()` — Express request/response doubles, including the streaming flags (`headersSent`, `writableEnded`) and a mock `WorkspaceClient`. Pass `obo: true` (or `obo: { userId, token, email }`) to set the forwarded identity headers `asUser` requires, instead of hand-adding them. `createMockResponse()` also captures everything a handler writes; pass it to `expectStream` (or call `sseResponse()`) to assert a streaming route's SSE. - `mockServiceContext(options?)` — spy the `ServiceContext` singleton so code that resolves the service principal or a user context gets test doubles. Call in `beforeEach`, and call the returned `restore()` in `afterEach`. +- `useServiceContextMock(options?)` — the same, in one line: it registers the `beforeEach` install and `afterEach` restore for you. Call it at the top of a `describe` block (not inside a test), and read the live `.current` handle from within a test: + ```ts + describe("my plugin", () => { + const ctx = useServiceContextMock(); + test("...", async () => { + await handler(createMockRequest({ obo: true }), res); + expect(ctx.current.createUserContextSpy).toHaveBeenCalled(); + }); + }); + ``` - `createSuccessfulSQLResponse(rows, columns)` / `createFailedSQLResponse(message)` — build SQL Warehouse statement responses. - `setupDatabricksEnv(overrides?)` — set `DATABRICKS_HOST` / `DATABRICKS_WAREHOUSE_ID` to test values. +- `resetTestCache()` — clear the shared cache singleton between (or within) tests; no-ops if the cache isn't initialized yet. ## Full example @@ -130,7 +167,7 @@ Instantiate the plugin **class** directly with `new`. The `analytics()` / `agent ```ts import { Plugin, type PluginManifest } from "@databricks/appkit"; -import { expectStream, createTestPluginContext } from "@databricks/appkit/testing"; +import { expectStream, createMockRequest, createTestPluginContext } from "@databricks/appkit/testing"; import { describe, expect, test } from "vitest"; // A small plugin that registers a route and streams two events. @@ -182,6 +219,9 @@ const mock = createTestPluginContext({ analytics: { query: [{ n: 1 }] } }); const plugin = new MyAgentPlugin({ dir: false }); await mock.attach(plugin); +// `obo` sets the forwarded identity headers `asUser` needs — without them the +// dispatch would (correctly) reject with "Missing user token". +const req = createMockRequest({ obo: true }); await plugin.runSomethingThatCallsAnalytics(req); expect(mock.toolCalls[0]).toMatchObject({ diff --git a/knip.json b/knip.json index 0e96b7df5..1f3d29fa1 100644 --- a/knip.json +++ b/knip.json @@ -7,6 +7,9 @@ "docs" ], "workspaces": { + "packages/appkit": { + "ignoreDependencies": ["vitest"] + }, "packages/appkit-ui": { "ignoreDependencies": ["tailwindcss", "tw-animate-css"] } diff --git a/packages/appkit/package.json b/packages/appkit/package.json index b98db7f10..914db2477 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -96,10 +96,17 @@ "semver": "7.7.3", "shared": "workspace:*", "vite": "npm:rolldown-vite@7.1.14", - "vitest": "3.2.4", "ws": "8.21.0", "zod": "4.3.6" }, + "peerDependencies": { + "vitest": ">=3" + }, + "peerDependenciesMeta": { + "vitest": { + "optional": true + } + }, "devDependencies": { "@opentelemetry/context-async-hooks": "2.8.0", "@types/express": "4.17.25", @@ -107,7 +114,8 @@ "@types/json-schema": "7.0.15", "@types/pg": "8.16.0", "@types/ws": "8.18.1", - "@vitejs/plugin-react": "5.1.1" + "@vitejs/plugin-react": "5.1.1", + "vitest": "3.2.4" }, "overrides": { "vite": "npm:rolldown-vite@7.1.14" diff --git a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts index d335a933c..811bb68e1 100644 --- a/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts +++ b/packages/appkit/src/plugins/agents/tests/dispatch-tool-call.test.ts @@ -325,8 +325,7 @@ describe("dispatchToolCall — toolkit timeout plumbing", () => { // plugin passes. const mock = createTestPluginContext({ analytics: { query: "rows" } }); const executeToolSpy = vi.spyOn(mock.ctx, "executeTool"); - // biome-ignore lint/suspicious/noExplicitAny: attach the real context to the plugin - (plugin as any).context = mock.ctx; + await mock.attach(plugin); const result = await callDispatch(plugin, { runState, @@ -372,8 +371,7 @@ describe("dispatchToolCall — toolkit timeout plumbing", () => { }), }, }); - // biome-ignore lint/suspicious/noExplicitAny: attach the real context - (plugin as any).context = mock.ctx; + await mock.attach(plugin); await expect( callDispatch(plugin, { diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts index 7b04e51be..e5cd06bee 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts @@ -1,6 +1,6 @@ import type express from "express"; import { describe, expect, test } from "vitest"; -import { createTestPluginContext } from "../../../testing"; +import { createMockRequest, createTestPluginContext } from "../../../testing"; /** * Dogfooding `@databricks/appkit/testing` for the cross-plugin tool-call (OBO) @@ -18,14 +18,6 @@ import { createTestPluginContext } from "../../../testing"; * workspace-client fixtures because that work lives behind those seams.) */ -function mockReq(headers: Record): express.Request { - return { - body: {}, - headers, - header: (name: string) => headers[name.toLowerCase()], - } as unknown as express.Request; -} - describe("analytics as a cross-plugin tool provider — dogfooding the kit", () => { test("a consumer dispatches analytics.query on-behalf-of the user", async () => { const rows = [{ customer: "Acme", revenue: 1_000_000 }]; @@ -38,10 +30,9 @@ describe("analytics as a cross-plugin tool provider — dogfooding the kit", () // Simulate what a consumer plugin (e.g. agents) does internally: resolve a // sibling plugin's tool through the shared PluginContext. - const req = mockReq({ - "x-forwarded-access-token": "user-token", - "x-forwarded-user": "analyst@example.com", - }); + const req = createMockRequest({ + obo: { userId: "analyst@example.com" }, + }) as unknown as express.Request; const result = await mock.ctx.executeTool(req, "analytics", "query", { sql: "SELECT * FROM top_customers", }); @@ -51,8 +42,7 @@ describe("analytics as a cross-plugin tool provider — dogfooding the kit", () echoedArgs: { sql: "SELECT * FROM top_customers" }, }); - // The kit proves the dispatch ran as the end user, not the service - // principal — and records who. + // Prove the dispatch ran as the end user, not the service principal. expect(mock.toolCalls).toHaveLength(1); expect(mock.toolCalls[0]).toMatchObject({ plugin: "analytics", @@ -67,8 +57,10 @@ describe("analytics as a cross-plugin tool provider — dogfooding the kit", () analytics: { query: () => ({ rows: [] }) }, }); + // No `obo` — a request with no forwarded token must be rejected. + const req = createMockRequest() as unknown as express.Request; await expect( - mock.ctx.executeTool(mockReq({}), "analytics", "query", {}), + mock.ctx.executeTool(req, "analytics", "query", {}), ).rejects.toThrow(/Missing user token/); expect(mock.toolCalls).toHaveLength(0); }); @@ -85,12 +77,10 @@ describe("analytics as a cross-plugin tool provider — dogfooding the kit", () }, }); + const req = createMockRequest({ obo: true }) as unknown as express.Request; await expect( mock.ctx.executeTool( - mockReq({ - "x-forwarded-access-token": "t", - "x-forwarded-user": "u", - }), + req, "analytics", "query", {}, diff --git a/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts b/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts index ca6aee5a0..a7ce5ef6a 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts @@ -1,7 +1,11 @@ import type express from "express"; -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, test, vi } from "vitest"; import { ServiceContext } from "../../../context"; -import { createTestPluginContext, expectStream } from "../../../testing"; +import { + createMockResponse, + createTestPluginContext, + expectStream, +} from "../../../testing"; import { type GeniePlugin, genie } from "../genie"; /** @@ -11,10 +15,10 @@ import { type GeniePlugin, genie } from "../genie"; * suite drives that real handler and asserts the emitted event ORDER with * `expectStream` — the streaming-assertion path the kit is meant to make easy. * - * See internal/testing-kit-dogfooding.md for the developer-experience notes - * this exercise produced (notably: the kit's `createMockResponse` does not - * capture written SSE bytes, so a small capturing response is needed to bridge - * a `res.write`-based handler into `expectStream`). + * The kit's `createMockResponse` captures the SSE bytes the handler writes, and + * `expectStream` reads them straight back: `expectStream(res).toEmit(...)`. No + * hand-rolled capturing response is needed. See + * internal/testing-kit-dogfooding.md for the wider developer-experience notes. */ // The base Plugin reads the cache singleton on attach; a tiny in-memory stub @@ -35,48 +39,6 @@ vi.mock("../../../cache", () => ({ CacheManager: { getInstanceSync: vi.fn(() => mockCacheInstance) }, })); -/** - * Collects the SSE bytes a handler writes and exposes them as a `Response`, - * so a `res.write`-based handler can be asserted with `expectStream`. - * - * NOTE: this bridge is exactly the friction the dogfooding writeup flags — the - * kit's own `createMockResponse` throws written chunks away, so streaming - * handler tests need this until the kit ships a capturing response. - */ -function createCapturingResponse() { - const chunks: string[] = []; - const listeners: Record void>> = {}; - const res = { - headersSent: false, - writableEnded: false, - statusCode: 200, - status: vi.fn().mockReturnThis(), - json: vi.fn().mockReturnThis(), - setHeader: vi.fn().mockReturnThis(), - flushHeaders: vi.fn().mockReturnThis(), - write: vi.fn((chunk: unknown) => { - chunks.push(String(chunk)); - return true; - }), - end: vi.fn(function (this: { writableEnded: boolean }) { - this.writableEnded = true; - for (const fn of listeners.close ?? []) fn(); - return this; - }), - on: vi.fn((event: string, fn: () => void) => { - listeners[event] ??= []; - listeners[event].push(fn); - return res; - }), - off: vi.fn().mockReturnThis(), - destroy: vi.fn().mockReturnThis(), - }; - return { - res: res as unknown as express.Response, - toResponse: () => new Response(chunks.join("")), - }; -} - function mockReq(body: unknown): express.Request { const headers: Record = { "x-forwarded-access-token": "user-token", @@ -153,7 +115,7 @@ describe("genie plugin — dogfooding the testing kit", () => { }); test("_handleSendMessage streams status -> message -> complete in order", async () => { - const { res, toResponse } = createCapturingResponse(); + const res = createMockResponse(); await ( plugin as unknown as { @@ -162,9 +124,11 @@ describe("genie plugin — dogfooding the testing kit", () => { w: express.Response, ) => Promise; } - )._handleSendMessage(mockReq({ content: "top customers?" }), res); + )._handleSendMessage( + mockReq({ content: "top customers?" }), + res as unknown as express.Response, + ); - // The kit's expectStream parses the real SSE the handler wrote. - await expectStream(toResponse()).toEmit("status", "message", "complete"); + await expectStream(res).toEmit("status", "message", "complete"); }); }); diff --git a/packages/appkit/src/plugins/serving/tests/serving.test.ts b/packages/appkit/src/plugins/serving/tests/serving.test.ts index 216c2b727..697da84c2 100644 --- a/packages/appkit/src/plugins/serving/tests/serving.test.ts +++ b/packages/appkit/src/plugins/serving/tests/serving.test.ts @@ -3,8 +3,8 @@ import { createMockRequest, createMockResponse, createMockRouter, - mockServiceContext, setupDatabricksEnv, + useServiceContextMock, } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ServiceContext } from "../../../context/service-context"; @@ -48,18 +48,18 @@ vi.mock("../../../connectors/serving/client", () => ({ })); describe("Serving Plugin", () => { - let serviceContextMock: Awaited>; + // The service-context spies' setup/teardown are handled by this hook (auto + // beforeEach install + afterEach restore); the block only adds its own env + // and singleton-reset setup around it. + useServiceContextMock(); - beforeEach(async () => { + beforeEach(() => { setupDatabricksEnv(); process.env.DATABRICKS_SERVING_ENDPOINT_NAME = "test-endpoint"; ServiceContext.reset(); - - serviceContextMock = await mockServiceContext(); }); afterEach(() => { - serviceContextMock?.restore(); delete process.env.DATABRICKS_SERVING_ENDPOINT_NAME; vi.restoreAllMocks(); }); diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts index e196e64fb..1ee8b05e5 100644 --- a/packages/appkit/src/testing/expect-stream.ts +++ b/packages/appkit/src/testing/expect-stream.ts @@ -10,17 +10,39 @@ export interface StreamEvent { [key: string]: unknown; } +/** + * A response double that captured what a streaming handler wrote and can + * replay it as a real `Response`. {@link createMockResponse} returns one; this + * structural type lets {@link expectStream} accept it without importing the + * fixtures module (which would form a cycle). + */ +export interface CapturedSSEResponse { + sseResponse(): Response; +} + /** * Anything {@link expectStream} can consume: * - an async event stream (an adapter's `run()`, an SSE reader), * - an already-collected array of events, - * - an SSE `Response` (or a promise of one) — its body is parsed into events. + * - an SSE `Response` (or a promise of one) — its body is parsed into events, + * - a captured mock response ({@link createMockResponse}) — its written SSE + * bytes are parsed into events. */ export type StreamSource = | AsyncIterable | Iterable | Response - | Promise; + | Promise + | CapturedSSEResponse; + +/** Does `value` expose a `sseResponse()` — i.e. is it a captured mock response? */ +function isCapturedSSEResponse(value: unknown): value is CapturedSSEResponse { + return ( + typeof value === "object" && + value !== null && + typeof (value as CapturedSSEResponse).sseResponse === "function" + ); +} /** Assertions over the events collected from a {@link StreamSource}. */ export interface StreamAssertion { @@ -110,11 +132,30 @@ async function collectEventsInner( ): Promise { const resolved = await source; + // A raw SSE body string is a trap: a string is itself an iterable, so it + // would be walked one character at a time. Reject it with a pointer to the + // right input rather than silently producing per-character "events". + if (typeof resolved === "string") { + throw new Error( + "expectStream: received a raw string. Pass a Response, a captured " + + "response from createMockResponse(), or call its sseResponse() — " + + "not the SSE body text (a string iterates one character at a time).", + ); + } + if (resolved instanceof Response) { const text = await resolved.text(); return parseSSEBody(text); } + // A captured mock response ({@link createMockResponse}) — replay the SSE it + // recorded. Checked before the generic iterable branches (it is a plain + // object without an iterator) so a streaming route reads back as events. + if (isCapturedSSEResponse(resolved)) { + const text = await resolved.sseResponse().text(); + return parseSSEBody(text); + } + if (resolved && typeof resolved === "object") { if (Symbol.asyncIterator in resolved) { const events: StreamEvent[] = []; @@ -202,6 +243,13 @@ function isSubsequence(actual: string[], expected: string[]): boolean { * await expectStream(res).toEmit("warehouse_status", "result"); * ``` * + * @example A plugin's streaming route (via {@link createMockResponse}) + * ```ts + * const res = createMockResponse(); + * await plugin._handleStream(req, res); // writes SSE to res + * await expectStream(res).toEmit("status", "result"); + * ``` + * * @example Guard against a non-terminating stream * ```ts * await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 240fdbc6e..0b8d8fe93 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -1,9 +1,9 @@ import type { Span, SpanOptions } from "@opentelemetry/api"; import type { IAppRouter } from "shared"; -import { vi } from "vitest"; +import { afterEach, beforeEach, vi } from "vitest"; +import { CacheManager } from "../cache"; import type { ServiceContextState } from "../context/service-context"; import { ServiceContext } from "../context/service-context"; -import type { UserContext } from "../context/user-context"; import type { InstrumentConfig, ITelemetry } from "../telemetry/types"; // Test fixtures intentionally use loose shapes; `noExplicitAny` is disabled @@ -103,18 +103,64 @@ export function createMockRouter(): { } /** - * Creates a mock Express request object. Carries a default mock - * WorkspaceClient (SQL succeeds, warehouse is RUNNING) on both the user and - * service-principal client slots; override any field via `overrides`. + * On-behalf-of shorthand for {@link createMockRequest}. `true` uses the default + * test user; an object picks the identity. Sets the forwarded headers the real + * `Plugin.asUser` reads (`x-forwarded-access-token`, `x-forwarded-user`, and — + * when given — `x-forwarded-email`), so an OBO test is one flag instead of + * hand-rolled headers. + */ +export type OboOption = + | boolean + | { + /** `x-forwarded-user` — defaults to `"test-user"`. */ + userId?: string; + /** `x-forwarded-access-token` — defaults to `"test-user-token"`. */ + token?: string; + /** `x-forwarded-email` — omitted unless provided. */ + email?: string; + }; + +/** Build the forwarded identity headers an `obo` option implies. */ +function oboHeaders(obo: Exclude): Record { + const opts = obo === true ? {} : obo; + const headers: Record = { + "x-forwarded-access-token": opts.token ?? "test-user-token", + "x-forwarded-user": opts.userId ?? "test-user", + }; + if (opts.email) headers["x-forwarded-email"] = opts.email; + return headers; +} + +/** + * Creates a mock Express request. Pass `overrides` to set `params`, `query`, + * `body`, `headers`, etc. + * + * For on-behalf-of tests, pass `obo` instead of hand-adding forwarded headers — + * `createMockRequest({ obo: true })` sets the identity headers the real + * `asUser` requires. Any explicit `headers` you also pass win over the ones + * `obo` generates, so you can override a single field. + * + * @example + * ```ts + * createMockRequest({ obo: true }); // default test user + token + * createMockRequest({ obo: { userId: "alice" } }); // pick the user + * ``` */ export function createMockRequest(overrides: Any = {}) { const mockWorkspaceClient = createMockWorkspaceClient(); + const { obo, headers: headerOverrides, ...rest } = overrides; + + // `obo` seeds the forwarded identity headers; an explicit `headers` override + // still wins (merged last) so a test can tweak or drop a single field. + const headers = { + ...(obo ? oboHeaders(obo) : {}), + ...headerOverrides, + }; const req = { params: {}, query: {}, body: {}, - headers: {}, userWorkspaceClient: mockWorkspaceClient, serviceWorkspaceClient: mockWorkspaceClient, getWarehouseId: vi.fn().mockResolvedValue("test-warehouse-id"), @@ -122,7 +168,10 @@ export function createMockRequest(overrides: Any = {}) { header: function (name: string) { return this.headers[name.toLowerCase()]; }, - ...overrides, + // `...rest` keeps the original override power over every default above; + // `headers` is applied last as the one managed field (obo + overrides). + ...rest, + headers, }; return req; } @@ -131,9 +180,21 @@ export function createMockRequest(overrides: Any = {}) { * Creates a mock Express response object. `write`/`send`/`setHeader` flip * `headersSent`, `end` flips `writableEnded` and fires any `close` listener — * enough for streaming handlers that branch on those flags. + * + * Every chunk passed to `write` (and a final chunk to `end`) is captured, so a + * streaming route's real SSE output can be replayed: pass the response straight + * to {@link expectStream}, or call `sseResponse()` for a real `Response`. + * + * @example Assert what a streaming route emitted + * ```ts + * const res = createMockResponse(); + * await plugin._handleStream(req, res); + * await expectStream(res).toEmit("status", "result"); + * ``` */ export function createMockResponse() { const eventListeners: Record void>> = {}; + const chunks: string[] = []; const res = { // Flips to true once headers/body have gone out — mirrors Express so @@ -147,7 +208,12 @@ export function createMockResponse() { return this; }), sendStatus: vi.fn().mockReturnThis(), - end: vi.fn(function (this: Any) { + end: vi.fn(function (this: Any, chunk?: unknown) { + // Express allows `end(chunk)` and `end(callback)`; capture only a data + // chunk, never the completion callback. + if (chunk != null && typeof chunk !== "function") { + chunks.push(String(chunk)); + } this.writableEnded = true; if (eventListeners.close) { for (const handler of eventListeners.close) { @@ -156,8 +222,11 @@ export function createMockResponse() { } return this; }), - write: vi.fn(function (this: Any) { + write: vi.fn(function (this: Any, chunk?: unknown) { this.headersSent = true; + if (chunk != null) chunks.push(String(chunk)); + // Return `this` (truthy) rather than a boolean: handlers that gate on + // backpressure (`if (res.write(buf)) …`) then take the no-wait path. return this; }), setHeader: vi.fn(function (this: Any) { @@ -190,6 +259,15 @@ export function createMockResponse() { return this; }), writableEnded: false, + /** + * The SSE body captured so far, as a real `Response` — the bridge from a + * `res.write`-based handler into {@link expectStream}. `expectStream` + * detects this method and calls it for you, so `expectStream(res)` and + * `expectStream(res.sseResponse())` are equivalent. + */ + sseResponse(): Response { + return new Response(chunks.join("")); + }, }; return res; } @@ -204,6 +282,36 @@ export function setupDatabricksEnv(overrides: Record = {}) { Object.assign(process.env, overrides); } +/** + * Clears AppKit's process-wide cache singleton so cached values don't leak + * between tests in the same file. + * + * The cache `attach()` seeds is shared by every test in a file (Vitest isolates + * files, not tests within a file). Call this in `beforeEach` when one test's + * cached value must not be seen by the next, or mid-test to force a cache miss + * before asserting a subsequent hit. + * + * No-ops when the cache has not been initialized yet, so it is safe to call + * before any `attach()`. + * + * @example + * ```ts + * beforeEach(async () => { + * await resetTestCache(); + * }); + * ``` + */ +export async function resetTestCache(): Promise { + let cache: ReturnType; + try { + cache = CacheManager.getInstanceSync(); + } catch { + // Not initialized yet — nothing to clear. + return; + } + await cache.clear(); +} + /** * Context options for running tests with mocked service/user context */ @@ -245,34 +353,19 @@ export function createMockWorkspaceClient() { } /** - * Builds a {@link ServiceContextState} for testing without touching the - * singleton. Use with {@link mockServiceContext} to install it. + * Builds a {@link ServiceContextState} value for testing without touching the + * singleton. Internal building block for {@link mockServiceContext}, which + * installs the state as spies — that installer is the public entry point. */ -export function createMockServiceContext(options: TestContextOptions = {}) { - const serviceContext: ServiceContextState = { - client: (options.serviceDatabricksClient || - createMockWorkspaceClient()) as Any, - serviceUserId: options.serviceUserId || "test-service-user", - warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), - workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), - }; - - return serviceContext; -} - -/** - * Creates a mock UserContext for testing. - */ -export function createMockUserContext( +function buildServiceContextState( options: TestContextOptions = {}, -): UserContext { +): ServiceContextState { return { - client: (options.userDatabricksClient || + client: (options.serviceDatabricksClient || createMockWorkspaceClient()) as Any, - userId: options.userId || "test-user", + serviceUserId: options.serviceUserId || "test-service-user", warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), - isUserContext: true, }; } @@ -285,7 +378,7 @@ export function createMockUserContext( * @returns The mock context plus the spies and a `restore()` helper. */ export function mockServiceContext(options: TestContextOptions = {}) { - const serviceContext = createMockServiceContext(options); + const serviceContext = buildServiceContextState(options); const getSpy = vi .spyOn(ServiceContext, "get") @@ -328,6 +421,63 @@ export function mockServiceContext(options: TestContextOptions = {}) { }; } +/** The handle {@link mockServiceContext} returns (spies + `restore`). */ +export type ServiceContextMock = ReturnType; + +/** + * Registers a fresh {@link mockServiceContext} before each test and restores it + * after — the whole `beforeEach`/`afterEach` dance in one line. + * + * Call it at the top of a `describe` block (or module top-level), NOT inside a + * test: Vitest's `beforeEach`/`afterEach` only register during collection, so a + * call from within a test body registers nothing for that test. + * + * Returns a **live** accessor, not the handle: each `beforeEach` builds fresh + * spies, so reading `.current` inside a test always sees that test's mock. A + * handle captured once would go stale after the first hook runs. + * + * @example + * ```ts + * describe("my plugin", () => { + * const ctx = useServiceContextMock({ warehouseId: "wh-1" }); + * + * test("resolves the warehouse", async () => { + * await myHandler(req, res); + * expect(ctx.current.getSpy).toHaveBeenCalled(); + * }); + * }); + * ``` + * + * @returns `{ current }` — the active {@link ServiceContextMock} for the test. + */ +export function useServiceContextMock(options: TestContextOptions = {}): { + readonly current: ServiceContextMock; +} { + let handle: ServiceContextMock | undefined; + + beforeEach(() => { + handle = mockServiceContext(options); + }); + + afterEach(() => { + handle?.restore(); + handle = undefined; + }); + + return { + get current(): ServiceContextMock { + if (!handle) { + throw new Error( + "useServiceContextMock: no active mock. Call useServiceContextMock() " + + "at the top of a describe block (not inside a test), and read " + + "`.current` from within a test.", + ); + } + return handle; + }, + }; +} + /** * Runs a test function within a mocked service context: installs the mock, * runs `fn`, and restores the singleton afterward. diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts index 83f2f371c..8565e4d5e 100644 --- a/packages/appkit/src/testing/index.ts +++ b/packages/appkit/src/testing/index.ts @@ -42,6 +42,7 @@ // path (../core/plugin-context) that is not part of the package's exports map. export type { PluginContext } from "../core/plugin-context"; export { + type CapturedSSEResponse, type ExpectStreamOptions, expectStream, parseSSEResponse, @@ -55,15 +56,17 @@ export { createMockRequest, createMockResponse, createMockRouter, - createMockServiceContext, createMockTelemetry, - createMockUserContext, createMockWorkspaceClient, createSuccessfulSQLResponse, mockServiceContext, + type OboOption, + resetTestCache, runWithRequestContext, + type ServiceContextMock, setupDatabricksEnv, type TestContextOptions, + useServiceContextMock, } from "./fixtures"; export { createTestPluginContext, diff --git a/packages/appkit/src/testing/tests/expect-stream.test.ts b/packages/appkit/src/testing/tests/expect-stream.test.ts index c91cd176f..78628ce43 100644 --- a/packages/appkit/src/testing/tests/expect-stream.test.ts +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "vitest"; import { expectStream, parseSSEResponse } from "../expect-stream"; +import { createMockResponse } from "../fixtures"; async function* asyncEvents(events: T[]): AsyncGenerator { for (const event of events) { @@ -149,6 +150,82 @@ describe("expectStream — SSE Response", () => { expectStream(res).toEmitExactly("warehouse_status", "result"), ).resolves.toEqual(["warehouse_status", "result"]); }); + + // Data payloads that are not JSON objects. A JSON object spreads its fields + // onto the event; anything else (scalar, array, non-JSON, multi-line) lands + // under a `data` key. These pin the four non-object branches of parseSSEBody. + test("a scalar JSON data value lands under `data`", async () => { + const res = new Response("event: n\ndata: 42\n\n"); + const events = await expectStream(res).collect(); + expect(events[0]).toEqual({ type: "n", data: 42 }); + }); + + test("an array JSON data value lands under `data` (not spread)", async () => { + const res = new Response("event: xs\ndata: [1,2,3]\n\n"); + const events = await expectStream(res).collect(); + expect(events[0]).toEqual({ type: "xs", data: [1, 2, 3] }); + }); + + test("a non-JSON data value is kept as a raw string", async () => { + const res = new Response("event: note\ndata: plain text\n\n"); + const events = await expectStream(res).collect(); + expect(events[0]).toEqual({ type: "note", data: "plain text" }); + }); + + test("multiple data: lines in one frame are joined with newlines", async () => { + // Per the SSE spec, consecutive `data:` lines join with `\n`. Here the + // joined value is not JSON, so it stays a string. + const res = new Response( + "event: multi\ndata: line one\ndata: line two\n\n", + ); + const events = await expectStream(res).collect(); + expect(events[0]).toEqual({ type: "multi", data: "line one\nline two" }); + }); +}); + +describe("expectStream — captured mock response", () => { + // Write SSE frames the way the real SSEWriter does: three writes per frame + // (`id:`, `event:`, `data:`), split across calls, terminated by a blank line. + function writeFrame( + res: ReturnType, + id: number, + event: string, + data: unknown, + ) { + res.write(`id: ${id}\n`); + res.write(`event: ${event}\n`); + res.write(`data: ${JSON.stringify(data)}\n\n`); + } + + test("reads the SSE a handler wrote straight from the mock response", async () => { + const res = createMockResponse(); + writeFrame(res, 0, "warehouse_status", { state: "RUNNING" }); + writeFrame(res, 1, "result", { rows: [] }); + res.end(); + + await expect( + expectStream(res).toEmitExactly("warehouse_status", "result"), + ).resolves.toEqual(["warehouse_status", "result"]); + }); + + test("sseResponse() exposes the same bytes as a real Response", async () => { + const res = createMockResponse(); + writeFrame(res, 0, "result", { ok: true }); + + const events = await expectStream(res.sseResponse()).collect(); + expect(events[0]).toMatchObject({ type: "result", ok: true }); + }); + + test("captures a final chunk passed to end()", async () => { + const res = createMockResponse(); + res.write(`event: a\ndata: {}\n\n`); + res.end(`event: b\ndata: {}\n\n`); + + await expect(expectStream(res).toEmitExactly("a", "b")).resolves.toEqual([ + "a", + "b", + ]); + }); }); describe("expectStream — invalid source", () => { @@ -158,6 +235,15 @@ describe("expectStream — invalid source", () => { expectStream(42 as unknown as never).collect(), ).rejects.toThrow(/async iterable, an iterable, or a Response/); }); + + test("rejects a raw SSE body string with an actionable error", async () => { + // A string is itself iterable (one char at a time), so silently walking it + // would produce per-character "events". The guard must point to the fix. + const body = `event: result\ndata: {"ok":true}\n\n`; + await expect( + expectStream(body as unknown as never).collect(), + ).rejects.toThrow(/raw string.*sseResponse/s); + }); }); describe("expectStream — timeout", () => { diff --git a/packages/appkit/src/testing/tests/fixtures.test.ts b/packages/appkit/src/testing/tests/fixtures.test.ts new file mode 100644 index 000000000..cbc179d1f --- /dev/null +++ b/packages/appkit/src/testing/tests/fixtures.test.ts @@ -0,0 +1,125 @@ +import { afterEach, describe, expect, test, vi } from "vitest"; +import { CacheManager } from "../../cache"; +import { InMemoryStorage } from "../../cache/storage"; +import { ServiceContext } from "../../context"; +import { + createMockRequest, + resetTestCache, + useServiceContextMock, +} from "../fixtures"; + +describe("createMockRequest — obo option", () => { + test("no obo leaves the forwarded identity headers unset", () => { + const req = createMockRequest(); + expect(req.header("x-forwarded-access-token")).toBeUndefined(); + expect(req.header("x-forwarded-user")).toBeUndefined(); + }); + + test("obo: true sets the default test identity headers", () => { + const req = createMockRequest({ obo: true }); + expect(req.header("x-forwarded-access-token")).toBe("test-user-token"); + expect(req.header("x-forwarded-user")).toBe("test-user"); + // email is omitted unless asked for. + expect(req.header("x-forwarded-email")).toBeUndefined(); + }); + + test("obo object picks the identity, including email", () => { + const req = createMockRequest({ + obo: { userId: "alice", token: "tok-1", email: "alice@example.com" }, + }); + expect(req.header("x-forwarded-user")).toBe("alice"); + expect(req.header("x-forwarded-access-token")).toBe("tok-1"); + expect(req.header("x-forwarded-email")).toBe("alice@example.com"); + }); + + test("case-insensitive header lookup mirrors Express", () => { + const req = createMockRequest({ obo: { userId: "bob" } }); + expect(req.header("X-Forwarded-User")).toBe("bob"); + }); + + test("an explicit headers override wins over the obo-generated header", () => { + const req = createMockRequest({ + obo: { userId: "alice" }, + headers: { "x-forwarded-user": "override" }, + }); + // The explicit override wins; the obo token it did not touch remains. + expect(req.header("x-forwarded-user")).toBe("override"); + expect(req.header("x-forwarded-access-token")).toBe("test-user-token"); + }); + + test("other overrides (params, body) still apply alongside obo", () => { + const req = createMockRequest({ + obo: true, + params: { alias: "demo" }, + body: { content: "hi" }, + }); + expect(req.params).toEqual({ alias: "demo" }); + expect(req.body).toEqual({ content: "hi" }); + expect(req.header("x-forwarded-user")).toBe("test-user"); + }); +}); + +describe("resetTestCache", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + test("no-ops when the cache is not initialized", async () => { + // Force the uninitialized branch deterministically (there is no public + // un-initialize), so the try/catch is exercised regardless of test order. + vi.spyOn(CacheManager, "getInstanceSync").mockImplementation(() => { + throw new Error("not initialized"); + }); + await expect(resetTestCache()).resolves.toBeUndefined(); + }); + + test("clears a populated cache", async () => { + // Seed the real singleton the way attach() does, then prove reset empties it. + const cache = await CacheManager.getInstance({ + storage: new InMemoryStorage({}), + }); + await cache.set("k", { hello: "world" }); + expect(await cache.get("k")).toEqual({ hello: "world" }); + + await resetTestCache(); + + expect(await cache.get("k")).toBeNull(); + }); +}); + +describe("useServiceContextMock", () => { + const ctx = useServiceContextMock({ warehouseId: "wh-1" }); + + test(".current exposes the active mock, installed for this test", () => { + // The spy is live: the real singleton getter is replaced. + expect(vi.isMockFunction(ServiceContext.get)).toBe(true); + expect(ctx.current.serviceContext.serviceUserId).toBe("test-service-user"); + // Record a call so the next test can prove it did NOT leak across the + // afterEach restore + fresh beforeEach install. + ServiceContext.get(); + expect(ctx.current.getSpy).toHaveBeenCalledTimes(1); + }); + + test("each test gets a FRESH mock (the accessor is live, not a snapshot)", () => { + // If `.current` returned a stale handle from the first test, this spy would + // already show the call recorded above. A fresh install starts at zero. + expect(ctx.current.getSpy).toHaveBeenCalledTimes(0); + // And options are re-applied each time. + expect(vi.isMockFunction(ServiceContext.get)).toBe(true); + }); +}); + +describe("useServiceContextMock — restores after the block", () => { + // A nested block that uses the hook; after it, the real method is back. + describe("inner", () => { + useServiceContextMock(); + test("spies while active", () => { + expect(vi.isMockFunction(ServiceContext.isInitialized)).toBe(true); + }); + }); + + test("the singleton is un-spied outside the hooked block", () => { + // afterEach in the inner block restored the original method. + expect(vi.isMockFunction(ServiceContext.isInitialized)).toBe(false); + }); +}); diff --git a/packages/appkit/src/testing/tests/test-plugin-context.test.ts b/packages/appkit/src/testing/tests/test-plugin-context.test.ts index 9c845454b..df43b2423 100644 --- a/packages/appkit/src/testing/tests/test-plugin-context.test.ts +++ b/packages/appkit/src/testing/tests/test-plugin-context.test.ts @@ -193,6 +193,53 @@ describe("createTestPluginContext — executeTool runs the REAL user-scoping pat ); expect(result).toBeNull(); }); + + test("reports a tool named like an Object.prototype method as missing", async () => { + // The lookup guard uses Object.hasOwn, not `tools[name] === undefined`, so + // a tool named "constructor"/"toString"/etc. does NOT resolve to the + // inherited prototype method — it is reported missing like any other. This + // pins that guard against being weakened to `in` / `=== undefined`. + const mock = createTestPluginContext({ analytics: { query: [] } }); + + for (const inherited of ["constructor", "toString", "hasOwnProperty"]) { + await expect( + mock.ctx.executeTool(mockReq(), "analytics", inherited, {}), + ).rejects.toThrow(new RegExp(`no fake tool "${inherited}"`)); + } + // None of them reached a tool. + expect(mock.toolCalls.every((c) => c.args !== undefined)).toBe(true); + }); +}); + +describe("createTestPluginContext — asUser dev-mode branch", () => { + test("in development, a token-less request is allowed through (no throw)", async () => { + // The fake asUser mirrors Plugin.asUser's dev-mode behavior: under + // NODE_ENV=development a missing token skips impersonation instead of + // throwing. The rest of the suite runs under NODE_ENV=test, so this is the + // only place that branch is exercised. + const prev = process.env.NODE_ENV; + process.env.NODE_ENV = "development"; + try { + const mock = createTestPluginContext({ analytics: { top_users: [] } }); + + // No forwarded headers at all — would reject in production. + const result = await mock.ctx.executeTool( + mockReq({}), + "analytics", + "top_users", + {}, + ); + + expect(result).toEqual([]); + // It still records the dispatch as an OBO call; userId is unset because + // no user header was present (dev skips impersonation, does not invent one). + expect(mock.toolCalls).toHaveLength(1); + expect(mock.toolCalls[0]).toMatchObject({ asUser: true }); + expect(mock.toolCalls[0]?.userId).toBeUndefined(); + } finally { + process.env.NODE_ENV = prev; + } + }); }); describe("createTestPluginContext — telemetry seam", () => { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index ecbcc2a50..7e8bc4be8 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -344,9 +344,6 @@ importers: vite: specifier: npm:rolldown-vite@7.1.14 version: rolldown-vite@7.1.14(@types/node@25.2.3)(esbuild@0.25.10)(jiti@2.6.1)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) - vitest: - specifier: 3.2.4 - version: 3.2.4(@types/debug@4.1.12)(@types/node@25.2.3)(jiti@2.6.1)(jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6))(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) ws: specifier: 8.21.0 version: 8.21.0(bufferutil@4.0.9) @@ -375,6 +372,9 @@ importers: '@vitejs/plugin-react': specifier: 5.1.1 version: 5.1.1(rolldown-vite@7.1.14(@types/node@25.2.3)(esbuild@0.25.10)(jiti@2.6.1)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2)) + vitest: + specifier: 3.2.4 + version: 3.2.4(@types/debug@4.1.12)(@types/node@25.2.3)(jiti@2.6.1)(jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6))(lightningcss@1.30.2)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.2) packages/appkit-ui: dependencies: diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index 63830f43b..feae223fd 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -16,9 +16,7 @@ export { createMockRequest, createMockResponse, createMockRouter, - createMockServiceContext, createMockTelemetry, - createMockUserContext, createMockWorkspaceClient, createSuccessfulSQLResponse, mockServiceContext, @@ -26,4 +24,5 @@ export { runWithRequestContext, setupDatabricksEnv, type TestContextOptions, + useServiceContextMock, } from "../packages/appkit/src/testing"; From 90288bb768bbe5cc829f9a3e0cb15b0e63da5b28 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Thu, 13 Aug 2026 13:37:58 +0200 Subject: [PATCH 18/22] docs(appkit): move the testing guide under Plugins The testing kit is entirely plugin-scoped (createTestPluginContext, attach(plugin), plugin route/tool/SSE assertions), and the page's own cross-links already pointed into plugins/. Move it next to custom-plugins and fix the relative links. Keep the heading as 'Testing'; the Plugins section supplies the context. Signed-off-by: Galymzhan --- docs/docs/{development => plugins}/testing.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename docs/docs/{development => plugins}/testing.md (100%) diff --git a/docs/docs/development/testing.md b/docs/docs/plugins/testing.md similarity index 100% rename from docs/docs/development/testing.md rename to docs/docs/plugins/testing.md From ac0f0cb668bb2471943ceaadc9ee742186ba1025 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Thu, 13 Aug 2026 14:03:13 +0200 Subject: [PATCH 19/22] test(appkit): fold dogfood tests into plugin suites Address round-2 review: the kit should be the default way to test a plugin, not a parallel '*.kit.test.ts' track. - Fold the three cross-plugin executeTool OBO tests into analytics.test.ts and delete analytics.kit.test.ts. - Upgrade genie.test.ts's SSE test to assert event ORDER via expectStream on genie's real event names (message_start, status, message_result, query_result), replacing brittle write.mock.calls substring checks, and delete genie.kit.test.ts. - Trim the heavy comment narration from the folded-in tests. - Re-export createTestPluginContext and expectStream from the test-helpers shim. - Finish the testing-guide move under plugins/ (sidebar position + links). Signed-off-by: Galymzhan --- docs/docs/plugins/testing.md | 8 +- .../analytics/tests/analytics.kit.test.ts | 92 ------------ .../plugins/analytics/tests/analytics.test.ts | 62 ++++++++ .../src/plugins/genie/tests/genie.kit.test.ts | 134 ------------------ .../src/plugins/genie/tests/genie.test.ts | 28 ++-- tools/test-helpers.ts | 2 + 6 files changed, 79 insertions(+), 247 deletions(-) delete mode 100644 packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts delete mode 100644 packages/appkit/src/plugins/genie/tests/genie.kit.test.ts diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md index 874abb263..bac73a7aa 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -1,5 +1,5 @@ --- -sidebar_position: 7 +sidebar_position: 8 --- # Testing @@ -233,6 +233,6 @@ expect(mock.toolCalls[0]).toMatchObject({ ## See also -- [Local development](./local-development.mdx) — run your app with hot reload while iterating. -- [Custom plugins](../plugins/custom-plugins.md) — build the plugins you test with this kit. -- [Execution context](../plugins/execution-context.md) — how `asUser` and the service principal differ at runtime. +- [Custom plugins](./custom-plugins.md) — build the plugins you test with this kit. +- [Execution context](./execution-context.md) — how `asUser` and the service principal differ at runtime. +- [Local development](../development/local-development.mdx) — run your app with hot reload while iterating. diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts deleted file mode 100644 index e5cd06bee..000000000 --- a/packages/appkit/src/plugins/analytics/tests/analytics.kit.test.ts +++ /dev/null @@ -1,92 +0,0 @@ -import type express from "express"; -import { describe, expect, test } from "vitest"; -import { createMockRequest, createTestPluginContext } from "../../../testing"; - -/** - * Dogfooding `@databricks/appkit/testing` for the cross-plugin tool-call (OBO) - * scenario — the case `createTestPluginContext` is purpose-built for. - * - * A plugin that consumes another plugin's tools calls - * `this.context.executeTool(pluginName, toolName, ...)`. Here we register a - * fake `analytics` provider and drive that dispatch, asserting both the result - * and that it resolved through the user's identity (the on-behalf-of path that - * silent `{ executeTool }` stubs could never verify). - * - * This needs ONLY the kit — no workspace, no ServiceContext, no network — which - * is the sweet spot noted in internal/testing-kit-dogfooding.md. (Driving - * analytics' own SQL handlers, by contrast, still needs the ServiceContext / - * workspace-client fixtures because that work lives behind those seams.) - */ - -describe("analytics as a cross-plugin tool provider — dogfooding the kit", () => { - test("a consumer dispatches analytics.query on-behalf-of the user", async () => { - const rows = [{ customer: "Acme", revenue: 1_000_000 }]; - const mock = createTestPluginContext({ - analytics: { - // A canned result for the analytics `query` tool. - query: (args) => ({ rows, echoedArgs: args }), - }, - }); - - // Simulate what a consumer plugin (e.g. agents) does internally: resolve a - // sibling plugin's tool through the shared PluginContext. - const req = createMockRequest({ - obo: { userId: "analyst@example.com" }, - }) as unknown as express.Request; - const result = await mock.ctx.executeTool(req, "analytics", "query", { - sql: "SELECT * FROM top_customers", - }); - - expect(result).toEqual({ - rows, - echoedArgs: { sql: "SELECT * FROM top_customers" }, - }); - - // Prove the dispatch ran as the end user, not the service principal. - expect(mock.toolCalls).toHaveLength(1); - expect(mock.toolCalls[0]).toMatchObject({ - plugin: "analytics", - tool: "query", - asUser: true, - userId: "analyst@example.com", - }); - }); - - test("a token-less request is rejected before the tool runs", async () => { - const mock = createTestPluginContext({ - analytics: { query: () => ({ rows: [] }) }, - }); - - // No `obo` — a request with no forwarded token must be rejected. - const req = createMockRequest() as unknown as express.Request; - await expect( - mock.ctx.executeTool(req, "analytics", "query", {}), - ).rejects.toThrow(/Missing user token/); - expect(mock.toolCalls).toHaveLength(0); - }); - - test("the per-call timeout the caller forwards actually aborts a slow tool", async () => { - const mock = createTestPluginContext({ - analytics: { - query: (_args, signal) => - new Promise((_resolve, reject) => { - signal?.addEventListener("abort", () => - reject(new Error("aborted by timeout")), - ); - }), - }, - }); - - const req = createMockRequest({ obo: true }) as unknown as express.Request; - await expect( - mock.ctx.executeTool( - req, - "analytics", - "query", - {}, - undefined, - 5, // 5ms timeout - ), - ).rejects.toThrow(/aborted by timeout/); - }); -}); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index aff246145..0c33ca6e7 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -2,6 +2,7 @@ import { createMockRequest, createMockResponse, createMockRouter, + createTestPluginContext, mockServiceContext, setupDatabricksEnv, } from "@tools/test-helpers"; @@ -16,6 +17,7 @@ import { Vector, vectorFromArray, } from "apache-arrow"; +import type express from "express"; import { sql } from "shared"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ServiceContext } from "../../../context/service-context"; @@ -1825,3 +1827,63 @@ describe("Analytics Plugin", () => { }); }); }); + +describe("analytics as a cross-plugin tool provider", () => { + // A consumer plugin (e.g. agents) resolves analytics' tools through the + // shared PluginContext. These drive that dispatch and assert the on-behalf-of + // identity the real executeTool resolves — coverage a bare stub can't give. + test("dispatches analytics.query on behalf of the user", async () => { + const rows = [{ customer: "Acme", revenue: 1_000_000 }]; + const mock = createTestPluginContext({ + analytics: { query: (args) => ({ rows, echoedArgs: args }) }, + }); + + const req = createMockRequest({ + obo: { userId: "analyst@example.com" }, + }) as unknown as express.Request; + const result = await mock.ctx.executeTool(req, "analytics", "query", { + sql: "SELECT * FROM top_customers", + }); + + expect(result).toEqual({ + rows, + echoedArgs: { sql: "SELECT * FROM top_customers" }, + }); + expect(mock.toolCalls[0]).toMatchObject({ + plugin: "analytics", + tool: "query", + asUser: true, + userId: "analyst@example.com", + }); + }); + + test("rejects a token-less request before the tool runs", async () => { + const mock = createTestPluginContext({ + analytics: { query: () => ({ rows: [] }) }, + }); + const req = createMockRequest() as unknown as express.Request; + + await expect( + mock.ctx.executeTool(req, "analytics", "query", {}), + ).rejects.toThrow(/Missing user token/); + expect(mock.toolCalls).toHaveLength(0); + }); + + test("forwards the per-call timeout so a slow tool is aborted", async () => { + const mock = createTestPluginContext({ + analytics: { + query: (_args, signal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject(new Error("aborted by timeout")), + ); + }), + }, + }); + const req = createMockRequest({ obo: true }) as unknown as express.Request; + + await expect( + mock.ctx.executeTool(req, "analytics", "query", {}, undefined, 5), + ).rejects.toThrow(/aborted by timeout/); + }); +}); diff --git a/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts b/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts deleted file mode 100644 index a7ce5ef6a..000000000 --- a/packages/appkit/src/plugins/genie/tests/genie.kit.test.ts +++ /dev/null @@ -1,134 +0,0 @@ -import type express from "express"; -import { afterEach, beforeEach, describe, test, vi } from "vitest"; -import { ServiceContext } from "../../../context"; -import { - createMockResponse, - createTestPluginContext, - expectStream, -} from "../../../testing"; -import { type GeniePlugin, genie } from "../genie"; - -/** - * Dogfooding `@databricks/appkit/testing` on a real core plugin (genie). - * - * Genie's `_handleSendMessage` streams SSE via the base `executeStream`. This - * suite drives that real handler and asserts the emitted event ORDER with - * `expectStream` — the streaming-assertion path the kit is meant to make easy. - * - * The kit's `createMockResponse` captures the SSE bytes the handler writes, and - * `expectStream` reads them straight back: `expectStream(res).toEmit(...)`. No - * hand-rolled capturing response is needed. See - * internal/testing-kit-dogfooding.md for the wider developer-experience notes. - */ - -// The base Plugin reads the cache singleton on attach; a tiny in-memory stub -// keeps this unit-level (mirrors the pattern in the sibling genie.test.ts). -const { mockCacheInstance } = vi.hoisted(() => ({ - mockCacheInstance: { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_k: unknown[], fn: (s?: AbortSignal) => Promise) => fn(), - ), - generateKey: vi.fn((...a: unknown[]) => JSON.stringify(a)), - }, -})); - -vi.mock("../../../cache", () => ({ - CacheManager: { getInstanceSync: vi.fn(() => mockCacheInstance) }, -})); - -function mockReq(body: unknown): express.Request { - const headers: Record = { - "x-forwarded-access-token": "user-token", - "x-forwarded-user": "alice", - }; - return { - params: { alias: "myspace" }, - query: {}, - body, - headers, - header: (name: string) => headers[name.toLowerCase()], - on: vi.fn(), - off: vi.fn(), - } as unknown as express.Request; -} - -describe("genie plugin — dogfooding the testing kit", () => { - let plugin: GeniePlugin; - let serviceContextMock: ReturnType; - - // A minimal ServiceContext stand-in (the kit's mockServiceContext fixture - // covers this, but genie's streaming path only needs a resolvable context). - function mockServiceContextLite() { - const state = { - client: {} as never, - serviceUserId: "sp", - warehouseId: Promise.resolve("wh"), - workspaceId: Promise.resolve("ws"), - }; - const get = vi.spyOn(ServiceContext, "get").mockReturnValue(state); - const isInit = vi - .spyOn(ServiceContext, "isInitialized") - .mockReturnValue(true); - return { - restore: () => { - get.mockRestore(); - isInit.mockRestore(); - }, - }; - } - - beforeEach(async () => { - process.env.DATABRICKS_HOST = "https://test.databricks.com"; - ServiceContext.reset(); - serviceContextMock = mockServiceContextLite(); - - // `genie(...)` returns a { plugin: ctor, config } descriptor for createApp; - // for a unit test, instantiate the class and attach a real PluginContext - // through the kit (seeds cache + flips isReady, the production path). - const GenieCtor = genie({}).plugin as unknown as new ( - c: unknown, - ) => GeniePlugin; - plugin = new GenieCtor({ spaces: { myspace: "space-1" }, timeout: 5000 }); - await createTestPluginContext().attach(plugin); - - // Fake the one real edge — the network connector — to yield a known event - // sequence. Everything else (executeStream, SSE writing) runs for real. - ( - plugin as unknown as { - genieConnector: { - streamSendMessage: (...a: unknown[]) => AsyncGenerator; - }; - } - ).genieConnector.streamSendMessage = async function* () { - yield { type: "status", status: "ASKING_AI" }; - yield { type: "message", content: "Here are your results" }; - yield { type: "complete" }; - }; - }); - - afterEach(() => { - serviceContextMock.restore(); - vi.restoreAllMocks(); - }); - - test("_handleSendMessage streams status -> message -> complete in order", async () => { - const res = createMockResponse(); - - await ( - plugin as unknown as { - _handleSendMessage: ( - r: express.Request, - w: express.Response, - ) => Promise; - } - )._handleSendMessage( - mockReq({ content: "top customers?" }), - res as unknown as express.Response, - ); - - await expectStream(res).toEmit("status", "message", "complete"); - }); -}); diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 0af6c25b3..c00b30418 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -2,6 +2,7 @@ import { createMockRequest, createMockResponse, createMockRouter, + expectStream, mockServiceContext, setupDatabricksEnv, } from "@tools/test-helpers"; @@ -333,23 +334,16 @@ describe("Genie Plugin", () => { "no-cache, no-transform", ); - // Verify SSE events are written - const writeCalls = mockRes.write.mock.calls.map((call: any[]) => call[0]); - const allWritten = writeCalls.join(""); - - // Should have message_start event - expect(allWritten).toContain("message_start"); - expect(allWritten).toContain("new-conv-id"); - - // Should have status events - expect(allWritten).toContain("status"); - expect(allWritten).toContain("ASKING_AI"); - - // Should have message_result event - expect(allWritten).toContain("message_result"); - - // Should have query_result event - expect(allWritten).toContain("query_result"); + // Assert the emitted SSE event ORDER via the kit's expectStream, which + // parses the SSE the handler actually wrote (captured by the mock + // response). This pins genie's real event sequence — a reorder or drop + // fails here, unlike a substring check. + await expectStream(mockRes).toEmit( + "message_start", + "status", + "message_result", + "query_result", + ); expect(mockRes.end).toHaveBeenCalled(); }); diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index feae223fd..5a88312da 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -19,6 +19,8 @@ export { createMockTelemetry, createMockWorkspaceClient, createSuccessfulSQLResponse, + createTestPluginContext, + expectStream, mockServiceContext, parseSSEResponse, runWithRequestContext, From 1ebdee180455e0bfbfe8e09474362ad3fe3819b3 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Thu, 13 Aug 2026 14:35:14 +0200 Subject: [PATCH 20/22] test(appkit): restore toHaveLength(1) on the analytics OBO dispatch test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dogfood fold trimmed expect(mock.toolCalls).toHaveLength(1), so a double-dispatch would no longer fail the happy-path test — and it was inconsistent with the token-less sibling that kept toHaveLength(0). Restore it. Signed-off-by: Galymzhan --- packages/appkit/src/plugins/analytics/tests/analytics.test.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 0c33ca6e7..2c62e4753 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1849,6 +1849,7 @@ describe("analytics as a cross-plugin tool provider", () => { rows, echoedArgs: { sql: "SELECT * FROM top_customers" }, }); + expect(mock.toolCalls).toHaveLength(1); expect(mock.toolCalls[0]).toMatchObject({ plugin: "analytics", tool: "query", From 5cdd207c64cb8578b7a7a5e5f4454ac35f912f68 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Thu, 13 Aug 2026 14:45:47 +0200 Subject: [PATCH 21/22] test(appkit): re-assert genie SSE payloads after the expectStream swap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The toEmit swap pinned event order but dropped the payload values the old substring checks covered (conversationId=new-conv-id, status=ASKING_AI), which aren't asserted elsewhere. Restore them structurally via collect() + toMatchObject — keeping the ordering guarantee without brittle substrings. Signed-off-by: Galymzhan --- .../appkit/src/plugins/genie/tests/genie.test.ts | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index c00b30418..d4f0d0246 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -334,16 +334,23 @@ describe("Genie Plugin", () => { "no-cache, no-transform", ); - // Assert the emitted SSE event ORDER via the kit's expectStream, which - // parses the SSE the handler actually wrote (captured by the mock - // response). This pins genie's real event sequence — a reorder or drop - // fails here, unlike a substring check. + // Assert the emitted SSE via the kit's expectStream, which parses the SSE + // the handler actually wrote (captured by the mock response). toEmit pins + // the real event ORDER; collect() lets us also pin the key payload values + // structurally, not by brittle substring match. await expectStream(mockRes).toEmit( "message_start", "status", "message_result", "query_result", ); + const events = await expectStream(mockRes).collect(); + expect(events.find((e) => e.type === "message_start")).toMatchObject({ + conversationId: "new-conv-id", + }); + expect(events.find((e) => e.type === "status")).toMatchObject({ + status: "ASKING_AI", + }); expect(mockRes.end).toHaveBeenCalled(); }); From 448947fb493881ef8f2aa8fff0e5009049700db5 Mon Sep 17 00:00:00 2001 From: Galymzhan Date: Fri, 14 Aug 2026 10:28:12 +0200 Subject: [PATCH 22/22] fix(appkit): drop fabricated workspace-client fields from createMockRequest MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit createMockRequest returned userWorkspaceClient, serviceWorkspaceClient, getWarehouseId and getWorkspaceId — fields no production code reads (plugins resolve those through getWorkspaceClient()/getWarehouseId() from src/context, which mockServiceContext stands in for). Publishing them via @databricks/appkit/testing would make four inert fields a permanent public promise. The two warehouse cold-start tests (analytics + metric) overrode mockReq.serviceWorkspaceClient.warehouses.get, which the route never reads — so they passed on the default RUNNING client without exercising the warehouse path at all. Route the warehouse client through mockServiceContext (the real seam) so the tests are live, and drop the 'mock WorkspaceClient' claim from the testing guide. Signed-off-by: Galymzhan --- docs/docs/plugins/testing.md | 2 +- .../plugins/analytics/tests/analytics.test.ts | 41 ++++++++++--------- .../plugins/analytics/tests/metric.test.ts | 31 ++++++++------ packages/appkit/src/testing/fixtures.ts | 5 --- 4 files changed, 41 insertions(+), 38 deletions(-) diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md index bac73a7aa..8adb47132 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -145,7 +145,7 @@ await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); The kit re-exports the request/response/context fixtures AppKit uses internally: -- `createMockRequest(overrides?)` / `createMockResponse()` — Express request/response doubles, including the streaming flags (`headersSent`, `writableEnded`) and a mock `WorkspaceClient`. Pass `obo: true` (or `obo: { userId, token, email }`) to set the forwarded identity headers `asUser` requires, instead of hand-adding them. `createMockResponse()` also captures everything a handler writes; pass it to `expectStream` (or call `sseResponse()`) to assert a streaming route's SSE. +- `createMockRequest(overrides?)` / `createMockResponse()` — Express request/response doubles, including the streaming flags (`headersSent`, `writableEnded`). Pass `obo: true` (or `obo: { userId, token, email }`) to set the forwarded identity headers `asUser` requires, instead of hand-adding them. `createMockResponse()` also captures everything a handler writes; pass it to `expectStream` (or call `sseResponse()`) to assert a streaming route's SSE. (Plugins resolve the workspace client through `getWorkspaceClient()`, not the request — use `mockServiceContext` to control it.) - `mockServiceContext(options?)` — spy the `ServiceContext` singleton so code that resolves the service principal or a user context gets test doubles. Call in `beforeEach`, and call the returned `restore()` in `afterEach`. - `useServiceContextMock(options?)` — the same, in one line: it registers the `beforeEach` install and `afterEach` restore for you. Call it at the top of a `describe` block (not inside a test), and read the live `.current` handle from within a test: ```ts diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 2c62e4753..c6166fc06 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1649,7 +1649,7 @@ describe("Analytics Plugin", () => { } }); - test("emits warehouse_status events before the result for a STARTING warehouse", async () => { + test("emits warehouse_status events before the result", async () => { const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -1667,27 +1667,30 @@ describe("Analytics Plugin", () => { const handler = getHandler("POST", "/query/:query_key"); - // Override the default RUNNING mock with a STARTING -> RUNNING sequence - // so the route streams a warehouse_status event before the result. - const warehouseGet = vi - .fn() - .mockResolvedValueOnce({ state: "STARTING" }) - .mockResolvedValueOnce({ state: "RUNNING" }); + // The route resolves its warehouse client via getWorkspaceClient() -> + // ServiceContext (NOT the request), so install it there. A warehouse that + // is already RUNNING still emits one warehouse_status event before the + // result — which is what this test pins, without a poll/sleep cycle. + const warehouseGet = vi.fn().mockResolvedValue({ state: "RUNNING" }); + serviceContextMock.restore(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }), + }, + warehouses: { get: warehouseGet, start: vi.fn() }, + }, + }); const mockReq = createMockRequest({ params: { query_key: "test_query" }, body: { parameters: {} }, }); - mockReq.serviceWorkspaceClient.warehouses.get = warehouseGet; - mockReq.userWorkspaceClient.warehouses.get = warehouseGet; const mockRes = createMockResponse(); - // The connector polls every 3s between warehouse state checks; use fake - // timers so the test doesn't actually sleep. - vi.useFakeTimers(); - const handlerPromise = handler(mockReq, mockRes); - await vi.runAllTimersAsync(); - await handlerPromise; - vi.useRealTimers(); + await handler(mockReq, mockRes); // Inspect the SSE writes: a `warehouse_status` event must precede the // `result` event. @@ -1704,11 +1707,9 @@ describe("Analytics Plugin", () => { expect(resultIdx).toBeGreaterThanOrEqual(0); expect(warehouseIdx).toBeLessThan(resultIdx); - // The status payload should include the state field. + // The status payload should include the RUNNING state. expect(mockRes.write).toHaveBeenCalledWith( - expect.stringMatching( - /"type":"warehouse_status".*"state":"(STARTING|RUNNING)"/, - ), + expect.stringMatching(/"type":"warehouse_status".*"state":"RUNNING"/), ); expect(executeMock).toHaveBeenCalledTimes(1); diff --git a/packages/appkit/src/plugins/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index e704ae309..2337dec32 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -512,7 +512,7 @@ describe("analytics metric route", () => { expect(mockRes.end).toHaveBeenCalled(); }); - test("emits warehouse_status before result for a STARTING warehouse", async () => { + test("emits warehouse_status before result", async () => { const plugin = pluginForDir( config, registryDir({ @@ -533,23 +533,30 @@ describe("analytics metric route", () => { plugin.injectRoutes(router); const handler = getHandler("POST", "/metric/:key"); - const warehouseGet = vi - .fn() - .mockResolvedValueOnce({ state: "STARTING" }) - .mockResolvedValueOnce({ state: "RUNNING" }); + // The route resolves its warehouse client via getWorkspaceClient() -> + // ServiceContext (NOT the request), so install it there. A warehouse that + // is already RUNNING still emits one warehouse_status event before the + // result — which is what this test pins, without a poll/sleep cycle. + const warehouseGet = vi.fn().mockResolvedValue({ state: "RUNNING" }); + serviceContextMock.restore(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }), + }, + warehouses: { get: warehouseGet, start: vi.fn() }, + }, + }); const mockReq = createMockRequest({ params: { key: "revenue" }, body: { measures: ["arr"] }, }); - mockReq.serviceWorkspaceClient.warehouses.get = warehouseGet; - mockReq.userWorkspaceClient.warehouses.get = warehouseGet; const mockRes = createMockResponse(); - vi.useFakeTimers(); - const handlerPromise = handler(mockReq, mockRes); - await vi.runAllTimersAsync(); - await handlerPromise; - vi.useRealTimers(); + await handler(mockReq, mockRes); const eventLines = (mockRes.write as any).mock.calls .map((call: any[]) => call[0] as string) diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 0b8d8fe93..f1c554314 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -147,7 +147,6 @@ function oboHeaders(obo: Exclude): Record { * ``` */ export function createMockRequest(overrides: Any = {}) { - const mockWorkspaceClient = createMockWorkspaceClient(); const { obo, headers: headerOverrides, ...rest } = overrides; // `obo` seeds the forwarded identity headers; an explicit `headers` override @@ -161,10 +160,6 @@ export function createMockRequest(overrides: Any = {}) { params: {}, query: {}, body: {}, - userWorkspaceClient: mockWorkspaceClient, - serviceWorkspaceClient: mockWorkspaceClient, - getWarehouseId: vi.fn().mockResolvedValue("test-warehouse-id"), - getWorkspaceId: vi.fn().mockResolvedValue("test-workspace-id"), header: function (name: string) { return this.headers[name.toLowerCase()]; },