diff --git a/apps/dev-playground/client/package-lock.json b/apps/dev-playground/client/package-lock.json index b24fd982d..d3eec358d 100644 --- a/apps/dev-playground/client/package-lock.json +++ b/apps/dev-playground/client/package-lock.json @@ -49,31 +49,6 @@ "vite": "npm:rolldown-vite@7.1.14" } }, - "../../../packages/appkit-ui": { - "name": "@databricks/appkit-ui", - "version": "1.0.0", - "extraneous": true, - "dependencies": { - "clsx": "^2.1.1", - "shared": "workspace:*", - "tailwind-merge": "^3.4.0" - }, - "devDependencies": { - "@types/react": "^19.0.0", - "@types/react-dom": "^19.0.0", - "react": "^19.0.0", - "react-dom": "^19.0.0", - "recharts": "^3.4.1" - }, - "peerDependencies": { - "react": "^18.0.0 || ^19.0.0", - "react-dom": "^18.0.0 || ^19.0.0", - "recharts": "^2.0.0 || ^3.0.0" - } - }, - "../../../packages/appkit-ui/dist": { - "extraneous": true - }, "node_modules/@alloc/quick-lru": { "version": "5.2.0", "resolved": "https://registry.npmjs.org/@alloc/quick-lru/-/quick-lru-5.2.0.tgz", @@ -6967,27 +6942,6 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } - }, - "package": { - "name": "@databricks/appkit-ui", - "version": "1.0.0", - "extraneous": true, - "dependencies": { - "clsx": "^2.1.1", - "tailwind-merge": "^3.4.0" - }, - "devDependencies": { - "@types/react": "^19.0.0", - "@types/react-dom": "^19.0.0", - "react": "^19.0.0", - "react-dom": "^19.0.0", - "recharts": "^3.4.1" - }, - "peerDependencies": { - "react": "^18.0.0 || ^19.0.0", - "react-dom": "^18.0.0 || ^19.0.0", - "recharts": "^2.0.0 || ^3.0.0" - } } } } diff --git a/docs/docs/api/appkit/Function.createApp.md b/docs/docs/api/appkit/Function.createApp.md index bc656537d..4bc8aca60 100644 --- a/docs/docs/api/appkit/Function.createApp.md +++ b/docs/docs/api/appkit/Function.createApp.md @@ -8,7 +8,7 @@ function createApp(config: { onPluginsReady?: (appkit: PluginMap) => void | Promise; plugins?: T; telemetry?: TelemetryConfig; -}): Promise>; +}): Promise>; ``` Bootstraps AppKit with the provided configuration. @@ -41,7 +41,7 @@ with an `asUser(req)` method for user-scoped execution. ## Returns -`Promise`\<`PluginMap`\<`T`\>\> +`Promise`\<[`AppHandle`](TypeAlias.AppHandle.md)\<`T`\>\> A `PluginMap` keyed by plugin name with typed exports diff --git a/docs/docs/api/appkit/Function.createWorkspaceClient.md b/docs/docs/api/appkit/Function.createWorkspaceClient.md index 8ad87f41d..de3f98837 100644 --- a/docs/docs/api/appkit/Function.createWorkspaceClient.md +++ b/docs/docs/api/appkit/Function.createWorkspaceClient.md @@ -18,8 +18,8 @@ Host resolution: | Parameter | Type | | ------ | ------ | -| `opts` | [`WorkspaceClientOptions`](Interface.WorkspaceClientOptions.md) | +| `opts` | `WorkspaceClientOptions` | ## Returns -[`WorkspaceClient`](Interface.WorkspaceClient.md) +`WorkspaceClient` diff --git a/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md b/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md index 29ef96aac..56229ac37 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClientOptions.md @@ -38,6 +38,16 @@ Databricks host, e.g. https://my-workspace.cloud.databricks.com. Defaults to DAT *** +### profile? + +```ts +optional profile: string; +``` + +`~/.databrickscfg` profile name. Used when no host/token is provided. + +*** + ### token? ```ts diff --git a/docs/docs/api/appkit/TypeAlias.AppHandle.md b/docs/docs/api/appkit/TypeAlias.AppHandle.md new file mode 100644 index 000000000..6d7c304fb --- /dev/null +++ b/docs/docs/api/appkit/TypeAlias.AppHandle.md @@ -0,0 +1,56 @@ +# Type Alias: AppHandle\ + +```ts +type AppHandle = PluginMap & { + [asyncDispose]: Promise; + close: Promise; +}; +``` + +What `createApp()` returns: every plugin's exports keyed by manifest name, +plus the app's own teardown handle. + +`close()` releases what AppKit acquired — sockets, timers, pools, cache, and +telemetry — without terminating the process, so a host can embed AppKit and a +test can boot more than once in a file. + +`Symbol.asyncDispose` is exposed alongside it because a plugin's manifest name +can never be a symbol: `await using app = await createApp(...)` is safe even +if a plugin were somehow named `close`. + +## Type Declaration + +### \[asyncDispose\]() + +```ts +asyncDispose: Promise; +``` + +#### Returns + +`Promise`\<`void`\> + +### close() + +```ts +close(options?: { + timeoutMs?: number; +}): Promise; +``` + +#### Parameters + +| Parameter | Type | Description | +| ------ | ------ | ------ | +| `options?` | \{ `timeoutMs?`: `number`; \} | - | +| `options.timeoutMs?` | `number` | Overall teardown budget. Defaults to AppKit's programmatic budget, which is shorter than the signal path's. | + +#### Returns + +`Promise`\<`void`\> + +## Type Parameters + +| Type Parameter | +| ------ | +| `U` *extends* readonly [`PluginData`](TypeAlias.PluginData.md)\<`PluginConstructor`, `unknown`, `string`\>[] | diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index f39a52db2..873a84cd1 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -104,6 +104,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [AgentTool](TypeAlias.AgentTool.md) | Any tool an agent can invoke: inline function tools (`tool()`), hosted MCP tools (`mcpServer()` / raw hosted), toolkit references from plugins (`analytics().toolkit()`), or adapter-hosted Supervisor-API tools (`supervisorTools.*`). | | [AgentTools](TypeAlias.AgentTools.md) | Per-agent tool record. String keys map to inline tools, toolkit entries, hosted tools, etc. | | [AgentToolsFn](TypeAlias.AgentToolsFn.md) | Function form of `AgentDefinition.tools`. Receives the typed [Plugins](TypeAlias.Plugins.md) map and returns a tool record. Invoked exactly once at setup (or once per `runAgent` call in standalone mode); the result is cached as the agent's resolved tool record. | +| [AppHandle](TypeAlias.AppHandle.md) | What `createApp()` returns: every plugin's exports keyed by manifest name, plus the app's own teardown handle. | | [BaseSystemPromptOption](TypeAlias.BaseSystemPromptOption.md) | - | | [ConfigSchema](TypeAlias.ConfigSchema.md) | Configuration schema definition for plugin config. Re-exported from the standard JSON Schema Draft 7 types. | | [ExecutionResult](TypeAlias.ExecutionResult.md) | Discriminated union for plugin execution results. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 18a5333b1..f9253a963 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -433,6 +433,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/TypeAlias.AgentToolsFn", label: "AgentToolsFn" }, + { + type: "doc", + id: "api/appkit/TypeAlias.AppHandle", + label: "AppHandle" + }, { type: "doc", id: "api/appkit/TypeAlias.BaseSystemPromptOption", diff --git a/docs/docs/plugins/custom-plugins.md b/docs/docs/plugins/custom-plugins.md index ccff2eff5..343734c97 100644 --- a/docs/docs/plugins/custom-plugins.md +++ b/docs/docs/plugins/custom-plugins.md @@ -78,6 +78,14 @@ export const myPlugin = toPlugin(MyPlugin); JSON is the canonical authoring surface — it is what `appkit plugin sync` reads when aggregating manifests for templates. For the full v2.0 manifest contract (resources, discovery descriptors, scaffolding rules), see [Plugin manifest](./manifest.md). +:::note Reserved plugin names +`close` cannot be used as a plugin `name`. Plugin exports are installed as own +properties on the object `createApp()` returns, and an own property shadows a +prototype method — so a plugin named `close` would silently replace the app +handle's own `close()` and break teardown. `createApp()` rejects it with a +`ConfigurationError` naming the plugin instead of failing quietly at shutdown. +::: + ## Config-dependent resources The manifest defines resources as either `required` (always needed) or `optional` (may be needed). diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md new file mode 100644 index 000000000..9310eb817 --- /dev/null +++ b/docs/docs/plugins/testing.md @@ -0,0 +1,408 @@ +--- +sidebar_position: 8 +--- + +# 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 three entry points plus a set of fixture helpers: + +- **`createTestApp({ plugins })`** — boot a real app and call it over real HTTP. Start here. +- **`createTestPluginContext()`** — build a real `PluginContext` with faked edges and attach it to a plugin, with no boot and no socket. +- **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. +- **Fixtures** — `createMockRequest`, `createMockResponse`, `createMockWorkspaceClient`, `mockServiceContext`, and SQL response builders. + +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. + +## Testing your plugin + +`createTestApp({ plugins })` boots a **real** AppKit app — real Express wiring, real routes, real resource validation — and hands you methods to call it like a client would: + +```ts +import { createTestApp, expectStream } from "@databricks/appkit/testing"; + +test("my plugin answers a request", async () => { + const app = await createTestApp({ plugins: [myPlugin()] }); + try { + const res = await app.post("/api/my-plugin/thing", { body: { q: 1 }, obo: true }); + expect(res.status).toBe(200); + await expectStream(res).toEmit("status", "result"); + } finally { + await app.close(); + } +}); +``` + +No workspace, no credentials, no network. The harness pins a non-development `NODE_ENV`, binds an ephemeral port, installs a fake workspace client, and keeps the cache in memory so nothing reaches out. + +Paths are the full mounted route. A plugin's prefix is `/api/` plus its manifest name in kebab-case, so a plugin named `mySearch` serves at `/api/my-search/…`. + +### Which harness? + +| | `createTestApp` | `createTestPluginContext` | +| --- | --- | --- | +| Boots the app | Yes | No | +| Binds a socket | Yes (ephemeral port) | No | +| Express middleware, error handler | Real | Not involved | +| Resource / env validation | Real, and strict | Not involved | +| Workspace client | Faked and injected | Fake it yourself with `mockServiceContext` | +| Needs `close()` | **Yes** | No | +| Speed | Fast, but pays for a socket | Fastest | + +Use `createTestApp` for a plugin's HTTP behaviour end to end. Use `createTestPluginContext` to unit-test wiring — route registration, tool dispatch, timeout composition. Name harness suites `*.integration.test.ts`, matching the existing convention. + +### Faking what your plugin reads + +Declare responses by dotted path — `"."` on AppKit's workspace-client facade: + +```ts +const app = await createTestApp({ + plugins: [myPlugin()], + responses: { + "jobs.getRun": { state: "TERMINATED", result_state: "SUCCESS" }, + "statementExecution.executeStatement": { status: { state: "SUCCEEDED" } }, + "apiClient.request": { results: [] }, + }, +}); +``` + +A function value receives the call arguments, so you can script per-argument behaviour or reject to test an error path. Any path you **don't** declare resolves `undefined` rather than crashing — see [Mocking Databricks services](#mocking-databricks-services) for the trade-off that buys. + +For the response *shapes*, follow the service types on the Databricks SDK — the kit doesn't validate them, so a wrong shape fails in your plugin, not in the fake. + +`app.client` is the very object your handler resolves at runtime — reached inside a plugin via `getExecutionContext().client` — so you can assert calls on it: + +```ts +import { getMockFn } from "@databricks/appkit/testing"; + +expect(getMockFn(app.client, "jobs.getRun")).toHaveBeenCalledWith({ run_id: 42 }); +``` + +`getMockFn` exists because facade accessors are typed against the SDK, so `expect(app.client.jobs.getRun).toHaveBeenCalled()` won't typecheck. + +### Requests + +`app.get/post/put/patch/delete(path, options?)` return a native `Response`, so `expectStream` composes directly with no bridge. + +- `body` — a non-string value is JSON-encoded with `content-type: application/json`. A string is sent as-is. +- `headers` — merged last, so they win over anything the harness set. +- `obo` — `true` for the default test user, or `{ userId, token, email }`. Same shorthand as `createMockRequest({ obo })`, so a handler using `asUser(req)` resolves that identity. +- `signal` — forwarded to `fetch`. + +### Teardown + +The harness binds a socket and installs signal handlers, so **every boot needs a `close()`**. `close()` releases the socket, runs your plugin's `shutdown()` hooks, drops AppKit's singletons, and restores `process.env` to its pre-boot state. It's idempotent. + +Use `try/finally`, or let the runtime do it: + +```ts +await using app = await createTestApp({ plugins: [myPlugin()] }); +// released at scope exit +``` + +Skip the `close()` and you'll leak a listener per boot — Node warns at about six. + +### Satisfying declared resources + +The harness runs the real validator with a strict posture, so a plugin whose manifest requires a resource fails the boot unless its env var is set. Supply it with `env`: + +```ts +// Throws: MY_WAREHOUSE_ID is required by the manifest. +await createTestApp({ plugins: [myPlugin()] }); + +// Boots. +await createTestApp({ plugins: [myPlugin()], env: { MY_WAREHOUSE_ID: "w-1" } }); +``` + +That makes "my plugin declares its resources correctly" a genuine assertion. `env` is restored on `close()`. + +:::note What this does not check +The harness validates that required resources' **environment variables are present**. It does **not** validate config *values* against your manifest's `config.schema` — no runtime validator exists for that yet. A test that boots successfully tells you your resource declarations and env are wired up; it says nothing about whether your config values are well-formed. +::: + +### Other options + +- `server: false` — no socket. Plugin setup, validation, and teardown still run; the request methods throw if called. Useful when you only care that a plugin boots. +- `client` — supply your own workspace client instead of the built-in fake. You then own its `currentUser.me()`: AppKit reads `currentUser.id` during boot and can't start without it. +- `nodeEnv` — defaults to `"test"`. `"development"` is **refused**: dev mode routes the harness's ephemeral port through `get-port`, which throws on port `0`, and it also boots a real Vite server and relaxes validation. +- `cache` — defaults to in-memory. Overriding it is what would let the cache reach the network, so leave it alone unless that's the point of the test. +- `closeTimeoutMs` — teardown budget. + +## `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. `createTestPluginContext()` 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 { createTestPluginContext } from "@databricks/appkit/testing"; + +const mock = createTestPluginContext({ + 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 = 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, 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 + +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 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(); +``` + +`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. + +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, 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"; + +// 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(); +``` + +### 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: + +```ts +await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); +``` + +## Fixtures + +AppKit has two contexts, and they're faked by different tools. `PluginContext` is the mediator between plugins — routes, tool dispatch, user scoping — and `createTestPluginContext()` gives you the real thing with faked edges. `ServiceContext` is the **data plane**: it resolves the workspace client, the service principal, and the warehouse ID that plugins reach through `getWorkspaceClient()`. + +The kit now covers both. `createTestApp` fakes the data plane for you by injecting a mock workspace client at the real seam; below that, `mockServiceContext` spies the singleton directly, and `createMockWorkspaceClient` builds the client either of them installs. + +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`). 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 + 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. +- `resetAppKitSingletons()` — drop AppKit's process-wide singletons so a later `createApp` builds fresh ones. `createTestApp`'s `close()` already does this; you need it only if you call `createApp` yourself. Close first, then reset — it drops pointers, it doesn't release resources. +- `createTestPlugin(factory, config?)` — instantiate a plugin from its factory with the same config merge AppKit applies. See [Full example](#full-example). + +## Mocking Databricks services + +Every core plugin's real work goes through `getWorkspaceClient()`. `createMockWorkspaceClient()` fakes that whole surface, so a plugin touching `jobs`, `genie`, `servingEndpoints`, or `files` is testable without hand-building a nested client: + +```ts +import { createMockWorkspaceClient, getMockFn } from "@databricks/appkit/testing"; + +const client = createMockWorkspaceClient({ + responses: { "jobs.getRun": { state: "TERMINATED" } }, + config: { host: "https://my-test-host.example.com" }, +}); + +await client.jobs.getRun({ run_id: 1 }); // → { state: "TERMINATED" } +await client.genie.getMessage({ id: "m-1" }); // → undefined, does not throw +``` + +`createTestApp` installs one of these for you, so reach for it directly only when you're driving a plugin through `createTestPluginContext` or `mockServiceContext`. + +How it works, and what to expect: + +- The **facade is typed**, so `client.jbos` is a compile error. AppKit owns that 9-member interface, so it's a closed set, not an open-ended chase of the SDK. +- Each **service** is a proxy that mints a memoized mock per method. `client.jobs.getRun === client.jobs.getRun`, so call assertions are stable, and `toLegacyWorkspaceClient()` shares the same functions — one `responses` entry covers both views. +- `config.host` is a real **string** (not a mock), because AppKit builds URLs from it. `apiClient.userAgent()` is synchronous for the same reason, and `apiClient.request` resolves `{}` so destructuring its result doesn't throw. +- Sensible defaults are built in: SQL statements succeed, warehouses report `RUNNING`, and `currentUser.me()` returns a service user. Pass `defaults: false` to script everything yourself. + +:::caution The honest catch +An undeclared method resolves `undefined` instead of throwing. That's the point — your plugin survives touching services the test doesn't care about — but it means a call whose response you *forgot* to declare silently returns `undefined` rather than failing loudly, so a test can pass for the wrong reason. + +TypeScript covers more of this than you might expect: because each accessor is typed against the SDK's own service class, both a misspelled **service** (`client.jbos`) and a misspelled **method** (`client.jobs.getRunz`) are compile errors. The gap is a *real* method with no declared response — and any call that bypasses the types with a cast. + +One more divergence to know about: a service's methods are minted on access, so they are **callable but not enumerable**. `typeof client.jobs.getRun` is `"function"`, but `'getRun' in client.jobs` is `false` and `Object.keys(client.jobs)` is `[]`. Plugin code that feature-detects with `in` or reflects over a service will therefore take a different branch than it does in production. This is a deliberate trade: reporting those keys would make `util.inspect` probe each one, minting a mock per probe, which is the runaway recursion the default traps avoid. + +Separately, `createLakebasePool({ workspaceClient })` will build a pool whose password callback resolves to a mock: the pool exists but cannot connect. A Lakebase test needs a real database or a purpose-built fake pool, not this. +::: + +## Full example + +For a plugin you wrote, instantiate the class directly with `new`. The `analytics()` / `agents()` factory functions you pass to `createApp` return a *descriptor* for the app to construct, not an instance. + +When you want an instance from one of those factories, use `createTestPlugin` rather than reaching through the descriptor: + +```ts +import { createTestPlugin } from "@databricks/appkit/testing"; + +const plugin = createTestPlugin(genie, { spaceId: "s-1" }); + +// Not this — it skips DEFAULT_CONFIG and forgets `name`, so the instance is +// configured differently from the one production builds: +// const plugin = new (genie({}).plugin)({ spaceId: "s-1" }); +``` + +`createTestPlugin` applies the same merge AppKit does at registration: `DEFAULT_CONFIG`, then your config, then the manifest `name`. It's for this unit-test path only — `createTestApp` takes descriptors and builds the instances itself. + +```ts +import { Plugin, type PluginManifest } from "@databricks/appkit"; +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. +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 = createTestPluginContext(); + const plugin = new GreeterPlugin({}); + + await mock.attach(plugin); + await plugin.setup(); + + expect(mock.routes).toContainEqual( + expect.objectContaining({ method: "get", path: "/hello" }), + ); + }); + + 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 = 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({ + plugin: "analytics", + tool: "query", + asUser: true, +}); +``` + +## See also + +- [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/docs/static/appkit-ui/styles.gen.css b/docs/static/appkit-ui/styles.gen.css index 9e1d5c0c0..579cdbecd 100644 --- a/docs/static/appkit-ui/styles.gen.css +++ b/docs/static/appkit-ui/styles.gen.css @@ -398,6 +398,9 @@ .\!m-0 { margin: calc(var(--spacing) * 0) !important; } + .m-1 { + margin: calc(var(--spacing) * 1); + } .-mx-1 { margin-inline: calc(var(--spacing) * -1); } @@ -714,6 +717,9 @@ .w-\(--sidebar-width\) { width: var(--sidebar-width); } + .w-1 { + width: calc(var(--spacing) * 1); + } .w-1\/2 { width: calc(1/2 * 100%); } 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 a6b259762..feb83a151 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" @@ -94,6 +99,14 @@ "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", @@ -101,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" @@ -113,6 +127,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/cache/index.ts b/packages/appkit/src/cache/index.ts index a98b7dff4..9fc37687a 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -58,6 +58,11 @@ export class CacheManager { private readonly name: string = "cache-manager"; private static instance: CacheManager | null = null; private static initPromise: Promise | null = null; + /** + * Bumped by {@link reset}, so an initialization already in flight cannot + * publish its result over a caller that has since discarded the singleton. + */ + private static generation = 0; private storage: CacheStorage; private config: CacheConfig; @@ -126,9 +131,16 @@ export class CacheManager { } if (!CacheManager.initPromise) { + const generation = CacheManager.generation; CacheManager.initPromise = CacheManager.create(userConfig).then( (instance) => { - CacheManager.instance = instance; + // A reset() while this was in flight means the caller discarded this + // manager before it existed. Installing it anyway would resurrect it + // and hand the next boot storage the caller never asked for, so the + // result is returned to whoever is awaiting but not published. + if (CacheManager.generation === generation) { + CacheManager.instance = instance; + } return instance; }, ); @@ -557,6 +569,31 @@ export class CacheManager { await this.storage.close(); } + /** + * Drop the singleton so the next {@link getInstance} builds a fresh manager. + * + * Both fields must be cleared: `getInstance()` returns `initPromise` when + * `instance` is null, so clearing only `instance` would leave the next boot + * awaiting a promise that resolves to the dead manager. + * + * Clearing `initPromise` is not enough on its own either: an initialization + * already in flight would still run its continuation and re-publish the very + * instance being discarded. The generation counter is what makes the reset + * hold in that case. + * + * Deliberately does **not** call {@link close} — a reset is a pointer drop, + * `close()` is I/O. Callers close first, then reset, which is the order + * `LifecycleManager.close()` uses. Resetting without closing leaks whatever + * the old storage held (under `PersistentStorage` that is a `pg.Pool`). + * + * @internal + */ + static reset(): void { + CacheManager.instance = null; + CacheManager.initPromise = null; + CacheManager.generation += 1; + } + /** * Check if the storage is healthy * @returns Promise of true if the storage is healthy, false otherwise diff --git a/packages/appkit/src/cache/tests/cache-manager-reset.test.ts b/packages/appkit/src/cache/tests/cache-manager-reset.test.ts new file mode 100644 index 000000000..dcb4d0aa5 --- /dev/null +++ b/packages/appkit/src/cache/tests/cache-manager-reset.test.ts @@ -0,0 +1,149 @@ +import { afterEach, beforeEach, describe, expect, test } from "vitest"; + +import { CacheManager } from ".."; +import { InitializationError } from "../../errors"; +import { InMemoryStorage } from "../storage/memory"; + +/** + * `getInstance()` returns the existing instance, so after `cache.close()` the + * singleton still points at closed storage — under `PersistentStorage` an ended + * `pg.Pool`. Every test passes explicit `storage` so nothing probes Lakebase. + */ +describe("CacheManager.reset", () => { + beforeEach(() => { + CacheManager.reset(); + }); + + afterEach(() => { + CacheManager.reset(); + }); + + function storage() { + return new InMemoryStorage({ enabled: true, maxSize: 100 } as never); + } + + test("the next getInstance() builds a fresh instance, not the closed one", async () => { + const first = await CacheManager.getInstance({ storage: storage() }); + await first.close(); + + CacheManager.reset(); + const second = await CacheManager.getInstance({ storage: storage() }); + + expect(second).not.toBe(first); + + // The point of the fix: the fresh instance's storage is live, so a + // write-then-read round-trips instead of hitting closed storage. + const key = second.generateKey(["reset-probe"], "test-user"); + await second.set(key, { ok: true }); + await expect(second.get(key)).resolves.toEqual({ ok: true }); + }); + + test("without a reset, getInstance() keeps returning the same instance", async () => { + // The regression guard for the *unchanged* path: a single boot with no reset + // must behave exactly as before. + const first = await CacheManager.getInstance({ storage: storage() }); + const second = await CacheManager.getInstance({ storage: storage() }); + + expect(second).toBe(first); + }); + + test("reset clears an in-flight initPromise, not just the instance", async () => { + // Start initialization but do not await it, so `instance` is still null and + // only `initPromise` is set. Clearing just `instance` would leave the next + // caller awaiting a promise that resolves to the discarded manager — + // getInstance() returns initPromise when instance is null. + const pending = CacheManager.getInstance({ storage: storage() }); + + CacheManager.reset(); + + const first = await pending; + const second = await CacheManager.getInstance({ storage: storage() }); + + expect(second).not.toBe(first); + }); + + test("getInstanceSync throws after a reset", async () => { + await CacheManager.getInstance({ storage: storage() }); + expect(() => CacheManager.getInstanceSync()).not.toThrow(); + + CacheManager.reset(); + + // Reset is a pointer drop, so the sync accessor is back to its + // not-initialized contract rather than handing out a stale manager. + expect(() => CacheManager.getInstanceSync()).toThrow(InitializationError); + }); + + test("reset is safe when the cache was never initialized", () => { + expect(() => CacheManager.reset()).not.toThrow(); + expect(() => CacheManager.reset()).not.toThrow(); + }); + test("without a reset, the next boot reuses storage the last teardown closed", async () => { + // Models PersistentStorage, whose close() is `pool.end()` — permanent. + // InMemoryStorage.close() merely clears a Map and stays usable, which is why + // an in-memory test cannot show this and why the bug hid for so long. + function endableStorage() { + let ended = false; + const entries = new Map(); + const guard = () => { + if (ended) throw new Error("Cannot use a pool after calling end()"); + }; + return { + get: async (key: string) => { + guard(); + return (entries.get(key) ?? null) as never; + }, + set: async (key: string, entry: unknown) => { + guard(); + entries.set(key, entry); + }, + delete: async (key: string) => { + guard(); + entries.delete(key); + }, + clear: async () => { + guard(); + entries.clear(); + }, + has: async (key: string) => { + guard(); + return entries.has(key); + }, + size: async () => { + guard(); + return entries.size; + }, + isPersistent: () => true, + healthCheck: async () => !ended, + close: async () => { + ended = true; + }, + }; + } + + const first = await CacheManager.getInstance({ + storage: endableStorage() as never, + }); + await first.close(); + + // The bug, with no reset in between: getInstance() hands back the same + // manager, still pointing at storage that has been ended. + const stale = await CacheManager.getInstance({ + storage: endableStorage() as never, + }); + expect(stale).toBe(first); + await expect( + stale.set(stale.generateKey(["x"], "test-user"), { v: 1 }), + ).rejects.toThrow(/after calling end/); + + // The fix: reset drops the pointer, so the next boot builds over live + // storage and the same write succeeds. + CacheManager.reset(); + const fresh = await CacheManager.getInstance({ + storage: endableStorage() as never, + }); + expect(fresh).not.toBe(first); + const key = fresh.generateKey(["x"], "test-user"); + await fresh.set(key, { v: 1 }); + await expect(fresh.get(key)).resolves.toEqual({ v: 1 }); + }); +}); diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 0ccfd64e9..61b443344 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -1,4 +1,5 @@ import type { + AppHandle, BasePlugin, CacheConfig, InputPluginMap, @@ -11,6 +12,7 @@ import type { import { version as productVersion } from "../../package.json"; import { CacheManager } from "../cache"; import { ServiceContext } from "../context"; +import { ConfigurationError } from "../errors"; import { isInternalTelemetryEnabled, TelemetryReporter, @@ -27,10 +29,24 @@ import { isToolProvider, PluginContext } from "./plugin-context"; const logger = createLogger("appkit"); +/** + * Names a plugin manifest may not use, because `createAndRegisterPlugin` + * installs exports as **own** properties and an own property shadows a + * prototype method. A plugin named `close` would therefore silently break + * teardown rather than merely confusing the types, so registration fails loudly + * instead. + */ +const RESERVED_PLUGIN_NAMES = new Set(["close"]); + export class AppKit { #pluginInstances: Record = {}; #setupPromises: Promise[] = []; #context: PluginContext; + /** + * Retained so {@link close} can reach the shutdown sequence. Assigned once + * every plugin has started; `close()` before that point is a no-op teardown. + */ + #lifecycle: LifecycleManager | undefined; private constructor(config: { plugins: TPlugins }) { const { plugins, ...globalConfig } = config; @@ -80,6 +96,15 @@ export class AppKit { pluginData: OptionalConfigPluginDef, extraData?: Record, ) { + if (RESERVED_PLUGIN_NAMES.has(name)) { + throw new ConfigurationError( + `Plugin name "${name}" is reserved by the app handle returned from ` + + "createApp(). Rename the plugin in its manifest — an own property " + + "would shadow the handle's method and silently break app teardown.", + { context: { pluginName: name } }, + ); + } + const { plugin: Plugin, config: pluginConfig } = pluginData; const baseConfig = { ...config, @@ -188,10 +213,28 @@ export class AppKit { telemetry?: TelemetryConfig; cache?: CacheConfig; client?: WorkspaceClient; + /** + * Runs after plugin setup but **before** the server starts. + * + * Deliberately typed `PluginMap` rather than the `AppHandle` + * that `createApp` returns: at this point the app is not fully + * started, so offering `close()` here would invite tearing down a + * half-booted app. The value passed at runtime is the same object — + * the narrower type is the point, not an oversight. + */ + /** + * Runs after plugin setup but **before** the server starts. + * + * Deliberately typed `PluginMap` rather than the `AppHandle` + * that `createApp` returns: at this point the app is not fully + * started, so offering `close()` here would invite tearing down a + * half-booted app. The value passed at runtime is the same object — + * the narrower type is the point, not an oversight. + */ onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; } = {}, - ): Promise> { + ): Promise> { // Initialize core services TelemetryManager.initialize(config?.telemetry); await CacheManager.getInstance(config?.cache); @@ -225,7 +268,7 @@ export class AppKit { await Promise.all(instance.#setupPromises); await instance.#context.emitLifecycle("setup:complete"); - const handle = instance as unknown as PluginMap; + const handle = instance as unknown as AppHandle; if (config.onPluginsReady) { logger.debug("Running onPluginsReady hook"); @@ -246,11 +289,41 @@ export class AppKit { // plugin has started. Applies uniformly whether or not a server plugin // is present — server-less apps still get their telemetry flushed and // plugin shutdown() hooks run. - new LifecycleManager(instance.#context).installSignalHandlers(); + // + // Retained on the instance (rather than discarded) so the returned handle's + // close() can reach the same sequence without a signal. + instance.#lifecycle = new LifecycleManager(instance.#context); + instance.#lifecycle.installSignalHandlers(); return handle; } + /** + * Release everything this app acquired — sockets, timers, pools, cache, and + * telemetry — without terminating the process. + * + * Runs the same phases as a SIGTERM shutdown (plugin `abortActiveOperations` + * and `shutdown()` hooks, the `"shutdown"` lifecycle event, cache close, + * telemetry flush) and detaches the signal handlers this app installed. + * Idempotent: repeated calls await the same teardown. + * + * @param options.timeoutMs - Overall budget. Defaults to the shorter + * programmatic budget, not the production signal budget. + */ + async close(options: { timeoutMs?: number } = {}): Promise { + await this.#lifecycle?.close(options); + } + + /** + * Enables `await using app = await createApp(...)`, which releases the app at + * scope exit. A manifest name can never be a symbol, so this entry point + * cannot be shadowed by a plugin — unlike {@link close}, which is why + * `"close"` is a reserved plugin name. + */ + async [Symbol.asyncDispose](): Promise { + await this.close(); + } + private static bootstrapInternalTelemetry(): void { const serviceCtx = ServiceContext.get(); const reporter = TelemetryReporter.initialize({ @@ -385,6 +458,6 @@ export async function createApp< onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; } = {}, -): Promise> { +): Promise> { return AppKit._createApp(config); } diff --git a/packages/appkit/src/core/lifecycle-manager.ts b/packages/appkit/src/core/lifecycle-manager.ts index 84dcb4a9e..f8c681426 100644 --- a/packages/appkit/src/core/lifecycle-manager.ts +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -5,6 +5,7 @@ import { TelemetryReporter } from "../internal-telemetry"; import { createLogger } from "../logging/logger"; import { TelemetryManager } from "../telemetry"; import type { PluginContext } from "./plugin-context"; +import { resetCoreSingletons } from "./reset-singletons"; const logger = createLogger("lifecycle"); @@ -45,35 +46,47 @@ export class LifecycleManager { */ private static readonly PHASE_SHUTDOWN_TIMEOUT_MS = 2_000; + /** Shorter than the signal path's: a programmatic caller wants its await back. */ + private static readonly CLOSE_TIMEOUT_MS = 5_000; + /** - * Guards against re-entrant shutdown (e.g. SIGTERM followed by SIGINT). - * The flag set in `shutdown` must remain synchronous and first — any - * `await` before it would open a window for a second signal to re-enter - * the sequence. - */ - private isShuttingDown = false; - /** - * Name of the shutdown phase currently in flight, so the force-exit log - * can say where shutdown got stuck without extra bookkeeping. + * The in-flight teardown, memoized. A boolean guard would let a second caller + * return while teardown was still running — fine for a signal, wrong for + * `close()`, which must not resolve before resources are released. */ + private teardown: Promise | undefined; + /** Reported by the force-exit log so a stuck shutdown names its phase. */ private shutdownPhase = "not started"; + /** Retained so {@link close} removes its own listeners and nothing else. */ + private signalHandlers: [NodeJS.Signals, () => void][] = []; constructor(private readonly context: PluginContext) {} + /** Install the SIGTERM/SIGINT handlers that trigger {@link shutdown}. */ + installSignalHandlers(): void { + this.signalHandlers = [ + ["SIGTERM", () => void this.shutdown()], + ["SIGINT", () => void this.shutdown()], + ]; + for (const [signal, handler] of this.signalHandlers) { + process.once(signal, handler); + } + } + /** - * Install the SIGTERM/SIGINT handlers that trigger {@link shutdown}. - * - * Uses `process.once` (not `on`) so a repeated signal cannot register the - * handler twice; re-entrancy from a *different* signal is guarded by - * `isShuttingDown` inside {@link shutdown}. + * Detach only this instance's handlers — never `removeAllListeners`, so an + * embedding host keeps its own. */ - installSignalHandlers(): void { - process.once("SIGTERM", () => this.shutdown()); - process.once("SIGINT", () => this.shutdown()); + removeSignalHandlers(): void { + for (const [signal, handler] of this.signalHandlers) { + process.removeListener(signal, handler); + } + this.signalHandlers = []; } /** - * Run the graceful-shutdown sequence and exit the process. + * Run the graceful-shutdown sequence and **exit the process**. See + * {@link close} for the non-exiting twin. * * Phases: * 1. stop the internal-telemetry reporter @@ -86,22 +99,15 @@ export class LifecycleManager { * Exits 0 on completion (and on the force-exit backstop): a deliberate * shutdown is not a crash. Exit 1 is reserved for an unexpected error * thrown by the sequence itself. + * + * A second signal now awaits the first teardown rather than returning at once; + * the first caller still exits, so this is unobservable in production. */ async shutdown(): Promise { - // Must stay synchronous and first: any await before the flag is set - // would let a second signal re-enter the shutdown sequence. - if (this.isShuttingDown) return; - this.isShuttingDown = true; - - logger.info("Starting graceful shutdown..."); - - let exitCode = 0; - - // Force exit once the overall budget is spent. Exit 0 is deliberate: - // a force-timeout still happens on a routine deploy (deliberate - // shutdown, not a crash), and orchestrators record nonzero exits on - // deploys as crashes. The error log below is the stuck-shutdown - // signal instead of the exit code. + // Exit 0 on force-timeout: a stuck deploy shutdown is not a crash, and + // orchestrators read nonzero deploy exits as one. The error log is the + // signal instead. Lives here, not in runPhases, because close() must not + // inherit it. const forceExitTimer = setTimeout(() => { logger.error( "Graceful shutdown did NOT complete within the %dms budget (phase in flight: %s); force-exiting with code 0.", @@ -110,13 +116,74 @@ export class LifecycleManager { ); process.exit(0); }, LifecycleManager.SHUTDOWN_TIMEOUT_MS); - // unref so this backstop timer never by itself keeps the process alive. - // Any real pending teardown (OTEL export timer, DB pool sockets, the - // still-open HTTP listener) is a ref'd handle that holds the loop open - // until this fires; if nothing is ref'd, there is nothing left to tear - // down and exiting early is correct. + // unref'd so the backstop alone never holds the process open; real pending + // teardown is ref'd and keeps the loop alive until this fires. forceExitTimer.unref(); + const exitCode = await this.runOnce(); + + clearTimeout(forceExitTimer); + process.exit(exitCode); + } + + /** + * Release everything AppKit acquired **without terminating the process** — + * same phases and per-phase budgets as {@link shutdown}, no `process.exit`. + * + * Handlers are detached before the first `await`, so the SIGTERM-mid-close + * window is near zero; if one does land there the signal wins and this promise + * never settles. Never throws — a hung phase is logged and `close()` resolves + * once its budget is spent, so an `afterEach` cannot hang. + */ + async close(options: { timeoutMs?: number } = {}): Promise { + // Before the first await, so a later signal finds no AppKit listener. + this.removeSignalHandlers(); + + const timeoutMs = options.timeoutMs ?? LifecycleManager.CLOSE_TIMEOUT_MS; + + try { + await this.raceWithTimeout(this.runOnce(), timeoutMs, "close"); + } catch (err) { + logger.error( + "close() did not complete within the %dms budget (phase in flight: %s): %O", + timeoutMs, + this.shutdownPhase, + err, + ); + } + + // close() only — on the signal path the process is dying and this is pure cost. + resetCoreSingletons(); + } + + /** No `await` between read and assign — that gap is the re-entrancy window. */ + private runOnce(): Promise { + this.teardown ??= this.runPhases(); + return this.teardown; + } + + /** Run the phases and report an exit code; no process-termination concerns. */ + private async runPhases(): Promise { + logger.info("Starting graceful shutdown..."); + + // Captured before the first await and never re-read: close() may give up + // waiting and reset the singletons while these phases still run, so phase 5 + // would otherwise skip this app's pool or tear down the *next* app's. + let capturedCache: CacheManager | undefined; + try { + capturedCache = CacheManager.getInstanceSync(); + } catch { + // Never initialized — nothing to close in phase 5. + } + let capturedTelemetry: TelemetryManager | undefined; + try { + capturedTelemetry = TelemetryManager.getInstance(); + } catch { + // Unavailable or mocked away — nothing to flush. + } + + let exitCode = 0; + try { const plugins = Array.from(this.context.getPlugins().values()); @@ -174,7 +241,10 @@ export class LifecycleManager { // cache), so they run concurrently — each bounded so a stuck pool // drain or stalled OTLP export cannot eat the remaining budget. this.shutdownPhase = "cache storage close + telemetry flush"; - await Promise.all([this.closeCacheStorage(), this.flushTelemetry()]); + await Promise.all([ + this.closeCacheStorage(capturedCache), + this.flushTelemetry(capturedTelemetry), + ]); logger.info("Graceful shutdown complete"); } catch (err) { @@ -184,16 +254,14 @@ export class LifecycleManager { exitCode = 1; } - clearTimeout(forceExitTimer); - process.exit(exitCode); + return exitCode; } - /** Close the cache storage, bounded and error-isolated. */ - private async closeCacheStorage(): Promise { - let cache: CacheManager; - try { - cache = CacheManager.getInstanceSync(); - } catch { + /** Bounded and error-isolated. Takes the manager — see the capture in {@link runPhases}. */ + private async closeCacheStorage( + cache: CacheManager | undefined, + ): Promise { + if (!cache) { // Cache was never initialized — nothing to close. return; } @@ -208,11 +276,14 @@ export class LifecycleManager { } } - /** Flush and shut down the telemetry SDK, bounded and error-isolated. */ - private async flushTelemetry(): Promise { + /** Bounded and error-isolated. Takes the manager — see {@link closeCacheStorage}. */ + private async flushTelemetry( + telemetry: TelemetryManager | undefined, + ): Promise { + if (!telemetry) return; try { await this.raceWithTimeout( - TelemetryManager.getInstance().shutdown(), + telemetry.shutdown(), LifecycleManager.PHASE_SHUTDOWN_TIMEOUT_MS, "telemetry flush", ); diff --git a/packages/appkit/src/core/plugin-context.ts b/packages/appkit/src/core/plugin-context.ts index 2a89d33c3..4f86a5189 100644 --- a/packages/appkit/src/core/plugin-context.ts +++ b/packages/appkit/src/core/plugin-context.ts @@ -2,7 +2,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"); @@ -63,7 +67,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. diff --git a/packages/appkit/src/core/reset-singletons.ts b/packages/appkit/src/core/reset-singletons.ts new file mode 100644 index 000000000..6b98d3e91 --- /dev/null +++ b/packages/appkit/src/core/reset-singletons.ts @@ -0,0 +1,33 @@ +import { CacheManager } from "../cache"; +import { ServiceContext } from "../context"; +import { TelemetryReporter } from "../internal-telemetry"; +import { createLogger } from "../logging/logger"; +import { TelemetryManager } from "../telemetry"; + +const logger = createLogger("lifecycle"); + +/** + * Drop the four singletons `AppKit._createApp` initializes. + * + * Pointer drops, not teardown — callers close first, or the old app's storage and + * exporters leak. Core initializes all four, so core drops all four; a host that + * closes then calls `ServiceContext.get()` will get an `InitializationError`. + * + * @internal + */ +export function resetCoreSingletons(): void { + const resets: [string, () => void][] = [ + ["ServiceContext", () => ServiceContext.reset()], + ["CacheManager", () => CacheManager.reset()], + ["TelemetryReporter", () => TelemetryReporter._reset()], + ["TelemetryManager", () => TelemetryManager.reset()], + ]; + + for (const [name, reset] of resets) { + try { + reset(); + } catch (err) { + logger.error("Error resetting %s: %O", name, err); + } + } +} diff --git a/packages/appkit/src/core/tests/app-close.integration.test.ts b/packages/appkit/src/core/tests/app-close.integration.test.ts new file mode 100644 index 000000000..a298c4ec1 --- /dev/null +++ b/packages/appkit/src/core/tests/app-close.integration.test.ts @@ -0,0 +1,239 @@ +import { + getListeningPort, + mockServiceContext, + setupDatabricksEnv, +} from "@databricks/appkit/testing"; +import type { AppHandle, PluginManifest, PluginMap } from "shared"; +import { afterEach, beforeEach, describe, expect, test } from "vitest"; + +import { CacheManager } from "../../cache"; +import { ServiceContext } from "../../context/service-context"; +import { ConfigurationError } from "../../errors"; +import { Plugin, toPlugin } from "../../plugin"; +import { server as serverPlugin } from "../../plugins/server"; +import { createApp } from "../appkit"; + +/** + * Deliberately unmocked — the claim is that `close()` releases *real* resources, + * so a mocked lifecycle would assert nothing. `port: 0` keeps it parallel-safe. + */ + +/** Minimal plugin with a route, so there is something real to serve. */ +class ProbePlugin extends Plugin { + static manifest: PluginManifest = { + name: "probe", + displayName: "Probe", + version: "0.0.0", + description: "close() integration probe", + resources: { required: [] }, + } as unknown as PluginManifest; + + /** Set when the lifecycle actually ran this plugin's teardown. */ + shutdownCalls = 0; + + async shutdown(): Promise { + this.shutdownCalls += 1; + } + + exports() { + return { shutdownCalls: () => this.shutdownCalls }; + } +} +const probe = toPlugin(ProbePlugin); + +/** A plugin whose manifest name collides with the handle's own method. */ +class ClosePlugin extends Plugin { + static manifest: PluginManifest = { + name: "close", + displayName: "Close", + version: "0.0.0", + description: "reserved-name probe", + resources: { required: [] }, + } as unknown as PluginManifest; +} +const closeNamed = toPlugin(ClosePlugin); + +describe("app handle close()", () => { + let serviceContextMock: ReturnType; + + beforeEach(() => { + setupDatabricksEnv(); + ServiceContext.reset(); + serviceContextMock = mockServiceContext(); + }); + + afterEach(() => { + serviceContextMock?.restore(); + }); + + test("releases the bound socket and runs plugin teardown", async () => { + const termBaseline = process.listenerCount("SIGTERM"); + + const app = await createApp({ + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + + // AppKit installed its handlers, so the count went up. + expect(process.listenerCount("SIGTERM")).toBe(termBaseline + 1); + + const port = await getListeningPort(app.server.getServer()); + const baseUrl = `http://127.0.0.1:${port}`; + await expect( + fetch(`${baseUrl}/health`).then((r) => r.status), + ).resolves.toBe(200); + + await app.close(); + + // The plugin's own teardown hook ran... + expect(app.probe.shutdownCalls()).toBe(1); + // ...the listener is gone... + await expect(fetch(`${baseUrl}/health`)).rejects.toThrow(); + // ...and the signal handlers came back off, which is what keeps repeated + // boots from tripping MaxListenersExceededWarning. + expect(process.listenerCount("SIGTERM")).toBe(termBaseline); + }); + + test("is idempotent at the app level", async () => { + const app = await createApp({ + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + await getListeningPort(app.server.getServer()); + + await app.close(); + await expect(app.close()).resolves.toBeUndefined(); + + // The memo means the phases ran once, not twice. + expect(app.probe.shutdownCalls()).toBe(1); + }); + + test("plugin exports stay reachable by name alongside close()", async () => { + const app = await createApp({ + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + await getListeningPort(app.server.getServer()); + + try { + // Adding `close` to the handle must not shadow or be shadowed by the + // plugin accessors installed with defineProperty. + expect(typeof app.close).toBe("function"); + expect(typeof app.server.getServer).toBe("function"); + expect(typeof app.probe.shutdownCalls).toBe("function"); + expect(typeof app[Symbol.asyncDispose]).toBe("function"); + } finally { + await app.close(); + } + }); + + test("a server-less app still closes cleanly", async () => { + // No server plugin at all: nothing bound a socket, but plugin hooks and the + // telemetry flush still have to run, and close() must not hang. + const app = await createApp({ plugins: [probe()] }); + + await expect(app.close()).resolves.toBeUndefined(); + expect(app.probe.shutdownCalls()).toBe(1); + }); + + test("await using releases the app at scope exit", async () => { + let captured: number | undefined; + let probeHandle: { shutdownCalls: () => number } | undefined; + + { + await using app = await createApp({ + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + captured = await getListeningPort(app.server.getServer()); + probeHandle = app.probe; + await expect( + fetch(`http://127.0.0.1:${captured}/health`).then((r) => r.status), + ).resolves.toBe(200); + } + + // Scope exited, so asyncDispose ran the same teardown. + expect(probeHandle?.shutdownCalls()).toBe(1); + await expect( + fetch(`http://127.0.0.1:${captured}/health`), + ).rejects.toThrow(); + }); + + test("a plugin named close is rejected instead of silently shadowing", async () => { + // An own property wins over a prototype method, so without this guard the + // plugin would quietly replace teardown rather than fail. + await expect(createApp({ plugins: [closeNamed()] })).rejects.toThrow( + ConfigurationError, + ); + await expect(createApp({ plugins: [closeNamed()] })).rejects.toThrow( + /"close" is reserved|Plugin name "close" is reserved/, + ); + }); + test("boot, close, boot again in one file — the second app gets a live cache", async () => { + // The stated driver for the whole close() effort: two real boots, two real + // sockets, one process. + // + // Note what this does *not* prove. The cache here is InMemoryStorage, whose + // close() merely clears a Map and stays usable, so this passes with or + // without the singleton resets. The reset's necessity is proven in + // cache/tests/cache-manager-reset.test.ts against storage whose close() is + // terminal, the way PersistentStorage's pool.end() is. + const termBaseline = process.listenerCount("SIGTERM"); + + const first = await createApp({ + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + const firstPort = await getListeningPort(first.server.getServer()); + await expect( + fetch(`http://127.0.0.1:${firstPort}/health`).then((r) => r.status), + ).resolves.toBe(200); + await first.close(); + + // ServiceContext was reset by close(), so the mock has to be reinstalled — + // exactly what createTestApp will do for the caller. + serviceContextMock.restore(); + serviceContextMock = mockServiceContext(); + + const second = await createApp({ + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + const secondPort = await getListeningPort(second.server.getServer()); + + expect(secondPort).not.toBe(firstPort); + await expect( + fetch(`http://127.0.0.1:${secondPort}/health`).then((r) => r.status), + ).resolves.toBe(200); + + // The second boot's cache round-trips a write. + const cache = CacheManager.getInstanceSync(); + const key = cache.generateKey(["second-boot"], "test-user"); + await cache.set(key, { alive: true }); + await expect(cache.get(key)).resolves.toEqual({ alive: true }); + + await second.close(); + + await expect( + fetch(`http://127.0.0.1:${secondPort}/health`), + ).rejects.toThrow(); + // Two boots and two closes leave no listener residue. + expect(process.listenerCount("SIGTERM")).toBe(termBaseline); + }); + /** + * Enforced by `tsc --noEmit`, not at runtime: the widening to `AppHandle` is + * only source-compatible if it stays assignable to `PluginMap`, and a + * regression there would break existing callers without failing any assertion. + */ + describe("createApp return-type widening is source-compatible", () => { + test("an AppHandle still satisfies a PluginMap annotation", async () => { + const app = await createApp({ plugins: [probe()] }); + try { + // The pre-widening annotation, unchanged. + const asPluginMap: PluginMap<[ReturnType]> = app; + expect(typeof asPluginMap.probe.shutdownCalls).toBe("function"); + + // And the added members are visible on the widened type. + const asHandle: AppHandle<[ReturnType]> = app; + expect(typeof asHandle.close).toBe("function"); + expect(typeof asHandle[Symbol.asyncDispose]).toBe("function"); + } finally { + await app.close(); + } + }); + }); +}); diff --git a/packages/appkit/src/core/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts index 121e7eb5c..6d2747ca3 100644 --- a/packages/appkit/src/core/tests/lifecycle-manager.test.ts +++ b/packages/appkit/src/core/tests/lifecycle-manager.test.ts @@ -381,4 +381,288 @@ describe("LifecycleManager", () => { onceSpy.mockRestore(); }); }); + describe("close (the programmatic path)", () => { + test("runs the full teardown sequence without exiting the process", async () => { + const stop = vi.fn(); + vi.mocked(TelemetryReporter.getInstance).mockReturnValue({ + stop, + } as never); + const cacheClose = vi.fn().mockResolvedValue(undefined); + vi.mocked(CacheManager.getInstanceSync).mockReturnValue({ + close: cacheClose, + } as never); + const telemetryShutdown = vi.fn().mockResolvedValue(undefined); + vi.mocked(TelemetryManager.getInstance).mockReturnValue({ + shutdown: telemetryShutdown, + } as never); + + const abortActiveOperations = vi.fn(); + const shutdown = vi.fn().mockResolvedValue(undefined); + const ctx = contextWithPlugins({ + alpha: { name: "alpha", abortActiveOperations, shutdown } as never, + }); + const emit = vi.spyOn(ctx, "emitLifecycle"); + const manager = new LifecycleManager(ctx); + + await manager.close(); + + expect(stop).toHaveBeenCalledTimes(1); + expect(abortActiveOperations).toHaveBeenCalledTimes(1); + expect(shutdown).toHaveBeenCalledTimes(1); + expect(emit).toHaveBeenCalledWith("shutdown"); + expect(cacheClose).toHaveBeenCalledTimes(1); + expect(telemetryShutdown).toHaveBeenCalledTimes(1); + + // The whole point of the split. + expect(exitSpy).not.toHaveBeenCalled(); + }); + + test("is idempotent: teardown runs once and the second call awaits it", async () => { + let releaseShutdown: (() => void) | undefined; + // Set only once the plugin hook has actually finished. Asserting against + // this flag (rather than counting microtask ticks) is what makes the test + // sensitive to a guard that returns early while teardown is in flight. + let teardownFinished = false; + const shutdown = vi.fn( + () => + new Promise((resolve) => { + releaseShutdown = () => { + teardownFinished = true; + resolve(); + }; + }), + ); + const ctx = contextWithPlugins({ + alpha: { name: "alpha", shutdown } as never, + }); + const manager = new LifecycleManager(ctx); + + const observed: string[] = []; + const first = manager + .close() + .then(() => observed.push(`first:${teardownFinished}`)); + const second = manager + .close() + .then(() => observed.push(`second:${teardownFinished}`)); + + // A full macrotask turn, so a guard that resolves the second caller + // early has every chance to settle before the assertion below. + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(observed).toEqual([]); + + releaseShutdown?.(); + await Promise.all([first, second]); + + // Both callers must observe a *completed* teardown. The old boolean + // guard resolved the second caller with teardown still running. + expect(observed).toEqual( + expect.arrayContaining(["first:true", "second:true"]), + ); + expect(shutdown).toHaveBeenCalledTimes(1); + expect(exitSpy).not.toHaveBeenCalled(); + }); + + test("a signal arriving after close() joins the same teardown, not a second one", async () => { + const shutdown = vi.fn().mockResolvedValue(undefined); + const ctx = contextWithPlugins({ + alpha: { name: "alpha", shutdown } as never, + }); + const manager = new LifecycleManager(ctx); + manager.installSignalHandlers(); + + await manager.close(); + // The signal path after a close: teardown is memoized, so the phases do + // not run twice even though shutdown() is still callable. + await manager.shutdown(); + + expect(shutdown).toHaveBeenCalledTimes(1); + }); + + test("close() after a signal-initiated teardown awaits the in-flight one", async () => { + let releaseShutdown: (() => void) | undefined; + // Sentinel rather than a tick count: `close()` reaches the memo through + // raceWithTimeout, so "how many microtasks until it would have settled" is + // not a property the test can rely on. + let teardownFinished = false; + const shutdown = vi.fn( + () => + new Promise((resolve) => { + releaseShutdown = () => { + teardownFinished = true; + resolve(); + }; + }), + ); + const ctx = contextWithPlugins({ + alpha: { name: "alpha", shutdown } as never, + }); + const manager = new LifecycleManager(ctx); + + const signalPath = manager.shutdown(); + await Promise.resolve(); + + let closeSawFinishedTeardown: boolean | undefined; + const closePath = manager.close().then(() => { + closeSawFinishedTeardown = teardownFinished; + }); + + // A full macrotask turn, so a close() that resolved early would have. + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(closeSawFinishedTeardown).toBeUndefined(); + + releaseShutdown?.(); + await Promise.all([signalPath, closePath]); + + expect(shutdown).toHaveBeenCalledTimes(1); + // It joined the in-flight teardown rather than resolving alongside it. + expect(closeSawFinishedTeardown).toBe(true); + // The signal wanted the process dead, and still gets it. + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a rejecting plugin shutdown() is isolated and close() still resolves", async () => { + const ctx = contextWithPlugins({ + bad: { + name: "bad", + shutdown: vi.fn().mockRejectedValue(new Error("teardown blew up")), + } as never, + good: { + name: "good", + shutdown: vi.fn().mockResolvedValue(undefined), + } as never, + }); + const manager = new LifecycleManager(ctx); + + await expect(manager.close()).resolves.toBeUndefined(); + expect(mockLoggerError).toHaveBeenCalled(); + expect(exitSpy).not.toHaveBeenCalled(); + }); + + test("a hung teardown is bounded by close()'s budget, logged, and never exits", async () => { + vi.useFakeTimers(); + const ctx = contextWithPlugins({ + stuck: { + name: "stuck", + shutdown: vi.fn(() => new Promise(() => {})), + } as never, + }); + const manager = new LifecycleManager(ctx); + + const closing = manager.close({ timeoutMs: 50 }); + await vi.advanceTimersByTimeAsync(60); + await expect(closing).resolves.toBeUndefined(); + + // The error names the phase that was in flight, which is the whole + // reason the phase tracker is retained. + const logged = mockLoggerError.mock.calls + .map((c) => String(c[0])) + .join("\n"); + expect(logged).toContain("close() did not complete"); + const phases = mockLoggerError.mock.calls.flat().map(String).join(" "); + expect(phases).toContain("plugin shutdown() hooks"); + + // A hung teardown must not kill the process on the programmatic path. + expect(exitSpy).not.toHaveBeenCalled(); + }); + }); + + describe("signal-handler ownership", () => { + test("close() removes only this manager's listeners", async () => { + const foreign = vi.fn(); + process.on("SIGTERM", foreign); + const baseline = process.listenerCount("SIGTERM"); + + const a = new LifecycleManager(contextWithPlugins({})); + const b = new LifecycleManager(contextWithPlugins({})); + a.installSignalHandlers(); + b.installSignalHandlers(); + expect(process.listenerCount("SIGTERM")).toBe(baseline + 2); + + await a.close(); + + // b's pair survives, and so does the unrelated host listener. + expect(process.listenerCount("SIGTERM")).toBe(baseline + 1); + + await b.close(); + expect(process.listenerCount("SIGTERM")).toBe(baseline); + expect(process.listeners("SIGTERM")).toContain(foreign); + + process.removeListener("SIGTERM", foreign); + }); + + test("listener counts return to the pre-install baseline", async () => { + const termBaseline = process.listenerCount("SIGTERM"); + const intBaseline = process.listenerCount("SIGINT"); + + const manager = new LifecycleManager(contextWithPlugins({})); + manager.installSignalHandlers(); + await manager.close(); + + // This is what keeps repeated boots in one test file from tripping + // MaxListenersExceededWarning at ~6 un-closed apps. + expect(process.listenerCount("SIGTERM")).toBe(termBaseline); + expect(process.listenerCount("SIGINT")).toBe(intBaseline); + }); + + test("removeSignalHandlers is safe when none were installed", () => { + const manager = new LifecycleManager(contextWithPlugins({})); + expect(() => manager.removeSignalHandlers()).not.toThrow(); + }); + }); + describe("a teardown that outlives close()'s budget", () => { + test("phase 5 still closes the app's own cache and telemetry, not the next app's", async () => { + vi.useFakeTimers(); + + // The app being torn down owns these. + const ownCacheClose = vi.fn().mockResolvedValue(undefined); + const ownTelemetryShutdown = vi.fn().mockResolvedValue(undefined); + vi.mocked(CacheManager.getInstanceSync).mockReturnValue({ + close: ownCacheClose, + } as never); + vi.mocked(TelemetryManager.getInstance).mockReturnValue({ + shutdown: ownTelemetryShutdown, + } as never); + + // A plugin hook slower than close()'s budget but inside its own per-plugin + // budget — the files plugin's 10s drain reaches exactly this state. + let releaseHook: (() => void) | undefined; + const ctx = contextWithPlugins({ + slow: { + name: "slow", + shutdown: vi.fn( + () => + new Promise((resolve) => { + releaseHook = resolve; + }), + ), + } as never, + }); + const manager = new LifecycleManager(ctx); + + const closing = manager.close({ timeoutMs: 50 }); + await vi.advanceTimersByTimeAsync(60); + await expect(closing).resolves.toBeUndefined(); + + // close() has given up waiting and already dropped the singletons, so the + // static slots now answer with a *different* app's resources. + const nextCacheClose = vi.fn().mockResolvedValue(undefined); + const nextTelemetryShutdown = vi.fn().mockResolvedValue(undefined); + vi.mocked(CacheManager.getInstanceSync).mockReturnValue({ + close: nextCacheClose, + } as never); + vi.mocked(TelemetryManager.getInstance).mockReturnValue({ + shutdown: nextTelemetryShutdown, + } as never); + + // Now let the orphaned teardown finish and reach phase 5. + releaseHook?.(); + await vi.advanceTimersByTimeAsync(10); + + // It must act on what it captured at the start, never on the current slots. + expect(ownCacheClose).toHaveBeenCalledTimes(1); + expect(ownTelemetryShutdown).toHaveBeenCalledTimes(1); + expect(nextCacheClose).not.toHaveBeenCalled(); + expect(nextTelemetryShutdown).not.toHaveBeenCalled(); + }); + }); }); diff --git a/packages/appkit/src/index.ts b/packages/appkit/src/index.ts index eac0b27b9..40fa658b8 100644 --- a/packages/appkit/src/index.ts +++ b/packages/appkit/src/index.ts @@ -7,6 +7,7 @@ // Types from shared export type { + AppHandle, BasePluginConfig, CacheConfig, IAppRouter, 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 55a03934b..70a195f9c 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 @@ -2,6 +2,7 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; +import { createTestPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; /** @@ -36,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; } @@ -284,15 +291,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"); - (plugin as any).context = { executeTool }; - - const toolIndex = new Map([ + const toolkitToolIndex = () => + new Map([ [ "analytics.query", { @@ -308,19 +308,74 @@ 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 = createTestPluginContext({ analytics: { query: "rows" } }); + const executeToolSpy = vi.spyOn(mock.ctx, "executeTool"); + await mock.attach(plugin); + + 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 = createTestPluginContext({ + analytics: { + query: (_args, signal) => + new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => + reject(new Error("aborted by toolkit timeout")), + ); + }), + }, + }); + await mock.attach(plugin); + + 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/plugins/agents/tests/route-handler-errors.test.ts b/packages/appkit/src/plugins/agents/tests/route-handler-errors.test.ts index 1deacd5c5..4aa069a61 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 @@ -2,28 +2,9 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; +import { createTestPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; -// Partial-mock the tracing module: traceAgent/traceTool still run their -// callbacks, but the trace id is deterministic and run-linking is a spy. -const linkTraceToRun = vi.hoisted(() => vi.fn()); -let mockTraceId: string | undefined; -vi.mock("../mlflow", () => ({ - initAgentTracing: vi.fn(async () => {}), - traceAgent: ( - _name: string, - _inputs: unknown, - fn: (span: { setOutputs: () => void }) => Promise, - ) => fn({ setOutputs: () => {} }), - traceTool: ( - _name: string, - _inputs: unknown, - fn: (span: { setOutputs: () => void }) => Promise, - ) => fn({ setOutputs: () => {} }), - currentTraceId: () => mockTraceId, - linkTraceToRun, -})); - /** * Surface-level guarantees on the agents plugin's HTTP route handlers when * downstream dependencies fail. Prior to PR #305 review finding #1+#2, @@ -40,8 +21,6 @@ vi.mock("../mlflow", () => ({ */ beforeEach(() => { - linkTraceToRun.mockClear(); - mockTraceId = undefined; (CacheManager as any).instance = { get: vi.fn(), set: vi.fn(), @@ -82,12 +61,12 @@ function mockRes() { }; } -function seedPlugin(adapter: unknown = { async *run() {} }): AgentsPlugin { +function seedPlugin(): AgentsPlugin { const plugin = new AgentsPlugin({ dir: false }); (plugin as any).agents.set("default", { name: "default", instructions: "hi", - adapter, + adapter: { async *run() {} }, toolIndex: new Map(), }); (plugin as any).defaultAgentName = "default"; @@ -395,94 +374,25 @@ describe("POST /invocations & /responses — successful invoke", () => { text: "hello world", }); }); - - function seedEchoPlugin(): AgentsPlugin { - const plugin = seedPlugin({ - async *run() { - yield { type: "message_delta", content: "ok" }; - }, - }); - (plugin as any).threadStore = { - create: vi.fn().mockResolvedValue({ id: "t-new", messages: [] }), - addMessage: vi.fn(), - delete: vi.fn(), - }; - return plugin; - } - - async function invoke( - plugin: AgentsPlugin, - body: unknown, - ): Promise> { - const { res, json } = mockRes(); - await ( - plugin as unknown as { - _handleInvoke: ( - r: express.Request, - w: express.Response, - ) => Promise; - } - )._handleInvoke(mockReq(body), res); - return json.mock.calls[0]?.[0] as Record; - } - - test("links the trace to the run and echoes mlflow_trace_id when tracing is on", async () => { - mockTraceId = "tr-abc123"; - const plugin = seedEchoPlugin(); - - const payload = await invoke(plugin, { - input: "hi", - mlflowRunId: "run-99", - }); - - expect(linkTraceToRun).toHaveBeenCalledWith("run-99"); - expect(payload.mlflow_trace_id).toBe("tr-abc123"); - }); - - test("omits mlflow_trace_id and does not link when tracing is off", async () => { - mockTraceId = undefined; // currentTraceId() no-ops when disabled - const plugin = seedEchoPlugin(); - - const payload = await invoke(plugin, { input: "hi" }); - - expect(linkTraceToRun).not.toHaveBeenCalled(); - expect(payload).not.toHaveProperty("mlflow_trace_id"); - }); - - test("does not link when no run id is supplied even if tracing is on", async () => { - mockTraceId = "tr-standalone"; - const plugin = seedEchoPlugin(); - - const payload = await invoke(plugin, { input: "hi" }); - - expect(linkTraceToRun).not.toHaveBeenCalled(); - // Trace still exists and its id is surfaced — just not linked to a run. - expect(payload.mlflow_trace_id).toBe("tr-standalone"); - }); }); 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(); - (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 = createTestPluginContext(); + (plugin as any).context = mock.ctx; (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]); }); }); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts index e099c8350..c1ffb0806 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts @@ -1,13 +1,11 @@ -import type { Server } from "node:http"; - import { - createConfigurableMockWorkspaceClient, createFailedSQLResponse, createSuccessfulSQLResponse, - mockServiceContext, + createTestApp, + getMockFn, parseSSEResponse, - setupDatabricksEnv, -} from "@tools/test-helpers"; + type TestApp, +} from "@databricks/appkit/testing"; import { sql } from "shared"; import { afterAll, @@ -20,85 +18,38 @@ import { } from "vitest"; import { AppManager } from "../../../app"; -import { ServiceContext } from "../../../context/service-context"; -import { createApp } from "../../../core"; -import { server as serverPlugin } from "../../server"; import { analytics } from "../index"; const getAppQuerySpy = vi.spyOn(AppManager.prototype, "getAppQuery"); -/** - * Wait for the supplied server to finish binding, then return the OS-assigned - * port. Required when the test passes `port: 0` to `serverPlugin` — - * `app.server.start()` returns as soon as `listen()` is invoked but before the - * bind completes, so `server.address()` returns `null` until the `listening` - * event fires. - */ -async function getListeningPort(server: Server): Promise { - const addr = server.address(); - if (addr && typeof addr === "object" && typeof addr.port === "number") { - return addr.port; - } - await new Promise((resolve, reject) => { - server.once("listening", () => resolve()); - server.once("error", (err) => reject(err)); - }); - const ready = server.address(); - if (!ready || typeof ready !== "object") { - throw new Error("Server is listening but address() returned null"); - } - return ready.port; -} - describe("Analytics Plugin Integration", () => { - let server: Server; - let baseUrl: string; - let serviceContextMock: Awaited>; - let mockClient: ReturnType; + let app: TestApp<[ReturnType]>; + /** The SQL mock the analytics route drives, via the harness's client. */ + let executeStatement: ReturnType; + let getStatement: ReturnType; beforeAll(async () => { - setupDatabricksEnv(); - ServiceContext.reset(); - - mockClient = createConfigurableMockWorkspaceClient(); - serviceContextMock = await mockServiceContext({ - serviceDatabricksClient: mockClient.client, - }); - - const app = await createApp({ - plugins: [ - // port: 0 → OS assigns an ephemeral port. Avoids EADDRINUSE / cross-test - // route bleed when another integration test (e.g. server.integration) - // holds a fixed port concurrently in the shared vitest worker pool. - serverPlugin({ - port: 0, - host: "127.0.0.1", - }), - analytics({}), - ], - }); - - server = app.server.getServer(); - const port = await getListeningPort(server); - baseUrl = `http://127.0.0.1:${port}`; + // The harness owns the env setup, the singleton resets, the mock client, the + // server plugin on an ephemeral port, and the teardown. What used to be ~45 + // lines of setup plus a local getListeningPort helper is this call. + app = await createTestApp({ plugins: [analytics({})] }); + executeStatement = getMockFn( + app.client, + "statementExecution.executeStatement", + ); + getStatement = getMockFn(app.client, "statementExecution.getStatement"); }); afterAll(async () => { getAppQuerySpy?.mockRestore(); - serviceContextMock?.restore(); - if (server) { - await new Promise((resolve, reject) => { - server.close((err) => { - if (err) reject(err); - else resolve(); - }); - }); - } + await app?.close(); }); beforeEach(() => { - mockClient.mocks.executeStatement.mockReset(); - mockClient.mocks.getStatement.mockReset(); + // Reset drops the built-in canned SUCCEEDED default too, matching the + // "script it yourself" semantics this suite relied on before. + executeStatement.mockReset(); + getStatement.mockReset(); getAppQuerySpy.mockReset(); }); @@ -119,18 +70,13 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockResolvedValueOnce( + executeStatement.mockResolvedValueOnce( createSuccessfulSQLResponse(mockData, mockColumns), ); - const response = await fetch( - `${baseUrl}/api/analytics/query/test_query`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response = await app.post("/api/analytics/query/test_query", { + body: { parameters: {} }, + }); expect(response.status).toBe(200); expect(response.headers.get("Content-Type")).toBe( @@ -144,8 +90,8 @@ describe("Analytics Plugin Integration", () => { { name: "Bob", age: "25" }, ]); - expect(mockClient.mocks.executeStatement).toHaveBeenCalledTimes(1); - expect(mockClient.mocks.executeStatement).toHaveBeenCalledWith( + expect(executeStatement).toHaveBeenCalledTimes(1); + expect(executeStatement).toHaveBeenCalledWith( expect.objectContaining({ statement: testQuery, warehouse_id: "test-warehouse-id", @@ -162,26 +108,17 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockResolvedValueOnce( + executeStatement.mockResolvedValueOnce( createSuccessfulSQLResponse([["Alice"]], [{ name: "name" }]), ); - const response = await fetch( - `${baseUrl}/api/analytics/query/user_query`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - parameters: { - user_id: sql.string("123"), - }, - }), - }, - ); + const response = await app.post("/api/analytics/query/user_query", { + body: { parameters: { user_id: sql.string("123") } }, + }); expect(response.status).toBe(200); - const callArgs = mockClient.mocks.executeStatement.mock.calls[0][0]; + const callArgs = executeStatement.mock.calls[0][0]; expect(callArgs.parameters).toEqual( expect.arrayContaining([ expect.objectContaining({ @@ -198,20 +135,15 @@ describe("Analytics Plugin Integration", () => { test("should return 404 when query does not exist", async () => { getAppQuerySpy.mockResolvedValueOnce(null); - const response = await fetch( - `${baseUrl}/api/analytics/query/nonexistent`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response = await app.post("/api/analytics/query/nonexistent", { + body: { parameters: {} }, + }); expect(response.status).toBe(404); const data = await response.json(); expect(data).toEqual({ error: "Query not found" }); - expect(mockClient.mocks.executeStatement).not.toHaveBeenCalled(); + expect(executeStatement).not.toHaveBeenCalled(); }); }); @@ -222,14 +154,12 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockResolvedValue( + executeStatement.mockResolvedValue( createFailedSQLResponse("Table not found"), ); - const response = await fetch(`${baseUrl}/api/analytics/query/broken`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), + const response = await app.post("/api/analytics/query/broken", { + body: { parameters: {} }, }); expect(response.status).toBe(200); @@ -243,14 +173,10 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockRejectedValue( - new Error("Network error"), - ); + executeStatement.mockRejectedValue(new Error("Network error")); - const response = await fetch(`${baseUrl}/api/analytics/query/error`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), + const response = await app.post("/api/analytics/query/error", { + body: { parameters: {} }, }); expect(response.status).toBe(200); @@ -268,33 +194,23 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockResolvedValue( + executeStatement.mockResolvedValue( createSuccessfulSQLResponse([["cached_value"]], [{ name: "value" }]), ); - const response1 = await fetch( - `${baseUrl}/api/analytics/query/cache_test`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response1 = await app.post("/api/analytics/query/cache_test", { + body: { parameters: {} }, + }); const data1 = await parseSSEResponse(response1); - const response2 = await fetch( - `${baseUrl}/api/analytics/query/cache_test`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response2 = await app.post("/api/analytics/query/cache_test", { + body: { parameters: {} }, + }); const data2 = await parseSSEResponse(response2); expect(data1.data).toEqual([{ value: "cached_value" }]); expect(data2.data).toEqual([{ value: "cached_value" }]); - expect(mockClient.mocks.executeStatement).toHaveBeenCalledTimes(1); + expect(executeStatement).toHaveBeenCalledTimes(1); }); }); }); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index dd14b1d3d..5101f9424 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"; @@ -1648,7 +1650,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(); @@ -1666,27 +1668,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. @@ -1703,11 +1708,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); @@ -1826,3 +1829,64 @@ 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).toHaveLength(1); + 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/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index a49412ac1..bf721feee 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -744,7 +744,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({ @@ -765,23 +765,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/plugins/files/tests/plugin.integration.test.ts b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts index 3be68b315..0134c09e6 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts @@ -1,6 +1,10 @@ import http, { type Server } from "node:http"; -import { mockServiceContext, setupDatabricksEnv } from "@tools/test-helpers"; +import { + getListeningPort, + mockServiceContext, + setupDatabricksEnv, +} from "@databricks/appkit/testing"; import { afterAll, beforeAll, @@ -67,29 +71,6 @@ const MOCK_AUTH_HEADERS = { /** Volume key used in all integration tests. */ const VOL = "files"; -/** - * Wait for the supplied server to finish binding, then return the - * OS-assigned port. Required when tests pass `port: 0` to `serverPlugin` - * — `appkit.server.start()` returns as soon as `listen()` is invoked but - * before the bind completes, so `server.address()` returns `null` until - * the `listening` event fires. - */ -async function getListeningPort(server: Server): Promise { - const addr = server.address(); - if (addr && typeof addr === "object" && typeof addr.port === "number") { - return addr.port; - } - await new Promise((resolve, reject) => { - server.once("listening", () => resolve()); - server.once("error", (err) => reject(err)); - }); - const ready = server.address(); - if (!ready || typeof ready !== "object") { - throw new Error("Server is listening but address() returned null"); - } - return ready.port; -} - describe("Files Plugin Integration", () => { let server: Server; let baseUrl: string; diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 7e5afa0fb..2d867d335 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"; @@ -334,23 +335,23 @@ 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 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(); }); diff --git a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts index 783debc8a..c706d648b 100644 --- a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts @@ -12,38 +12,47 @@ import { import { mapParams } from "../params"; import { JobsPlugin, jobs } from "../plugin"; -const { mockClient, mockCacheInstance } = vi.hoisted(() => { - const mockJobsApi = { - runNow: vi.fn(), - submit: vi.fn(), - getRun: vi.fn(), - getRunOutput: vi.fn(), - cancelRun: vi.fn(), - listRuns: vi.fn(), - get: vi.fn(), - }; - - const mockClient = { - jobs: mockJobsApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; +const { mockClient, jobsApi, mockCacheInstance } = await vi.hoisted( + async () => { + // The testing kit's fake, not a hand-rolled literal: the seven jobs methods, + // `config.host` as a real string, and `config.authenticate` all come for free, + // and any *other* service this plugin grows into resolves instead of throwing. + // Imported inside the hoisted factory because the factory runs before the + // file's own imports are evaluated. + const { createMockWorkspaceClient, getMockFn } = + await import("../../../testing/mock-workspace-client"); + + const mockClient = createMockWorkspaceClient(); + + // Facade accessors are typed against the legacy SDK, so `.mockResolvedValue` + // on them would not typecheck. `getMockFn` is the typed handle; it mints + // idempotently, so these are the very functions the plugin will call. + const jobsApi = { + runNow: getMockFn(mockClient, "jobs.runNow"), + submit: getMockFn(mockClient, "jobs.submit"), + getRun: getMockFn(mockClient, "jobs.getRun"), + getRunOutput: getMockFn(mockClient, "jobs.getRunOutput"), + cancelRun: getMockFn(mockClient, "jobs.cancelRun"), + listRuns: getMockFn(mockClient, "jobs.listRuns"), + get: getMockFn(mockClient, "jobs.get"), + }; - const mockCacheInstance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_key: unknown[], fn: (signal?: AbortSignal) => Promise) => - fn(), - ), - generateKey: vi.fn(), - }; + const mockCacheInstance = { + get: vi.fn(), + set: vi.fn(), + delete: vi.fn(), + getOrExecute: vi.fn( + async ( + _key: unknown[], + fn: (signal?: AbortSignal) => Promise, + ) => fn(), + ), + generateKey: vi.fn(), + }; - return { mockJobsApi, mockClient, mockCacheInstance }; -}); + return { mockClient, jobsApi, mockCacheInstance }; + }, +); vi.mock("../../../workspace-client", async (importOriginal) => { const actual = @@ -290,7 +299,7 @@ describe("JobsPlugin", () => { test("runNow passes configured job_id to connector", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const exported = plugin.exports(); @@ -298,7 +307,7 @@ describe("JobsPlugin", () => { await handle.runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123 }), expect.anything(), ); @@ -307,7 +316,7 @@ describe("JobsPlugin", () => { test("runNow merges user params with configured job_id (no taskType)", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const exported = plugin.exports(); @@ -317,7 +326,7 @@ describe("JobsPlugin", () => { notebook_params: { key: "value" }, }); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123, notebook_params: { key: "value" }, @@ -349,7 +358,7 @@ describe("JobsPlugin", () => { test("runNow maps validated params to SDK fields when taskType is set", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({ jobs: { @@ -363,7 +372,7 @@ describe("JobsPlugin", () => { await handle.runNow({ key: "value" }); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123, notebook_params: { key: "value" }, @@ -375,7 +384,7 @@ describe("JobsPlugin", () => { test("runNow skips validation when no schema is configured", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -388,7 +397,7 @@ describe("JobsPlugin", () => { test("getRun wraps call in execute", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ + jobsApi.getRun.mockResolvedValue({ run_id: 1, state: { life_cycle_state: "TERMINATED" }, }); @@ -415,7 +424,7 @@ describe("JobsPlugin", () => { test("getJob wraps call in execute", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.get.mockResolvedValue({ job_id: 123 }); + jobsApi.get.mockResolvedValue({ job_id: 123 }); const plugin = new JobsPlugin({}); const executeSpy = vi.spyOn(plugin as any, "execute"); @@ -439,7 +448,7 @@ describe("JobsPlugin", () => { test("listRuns clamps caller-supplied limit before calling the SDK", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockReturnValue((async function* () {})()); + jobsApi.listRuns.mockReturnValue((async function* () {})()); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -447,7 +456,7 @@ describe("JobsPlugin", () => { await handle.listRuns({ limit: 10000 }); // SDK should receive the clamped limit, not the caller-supplied 10000. - expect(mockClient.jobs.listRuns).toHaveBeenCalledWith( + expect(jobsApi.listRuns).toHaveBeenCalledWith( expect.objectContaining({ limit: 100 }), expect.anything(), ); @@ -457,8 +466,8 @@ describe("JobsPlugin", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Pre-flight getRun verifies the run belongs to the configured jobId. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 1, job_id: 123 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 1, job_id: 123 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const executeSpy = vi.spyOn(plugin as any, "execute"); @@ -478,8 +487,8 @@ describe("JobsPlugin", () => { test("runAndWait yields status updates and terminates on TERMINATED", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); - mockClient.jobs.getRun + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.getRun .mockResolvedValueOnce({ run_id: 42, state: { life_cycle_state: "RUNNING" }, @@ -505,7 +514,7 @@ describe("JobsPlugin", () => { test("runAndWait throws when runNow returns no run_id", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({}); + jobsApi.runNow.mockResolvedValue({}); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -521,7 +530,7 @@ describe("JobsPlugin", () => { test("runNow returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockRejectedValue(new Error("API timeout")); + jobsApi.runNow.mockRejectedValue(new Error("API timeout")); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -538,9 +547,7 @@ describe("JobsPlugin", () => { test("cancelRun returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.cancelRun.mockRejectedValue( - new Error("Permission denied"), - ); + jobsApi.cancelRun.mockRejectedValue(new Error("Permission denied")); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -556,9 +563,7 @@ describe("JobsPlugin", () => { test("getRun returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockRejectedValue( - new Error("Internal server error"), - ); + jobsApi.getRun.mockRejectedValue(new Error("Internal server error")); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -574,7 +579,7 @@ describe("JobsPlugin", () => { test("listRuns returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockImplementation(() => { + jobsApi.listRuns.mockImplementation(() => { throw new Error("Auth failure"); }); @@ -594,7 +599,7 @@ describe("JobsPlugin", () => { const error = new Error("Detailed internal failure: db connection reset"); (error as any).statusCode = 403; - mockClient.jobs.getRun.mockRejectedValue(error); + jobsApi.getRun.mockRejectedValue(error); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -611,7 +616,7 @@ describe("JobsPlugin", () => { test("successful operations return ok result with data", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -628,7 +633,7 @@ describe("JobsPlugin", () => { test("getRun returns 404 when run.job_id does not match configured jobId", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -641,8 +646,8 @@ describe("JobsPlugin", () => { test("getRunOutput returns 404 when run belongs to another job", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); - mockClient.jobs.getRunOutput.mockResolvedValue({ logs: "nope" }); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.getRunOutput.mockResolvedValue({ logs: "nope" }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -651,14 +656,14 @@ describe("JobsPlugin", () => { expect(result.ok).toBe(false); if (!result.ok) expect(result.status).toBe(404); // Should never have called getRunOutput on the upstream SDK - expect(mockClient.jobs.getRunOutput).not.toHaveBeenCalled(); + expect(jobsApi.getRunOutput).not.toHaveBeenCalled(); }); test("cancelRun returns 404 when run belongs to another job", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -666,13 +671,13 @@ describe("JobsPlugin", () => { const result = await handle.cancelRun(99); expect(result.ok).toBe(false); if (!result.ok) expect(result.status).toBe(404); - expect(mockClient.jobs.cancelRun).not.toHaveBeenCalled(); + expect(jobsApi.cancelRun).not.toHaveBeenCalled(); }); test("getRun succeeds when run.job_id matches configured jobId", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ + jobsApi.getRun.mockResolvedValue({ run_id: 42, job_id: 123, state: { life_cycle_state: "TERMINATED" }, @@ -696,7 +701,7 @@ describe("JobsPlugin", () => { const { JobsConnector } = await import("../../../connectors/jobs"); const connector = new JobsConnector({}); - mockClient.jobs.get.mockResolvedValue({ job_id: 123 }); + jobsApi.get.mockResolvedValue({ job_id: 123 }); const controller = new AbortController(); await connector.getJob( @@ -718,8 +723,8 @@ describe("JobsPlugin", () => { test("runAndWait stops polling when signal is aborted", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); - mockClient.jobs.getRun.mockResolvedValue({ + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.getRun.mockResolvedValue({ run_id: 42, state: { life_cycle_state: "RUNNING" }, }); @@ -829,21 +834,21 @@ describe("JobsPlugin", () => { process.env.DATABRICKS_JOB_ETL = "100"; process.env.DATABRICKS_JOB_ML = "200"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 1 }); + jobsApi.runNow.mockResolvedValue({ run_id: 1 }); const plugin = new JobsPlugin({}); const exported = plugin.exports(); await exported("etl").runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 100 }), expect.anything(), ); - mockClient.jobs.runNow.mockClear(); + jobsApi.runNow.mockClear(); await exported("ml").runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 200 }), expect.anything(), ); @@ -1081,7 +1086,7 @@ describe("injectRoutes", () => { test("returns runId on successful non-streaming run", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1203,7 +1208,7 @@ describe("injectRoutes", () => { { run_id: 1, state: { life_cycle_state: "TERMINATED" } }, { run_id: 2, state: { life_cycle_state: "RUNNING" } }, ]; - mockClient.jobs.listRuns.mockReturnValue( + jobsApi.listRuns.mockReturnValue( (async function* () { for (const run of mockRuns) yield run; })(), @@ -1241,7 +1246,7 @@ describe("injectRoutes", () => { test("passes limit query param to listRuns", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockReturnValue((async function* () {})()); + jobsApi.listRuns.mockReturnValue((async function* () {})()); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1268,7 +1273,7 @@ describe("injectRoutes", () => { await handler(mockReq, mockRes); // Verify the connector was called with limit 5 - expect(mockClient.jobs.listRuns).toHaveBeenCalledWith( + expect(jobsApi.listRuns).toHaveBeenCalledWith( expect.objectContaining({ limit: 5 }), expect.anything(), ); @@ -1284,7 +1289,7 @@ describe("injectRoutes", () => { job_id: 123, state: { life_cycle_state: "TERMINATED" }, }; - mockClient.jobs.getRun.mockResolvedValue(mockRun); + jobsApi.getRun.mockResolvedValue(mockRun); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1351,7 +1356,7 @@ describe("injectRoutes", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Run exists upstream but is owned by job 456, not the configured 123. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1393,7 +1398,7 @@ describe("injectRoutes", () => { run_id: 42, state: { life_cycle_state: "TERMINATED" }, }; - mockClient.jobs.listRuns.mockReturnValue( + jobsApi.listRuns.mockReturnValue( (async function* () { yield mockRun; })(), @@ -1432,7 +1437,7 @@ describe("injectRoutes", () => { test("returns null status when no runs exist", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockReturnValue((async function* () {})()); + jobsApi.listRuns.mockReturnValue((async function* () {})()); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1469,8 +1474,8 @@ describe("injectRoutes", () => { test("cancels run and returns 204", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1540,8 +1545,8 @@ describe("injectRoutes", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Pre-flight getRun reports a run owned by a different job. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1570,7 +1575,7 @@ describe("injectRoutes", () => { expect(mockRes.status).toHaveBeenCalledWith(404); // Must not fall through to the cancel call or the 204. - expect(mockClient.jobs.cancelRun).not.toHaveBeenCalled(); + expect(jobsApi.cancelRun).not.toHaveBeenCalled(); expect(mockRes.end).not.toHaveBeenCalled(); }); @@ -1722,7 +1727,7 @@ describe("injectRoutes", () => { test("allows exactly MAX_UNVALIDATED_PARAM_KEYS (50) keys without schema", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({ jobs: { etl: { taskType: "notebook" } }, @@ -1760,13 +1765,13 @@ describe("injectRoutes", () => { // 50 keys is under the cap — request proceeds to the SDK. expect(mockRes.json).toHaveBeenCalledWith({ runId: 42 }); - expect(mockClient.jobs.runNow).toHaveBeenCalled(); + expect(jobsApi.runNow).toHaveBeenCalled(); }); test("allows undefined params", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1807,7 +1812,7 @@ describe("injectRoutes", () => { const error = new Error("Sensitive internal detail: token expired"); (error as any).statusCode = 403; - mockClient.jobs.runNow.mockRejectedValue(error); + jobsApi.runNow.mockRejectedValue(error); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1849,7 +1854,7 @@ describe("injectRoutes", () => { const error = new Error("Unauthorized"); (error as any).statusCode = 401; - mockClient.jobs.listRuns.mockImplementation(() => { + jobsApi.listRuns.mockImplementation(() => { throw error; }); @@ -1884,10 +1889,10 @@ describe("injectRoutes", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Pre-flight succeeds so we reach the actual cancel call. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); + jobsApi.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); const error = new Error("Forbidden"); (error as any).statusCode = 403; - mockClient.jobs.cancelRun.mockRejectedValue(error); + jobsApi.cancelRun.mockRejectedValue(error); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); 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/plugins/server/tests/server.integration.test.ts b/packages/appkit/src/plugins/server/tests/server.integration.test.ts index 6502af8ee..51036cbee 100644 --- a/packages/appkit/src/plugins/server/tests/server.integration.test.ts +++ b/packages/appkit/src/plugins/server/tests/server.integration.test.ts @@ -1,6 +1,10 @@ import type { Server } from "node:http"; -import { mockServiceContext, setupDatabricksEnv } from "@tools/test-helpers"; +import { + getListeningPort, + mockServiceContext, + setupDatabricksEnv, +} from "@databricks/appkit/testing"; import { afterAll, beforeAll, describe, expect, test } from "vitest"; // Set required env vars BEFORE imports that use them @@ -20,7 +24,9 @@ describe("ServerPlugin Integration", () => { let server: Server; let baseUrl: string; let serviceContextMock: Awaited>; - const TEST_PORT = 9876; // Use non-standard port to avoid conflicts + // This block alone pins a port, because it asserts the server honours a + // configured one. Every other block below uses an ephemeral port. + const TEST_PORT = 9876; beforeAll(async () => { setupDatabricksEnv(); @@ -37,7 +43,7 @@ describe("ServerPlugin Integration", () => { }); server = app.server.getServer(); - baseUrl = `http://127.0.0.1:${TEST_PORT}`; + baseUrl = `http://127.0.0.1:${await getListeningPort(server)}`; // Wait a bit for server to be ready await new Promise((resolve) => setTimeout(resolve, 100)); @@ -90,7 +96,6 @@ describe("ServerPlugin with custom plugin", () => { let server: Server; let baseUrl: string; let serviceContextMock: Awaited>; - const TEST_PORT = 9877; beforeAll(async () => { setupDatabricksEnv(); @@ -122,7 +127,7 @@ describe("ServerPlugin with custom plugin", () => { const app = await createApp({ plugins: [ serverPlugin({ - port: TEST_PORT, + port: 0, host: "127.0.0.1", }), testPlugin({}), @@ -130,9 +135,7 @@ describe("ServerPlugin with custom plugin", () => { }); server = app.server.getServer(); - baseUrl = `http://127.0.0.1:${TEST_PORT}`; - - await new Promise((resolve) => setTimeout(resolve, 100)); + baseUrl = `http://127.0.0.1:${await getListeningPort(server)}`; }); afterAll(async () => { @@ -174,7 +177,6 @@ describe("ServerPlugin with extend() via onPluginsReady", () => { let server: Server; let baseUrl: string; let serviceContextMock: Awaited>; - const TEST_PORT = 9878; beforeAll(async () => { setupDatabricksEnv(); @@ -184,7 +186,7 @@ describe("ServerPlugin with extend() via onPluginsReady", () => { const app = await createApp({ plugins: [ serverPlugin({ - port: TEST_PORT, + port: 0, host: "127.0.0.1", }), ], @@ -198,9 +200,7 @@ describe("ServerPlugin with extend() via onPluginsReady", () => { }); server = app.server.getServer(); - baseUrl = `http://127.0.0.1:${TEST_PORT}`; - - await new Promise((resolve) => setTimeout(resolve, 100)); + baseUrl = `http://127.0.0.1:${await getListeningPort(server)}`; }); afterAll(async () => { @@ -229,7 +229,6 @@ describe("createApp with async onPluginsReady callback", () => { let server: Server; let baseUrl: string; let serviceContextMock: Awaited>; - const TEST_PORT = 9885; beforeAll(async () => { setupDatabricksEnv(); @@ -239,7 +238,7 @@ describe("createApp with async onPluginsReady callback", () => { const app = await createApp({ plugins: [ serverPlugin({ - port: TEST_PORT, + port: 0, host: "127.0.0.1", }), ], @@ -254,9 +253,7 @@ describe("createApp with async onPluginsReady callback", () => { }); server = app.server.getServer(); - baseUrl = `http://127.0.0.1:${TEST_PORT}`; - - await new Promise((resolve) => setTimeout(resolve, 100)); + baseUrl = `http://127.0.0.1:${await getListeningPort(server)}`; }); afterAll(async () => { @@ -286,7 +283,6 @@ describe("ServerPlugin error handling for rejected async handlers", () => { let baseUrl: string; let serviceContextMock: Awaited>; let originalNodeEnv: string | undefined; - const TEST_PORT = 9879; const unhandledRejections: unknown[] = []; // Only count rejections raised by this suite's handlers — other suites in // the same worker may legitimately produce unrelated rejections. @@ -377,7 +373,7 @@ describe("ServerPlugin error handling for rejected async handlers", () => { const app = await createApp({ plugins: [ serverPlugin({ - port: TEST_PORT, + port: 0, host: "127.0.0.1", }), throwingPlugin({}), @@ -385,9 +381,7 @@ describe("ServerPlugin error handling for rejected async handlers", () => { }); server = app.server.getServer(); - baseUrl = `http://127.0.0.1:${TEST_PORT}`; - - await new Promise((resolve) => setTimeout(resolve, 100)); + baseUrl = `http://127.0.0.1:${await getListeningPort(server)}`; }); afterAll(async () => { diff --git a/packages/appkit/src/plugins/serving/tests/serving.test.ts b/packages/appkit/src/plugins/serving/tests/serving.test.ts index c273ff6a1..bca2f091a 100644 --- a/packages/appkit/src/plugins/serving/tests/serving.test.ts +++ b/packages/appkit/src/plugins/serving/tests/serving.test.ts @@ -4,8 +4,8 @@ import { createMockRequest, createMockResponse, createMockRouter, - mockServiceContext, setupDatabricksEnv, + useServiceContextMock, } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; @@ -50,18 +50,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/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index b19cd1a07..2a4852e12 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -162,10 +162,17 @@ export class TelemetryManager { /** * Flush and shut down the OpenTelemetry SDK. * - * Idempotent: the SDK reference is cleared synchronously and concurrent - * or repeated calls await the same in-flight flush. Awaited by the core - * lifecycle manager during graceful shutdown — that manager owns the - * process signal handlers, so telemetry no longer registers its own. + * Idempotent: the SDK reference is cleared synchronously and concurrent or + * repeated calls await the same in-flight flush. Awaited by the core lifecycle + * manager during graceful shutdown — that manager owns the process signal + * handlers, so telemetry no longer registers its own. + * + * Survives re-`initialize()`. `shutdownPromise` is deliberately *not* cleared + * when the flush settles, and that is safe: the memo is only ever returned + * after being reassigned for whatever SDK is currently live, so a stale + * resolved promise can only be returned when there is no SDK to flush. The + * covering test asserts every SDK across repeated + * initialize/shutdown cycles is flushed. */ async shutdown(): Promise { if (this.sdk) { @@ -182,4 +189,16 @@ export class TelemetryManager { return this.shutdownPromise; } + + /** + * Drop the singleton so the next {@link getInstance} builds a fresh manager. + * + * Does not flush: callers `shutdown()` first, then reset — the order + * `LifecycleManager.close()` uses. + * + * @internal + */ + static reset(): void { + TelemetryManager.instance = undefined; + } } diff --git a/packages/appkit/src/telemetry/tests/telemetry-manager-reset.test.ts b/packages/appkit/src/telemetry/tests/telemetry-manager-reset.test.ts new file mode 100644 index 000000000..06e6a8e45 --- /dev/null +++ b/packages/appkit/src/telemetry/tests/telemetry-manager-reset.test.ts @@ -0,0 +1,111 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +/** + * `_initialize` builds no SDK without `OTEL_EXPORTER_OTLP_ENDPOINT`, so these set + * it and mock `NodeSDK` to make the shutdown path observable. + * + * The never-cleared `shutdownPromise` was suspected of skipping a re-initialized + * SDK's flush. It does not — the memo is reassigned whenever an SDK is live — and + * the first test pins that so a future "cleanup" cannot change it. + */ + +const { sdkShutdown, NodeSDKMock } = vi.hoisted(() => { + const sdkShutdown = vi.fn().mockResolvedValue(undefined); + const NodeSDKMock = vi.fn(() => ({ + start: vi.fn(), + shutdown: sdkShutdown, + })); + return { sdkShutdown, NodeSDKMock }; +}); + +vi.mock("@opentelemetry/sdk-node", () => ({ NodeSDK: NodeSDKMock })); +vi.mock("@opentelemetry/auto-instrumentations-node", () => ({ + getNodeAutoInstrumentations: vi.fn(() => []), +})); +vi.mock("@opentelemetry/exporter-trace-otlp-proto", () => ({ + OTLPTraceExporter: vi.fn(() => ({})), +})); +vi.mock("@opentelemetry/exporter-metrics-otlp-proto", () => ({ + OTLPMetricExporter: vi.fn(() => ({})), +})); +vi.mock("@opentelemetry/exporter-logs-otlp-proto", () => ({ + OTLPLogExporter: vi.fn(() => ({})), +})); +vi.mock("@opentelemetry/resources", async () => { + const actual = await vi.importActual< + typeof import("@opentelemetry/resources") + >("@opentelemetry/resources"); + return { ...actual, detectResources: vi.fn(() => actual.emptyResource()) }; +}); + +import { TelemetryManager } from "../telemetry-manager"; + +describe("TelemetryManager re-bootability", () => { + let originalEndpoint: string | undefined; + + beforeEach(() => { + originalEndpoint = process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = "http://localhost:4318"; + vi.clearAllMocks(); + TelemetryManager.reset(); + }); + + afterEach(() => { + if (originalEndpoint === undefined) { + delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + } else { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = originalEndpoint; + } + TelemetryManager.reset(); + }); + + test("shutdown() twice across a re-initialize() flushes both SDKs", async () => { + TelemetryManager.initialize({}); + const manager = TelemetryManager.getInstance(); + expect(NodeSDKMock).toHaveBeenCalledTimes(1); + + await manager.shutdown(); + expect(sdkShutdown).toHaveBeenCalledTimes(1); + + // Re-initialize builds a *new* SDK, because shutdown() cleared `sdk`. + TelemetryManager.initialize({}); + expect(NodeSDKMock).toHaveBeenCalledTimes(2); + + await manager.shutdown(); + expect(sdkShutdown).toHaveBeenCalledTimes(2); + + // A third cycle, to pin the general property rather than one transition. + TelemetryManager.initialize({}); + await manager.shutdown(); + expect(NodeSDKMock).toHaveBeenCalledTimes(3); + expect(sdkShutdown).toHaveBeenCalledTimes(3); + }); + + test("concurrent shutdown() calls share one flush", async () => { + TelemetryManager.initialize({}); + const manager = TelemetryManager.getInstance(); + + await Promise.all([manager.shutdown(), manager.shutdown()]); + + // Clearing `sdk` synchronously is what makes this safe: the second caller + // finds no SDK and awaits the first caller's memo. + expect(sdkShutdown).toHaveBeenCalledTimes(1); + }); + + test("shutdown() with no SDK built resolves without flushing", async () => { + delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + TelemetryManager.initialize({}); + const manager = TelemetryManager.getInstance(); + + await expect(manager.shutdown()).resolves.toBeUndefined(); + expect(sdkShutdown).not.toHaveBeenCalled(); + }); + + test("reset() drops the singleton so the next getInstance() is fresh", () => { + const first = TelemetryManager.getInstance(); + TelemetryManager.reset(); + const second = TelemetryManager.getInstance(); + + expect(second).not.toBe(first); + }); +}); diff --git a/packages/appkit/src/testing/create-test-app.ts b/packages/appkit/src/testing/create-test-app.ts new file mode 100644 index 000000000..323d3b74c --- /dev/null +++ b/packages/appkit/src/testing/create-test-app.ts @@ -0,0 +1,357 @@ +/** + * Boot a real AppKit app with no workspace, credentials, or network, then call it + * over real HTTP. + */ + +import type { Server } from "node:http"; + +import type { + CacheConfig, + PluginConstructor, + PluginData, + PluginMap, +} from "shared"; + +import { InMemoryStorage } from "../cache/storage/memory"; +import { createApp } from "../core/appkit"; +import type { WorkspaceClient } from "../workspace-client"; +import type { OboOption } from "./fixtures"; +import { oboHeaders, setupDatabricksEnv } from "./fixtures"; +import type { CreateMockWorkspaceClientOptions } from "./mock-workspace-client"; +import { createMockWorkspaceClient } from "./mock-workspace-client"; +import { resetAppKitSingletons } from "./reset"; + +type Any = any; + +/** + * One baseline shared by every live harness app, reference-counted. + * + * A per-app snapshot does not compose: the second boot captures the first's + * mutations and whichever closes last re-applies them. Anchoring on the first + * boot and restoring on the last close makes the result order-independent. + */ +let envBaseline: NodeJS.ProcessEnv | undefined; +let liveHarnessApps = 0; + +/** Take the baseline on the first live app. */ +function acquireEnvBaseline(): void { + if (liveHarnessApps === 0) envBaseline = { ...process.env }; + liveHarnessApps += 1; +} + +/** Restore the baseline once no apps are left. */ +function releaseEnvBaseline(): void { + liveHarnessApps = Math.max(0, liveHarnessApps - 1); + if (liveHarnessApps > 0 || !envBaseline) return; + + const baseline = envBaseline; + envBaseline = undefined; + for (const key of Object.keys(process.env)) { + if (!(key in baseline)) delete process.env[key]; + } + Object.assign(process.env, baseline); +} + +/** Plugin descriptors, exactly as `createApp` takes them. */ +type Plugins = PluginData[]; + +/** Options for {@link createTestApp}. */ +export interface CreateTestAppOptions { + /** The plugins under test, as `createApp` takes them. */ + plugins?: T; + + /** Dotted-path responses for the built-in mock. Ignored when `client` is set. */ + responses?: CreateMockWorkspaceClientOptions["responses"]; + + /** + * Replaces the built-in mock. You then own `currentUser.me()` — boot reads + * `currentUser.id` and fails without it. + */ + client?: WorkspaceClient; + + /** Extra env for the boot, restored on `close()`; satisfies declared resources. */ + env?: Record; + + /** No socket; setup, validation, and teardown still run, request methods throw. */ + server?: false; + + /** + * Defaults to `"test"`. `"development"` is refused — it throws a `RangeError` + * in `get-port` on `port: 0`, boots Vite, and relaxes validation. + */ + nodeEnv?: string; + + /** Defaults to in-memory, which is what keeps boot offline. */ + cache?: CacheConfig; + + /** Teardown budget. Defaults to AppKit's programmatic budget. */ + closeTimeoutMs?: number; +} + +/** Per-request options for the {@link TestApp} HTTP methods. */ +export interface TestRequestOptions { + /** A non-string value is JSON-encoded with `content-type: application/json`. */ + body?: unknown; + /** Merged last, so they win over anything the harness sets. */ + headers?: Record; + /** Same convention as `createMockRequest({ obo })`. */ + obo?: OboOption; + /** Forwarded to `fetch`. */ + signal?: AbortSignal; +} + +/** A booted test app. */ +export interface TestApp { + /** + * Plugin exports by manifest name. Nested rather than spread because `get` and + * `delete` are plausible plugin names and would collide with the request methods. + */ + plugins: PluginMap; + /** The same object a handler resolves at runtime. */ + client: WorkspaceClient; + /** e.g. `http://127.0.0.1:54321`. Throws when `server: false`. */ + baseUrl: string; + /** The bound ephemeral port. Throws when `server: false`. */ + port: number; + /** The underlying HTTP server, or `undefined` with `server: false`. */ + server?: Server; + + /** Release the app and restore env. Idempotent. */ + close(): Promise; + [Symbol.asyncDispose](): Promise; + + get(path: string, options?: TestRequestOptions): Promise; + post(path: string, options?: TestRequestOptions): Promise; + put(path: string, options?: TestRequestOptions): Promise; + patch(path: string, options?: TestRequestOptions): Promise; + delete(path: string, options?: TestRequestOptions): Promise; +} + +/** + * `start()` returns once `listen()` is invoked, before the bind completes, so + * `address()` is null until the `listening` event fires. + * + * @internal + */ +export async function getListeningPort(server: Server): Promise { + const addr = server.address(); + if (addr && typeof addr === "object" && typeof addr.port === "number") { + return addr.port; + } + await new Promise((resolve, reject) => { + server.once("listening", () => resolve()); + server.once("error", (err) => reject(err)); + }); + const ready = server.address(); + if (!ready || typeof ready !== "object") { + throw new Error("Server is listening but address() returned null"); + } + return ready.port; +} + +/** + * Boot a real app — real Express wiring, routes, and resource validation — with + * no workspace, credentials, or network. `createTestPluginContext` is cheaper + * when you only need to unit-test wiring. + * + * Does **not** validate config values against `manifest.config.schema`; no + * runtime validator exists for that. + * + * @example + * ```ts + * const app = await createTestApp({ plugins: [myPlugin()] }); + * try { + * const res = await app.post("/api/my-plugin/thing", { body: { q: 1 }, obo: true }); + * await expectStream(res).toEmit("status", "result"); + * } finally { + * await app.close(); + * } + * ``` + */ +export async function createTestApp( + options: CreateTestAppOptions = {}, +): Promise> { + const { + plugins = [] as unknown as T, + responses, + client: suppliedClient, + env = {}, + server: serverOption, + nodeEnv = "test", + cache, + closeTimeoutMs, + } = options; + + if (nodeEnv === "development") { + throw new Error( + 'createTestApp: nodeEnv "development" is not supported. Dev mode routes ' + + "the harness's ephemeral `port: 0` through get-port, which throws a " + + "RangeError, and it also boots a real Vite dev server, downgrades " + + "resource validation to a warning, and stops filtering dev-only " + + "plugins. Pin a port explicitly with your own server plugin if you " + + "need dev behaviour.", + ); + } + + // Wholesale rather than a whitelist: plugins read vars we cannot enumerate. + acquireEnvBaseline(); + const restoreEnv = releaseEnvBaseline; + + let app: Awaited> | undefined; + + try { + process.env.NODE_ENV = nodeEnv; + + // Redundant while NODE_ENV is pinned, but keeps the throw-on-missing-resource + // contract if that pin ever changes. No opt-out: the warning path is + // dev-only, and dev is refused. + process.env.APPKIT_STRICT_VALIDATION = "true"; + + // The workspace ID short-circuits getWorkspaceId's SCIM probe, which would + // otherwise show up as an apiClient.request call. + setupDatabricksEnv({ + DATABRICKS_WORKSPACE_ID: "test-workspace-id", + ...env, + }); + + resetAppKitSingletons(); + + // Boot runs ServiceContext.createContext for real, which reads + // currentUser.id — the mock's built-in default is what lets it through. + const client = suppliedClient ?? createMockWorkspaceClient({ responses }); + + // createApp never auto-adds a server, so without this there is nothing to + // fetch. Lazily imported: the plugin runs dotenv.config() at module load, so + // a static import would mutate a consumer's env on import of this kit. + const hasServer = plugins.some((p) => p?.name === "server"); + if (serverOption === false && hasServer) { + // The plugin would still bind a socket while the handle denied one existed. + throw new Error( + "createTestApp: `server: false` conflicts with the server plugin in " + + "`plugins`. Drop one — omit `server: false` to use your plugin, or " + + "remove the plugin to boot without a socket.", + ); + } + const bootPlugins = [...plugins] as Plugins; + if (serverOption !== false && !hasServer) { + const { server: serverPlugin } = await import("../plugins/server"); + bootPlugins.push(serverPlugin({ port: 0, host: "127.0.0.1" })); + } + + // Both extras are load-bearing: without explicit storage the cache builds its + // own client and probes Lakebase over the network, and without the opt-out + // TelemetryReporter fires an apiClient.request on boot. + app = await createApp({ + plugins: bootPlugins as Any, + client, + cache: cache ?? { + storage: new InMemoryStorage({ enabled: true } as Any), + }, + disableInternalTelemetry: true, + }); + + const serverExports = (app as Any).server; + const httpServer: Server | undefined = + serverOption === false ? undefined : serverExports?.getServer?.(); + const port = httpServer ? await getListeningPort(httpServer) : undefined; + const baseUrl = port === undefined ? undefined : `http://127.0.0.1:${port}`; + + const bootedApp = app; + let closed: Promise | undefined; + + /** Memoized, so repeated calls are safe in nested `finally`s. */ + const close = () => { + closed ??= (async () => { + try { + await bootedApp.close( + closeTimeoutMs === undefined ? {} : { timeoutMs: closeTimeoutMs }, + ); + } finally { + // close() resets these already; belt and braces for a caller who + // supplied their own server plugin. + resetAppKitSingletons(); + restoreEnv(); + } + })(); + return closed; + }; + + const request = async ( + method: string, + path: string, + reqOptions: TestRequestOptions = {}, + ): Promise => { + if (baseUrl === undefined) { + throw new Error( + "createTestApp: no HTTP server was started (server: false), so " + + `${method} ${path} cannot be issued.`, + ); + } + + const headers: Record = {}; + if (reqOptions.obo) { + Object.assign(headers, oboHeaders(reqOptions.obo)); + } + + let body: string | undefined; + if (reqOptions.body !== undefined) { + if (typeof reqOptions.body === "string") { + body = reqOptions.body; + } else { + body = JSON.stringify(reqOptions.body); + headers["content-type"] = "application/json"; + } + } + + // Caller headers last, so an explicit content-type or identity wins. + Object.assign(headers, reqOptions.headers ?? {}); + + return fetch(new URL(path, baseUrl), { + method, + headers, + body, + signal: reqOptions.signal, + }); + }; + + return { + plugins: bootedApp as unknown as PluginMap, + client, + get baseUrl() { + if (baseUrl === undefined) { + throw new Error( + "createTestApp: no HTTP server was started (server: false).", + ); + } + return baseUrl; + }, + get port() { + if (port === undefined) { + throw new Error( + "createTestApp: no HTTP server was started (server: false).", + ); + } + return port; + }, + server: httpServer, + close, + [Symbol.asyncDispose]: close, + get: (path, o) => request("GET", path, o), + post: (path, o) => request("POST", path, o), + put: (path, o) => request("PUT", path, o), + patch: (path, o) => request("PATCH", path, o), + delete: (path, o) => request("DELETE", path, o), + }; + } catch (err) { + // Teardown must run from the failure path too, or the boot leaks env + // mutations and singletons into every later test in the file. + try { + await (app as Any)?.close?.(); + } catch { + // The boot error is the interesting one; don't let teardown mask it. + } + resetAppKitSingletons(); + restoreEnv(); + throw err; + } +} diff --git a/packages/appkit/src/testing/create-test-plugin.ts b/packages/appkit/src/testing/create-test-plugin.ts new file mode 100644 index 000000000..b969e2799 --- /dev/null +++ b/packages/appkit/src/testing/create-test-plugin.ts @@ -0,0 +1,27 @@ +import type { PluginConstructor, PluginData } from "shared"; + +/** + * Instantiate a plugin from its `toPlugin()` factory for use with + * `createTestPluginContext`. + * + * Merge order mirrors `AppKit.createAndRegisterPlugin` — `DEFAULT_CONFIG`, then + * the factory's config, then the manifest `name` — so the instance matches what + * production builds. Reaching through the descriptor by hand + * (`new (genie({}).plugin)({})`) skips both. + */ +export function createTestPlugin< + TClass extends PluginConstructor, + TConfig, + TName extends string, +>( + factory: (config?: TConfig) => PluginData, + config?: TConfig, +): InstanceType { + const { plugin: PluginClass, config: factoryConfig, name } = factory(config); + + return new PluginClass({ + ...(PluginClass.DEFAULT_CONFIG ?? {}), + ...(factoryConfig ?? {}), + name, + }) as InstanceType; +} diff --git a/packages/appkit/src/testing/expect-stream.ts b/packages/appkit/src/testing/expect-stream.ts new file mode 100644 index 000000000..1ee8b05e5 --- /dev/null +++ b/packages/appkit/src/testing/expect-stream.ts @@ -0,0 +1,331 @@ +/** + * 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; +} + +/** + * 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, + * - a captured mock response ({@link createMockResponse}) — its written SSE + * bytes are parsed into events. + */ +export type StreamSource = + | AsyncIterable + | Iterable + | Response + | 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 { + /** + * 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[] = []; + // 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 line of block.split("\n")) { + 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. + } + + // 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 = {}; + 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 }; + } + } + + // 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({ + ...parsed, + type: name ?? (parsed.type as string | undefined), + }); + } + + return events; +} + +async function collectEventsInner( + source: StreamSource, +): 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[] = []; + 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", + ); +} + +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; + for (const type of actual) { + if (i < expected.length && type === expected[i]) i++; + if (i === expected.length) break; + } + 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 createTestPluginContext} 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"); + * ``` + * + * @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"); + * ``` + * + * @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, + options: ExpectStreamOptions = {}, +): StreamAssertion { + const events = collectEvents(source, options.timeout); + + 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 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. + * + * @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 events = parseSSEBody(text); + const last = events.at(-1); + + if (!last) { + throw new Error(`No data found in SSE response: ${text}`); + } + + // `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/fixtures.ts b/packages/appkit/src/testing/fixtures.ts new file mode 100644 index 000000000..35f9a9720 --- /dev/null +++ b/packages/appkit/src/testing/fixtures.ts @@ -0,0 +1,564 @@ +import type { Span, SpanOptions } from "@opentelemetry/api"; +import type { IAppRouter } from "shared"; +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 { InstrumentConfig, ITelemetry } from "../telemetry/types"; +import { createMockWorkspaceClient } from "./mock-workspace-client"; + +// Test fixtures intentionally use loose shapes; `noExplicitAny` is disabled +// repo-wide (see .oxlintrc.json), so a local alias keeps the intent readable. +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 createTestPluginContext} 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}`], + }; +} + +/** + * 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. + * + * Exported so `createTestApp`'s request methods use the same convention as + * `createMockRequest` rather than a second one. + * + * @internal + */ +export 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 { 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: {}, + header: function (name: string) { + return this.headers[name.toLowerCase()]; + }, + // `...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; +} + +/** + * 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 + // 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, 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) { + handler(); + } + } + return this; + }), + 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) { + 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, + /** + * 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; +} + +/** + * 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); +} + +/** + * 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 + */ +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; +} + +/** + * 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. + */ +function buildServiceContextState( + options: TestContextOptions = {}, +): ServiceContextState { + return { + 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"), + }; +} + +/** + * 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 = buildServiceContextState(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); + + const createUserContextSpy = vi + .spyOn(ServiceContext, "createUserContext") + .mockImplementation((_token: string, userId: string, userName?: string) => { + return { + client: (options.userDatabricksClient || + createMockWorkspaceClient()) 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(); + }, + }; +} + +/** 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. + */ +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. + * + * @deprecated Use `createMockWorkspaceClient({ defaults: false })` with + * `getMockFn(client, "statementExecution.executeStatement")` instead — it fakes + * the whole facade rather than two services, so a plugin that reaches any other + * service does not crash. + * + * Left byte-for-byte unchanged rather than reimplemented on the new builder, + * because the semantics differ in a way its one remaining caller can observe: + * these bare `vi.fn()`s return `undefined` **synchronously**, whereas the new + * floor returns `Promise`. + */ +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..a46057450 --- /dev/null +++ b/packages/appkit/src/testing/index.ts @@ -0,0 +1,95 @@ +/** + * @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. + * + * Three entry points: + * - {@link createTestApp} — boot a real app with a faked data plane and call it + * over real HTTP. The recommended starting point. + * - {@link createTestPluginContext} — build a real `PluginContext` with faked edges + * and attach it to a plugin, with no boot and no socket. + * - {@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 { 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 = createTestPluginContext({ analytics: { query: fixtureRows } }); + * const plugin = new MyPlugin({}); + * await mock.attach(plugin); + * + * await expectStream(plugin.streamSomething(input)).toEmit( + * "tool_call", + * "message_delta", + * ); + * ``` + * + * @module + */ + +// 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"; +export { + createTestApp, + type CreateTestAppOptions, + getListeningPort, + type TestApp, + type TestRequestOptions, +} from "./create-test-app"; +export { + type CapturedSSEResponse, + type ExpectStreamOptions, + expectStream, + parseSSEResponse, + type StreamAssertion, + type StreamEvent, + type StreamSource, +} from "./expect-stream"; +export { + createConfigurableMockWorkspaceClient, + createFailedSQLResponse, + createMockRequest, + createMockResponse, + createMockRouter, + createMockTelemetry, + createSuccessfulSQLResponse, + mockServiceContext, + type OboOption, + resetTestCache, + runWithRequestContext, + type ServiceContextMock, + setupDatabricksEnv, + type TestContextOptions, + useServiceContextMock, +} from "./fixtures"; +export { + createMockWorkspaceClient, + type CreateMockWorkspaceClientOptions, + getMockFn, + type MockWorkspaceClient, +} from "./mock-workspace-client"; +export { createTestPlugin } from "./create-test-plugin"; +export { resetAppKitSingletons } from "./reset"; +export { + createTestPluginContext, + type FakeProvider, + type FakeProviders, + type FakeToolResponse, + type RecordedRoute, + type RecordedToolCall, + type TestPluginContext, +} from "./test-plugin-context"; diff --git a/packages/appkit/src/testing/mock-workspace-client.ts b/packages/appkit/src/testing/mock-workspace-client.ts new file mode 100644 index 000000000..a427be8f0 --- /dev/null +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -0,0 +1,262 @@ +/** + * A never-crash fake `WorkspaceClient`. Declared paths resolve their value; + * everything else resolves `undefined` instead of throwing. + */ + +import type { Mock } from "vitest"; +import { vi } from "vitest"; + +import type { WorkspaceClient } from "../workspace-client"; + +type Any = any; + +type LegacyClient = ReturnType; + +/** Options for {@link createMockWorkspaceClient}. */ +export interface CreateMockWorkspaceClientOptions { + /** + * Responses keyed by dotted path (`"jobs.getRun"`). A function value is called + * with the arguments, so a test can script behaviour or reject. + */ + responses?: Record; + + /** Seed `config`; `host` must stay a real string. */ + config?: Partial; + + /** Apply the canned defaults (SQL succeeds, warehouse RUNNING). Default true. */ + defaults?: boolean; +} + +export type MockWorkspaceClient = WorkspaceClient; + +/** + * Applied beneath caller-supplied `responses`. + * + * The first three must stay byte-identical to the old `fixtures.ts` values — 13 + * suites reach them implicitly via `mockServiceContext`. `currentUser.me` is + * required: `ServiceContext.createContext` reads `.id`, so `createApp({ client })` + * cannot boot without it. + */ +const DEFAULT_RESPONSES: Record = { + "statementExecution.executeStatement": { + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }, + "warehouses.get": { state: "RUNNING" }, + "warehouses.start": undefined, + "currentUser.me": { + id: "test-service-user", + userName: "test-service-user", + }, +}; + +/** The seven generically-proxied services; `config`/`apiClient` are seeded below. */ +const FACADE_SERVICES = [ + "files", + "warehouses", + "genie", + "jobs", + "statementExecution", + "servingEndpoints", + "currentUser", +] as const; + +/** + * Answered with `undefined` rather than a minted mock. `then` is load-bearing: + * without it a service is thenable, so `await client.jobs` hangs. + */ +const PASSTHROUGH_DENY: ReadonlySet = new Set([ + "then", + "catch", + "finally", + "toJSON", + "inspect", + "constructor", + "$$typeof", + "asymmetricMatch", +]); + +/** + * Symbols and denied names short-circuit before minting; anything already on the + * target (seeded members, `Object.prototype`) wins. + * + * `ownKeys`/`getOwnPropertyDescriptor` stay at their defaults on purpose — + * reporting keys makes `util.inspect` probe each one, minting a mock per probe. + */ +function neverCrashGet(namespace: string, mint: (path: string) => Mock) { + return (target: Any, prop: Any): Any => { + if (typeof prop === "symbol") return Reflect.get(target, prop); + if (PASSTHROUGH_DENY.has(prop)) return undefined; + if (prop in target) return target[prop]; + return mint(`${namespace}.${String(prop)}`); + }; +} + +// In a WeakMap, not on the client: a stray own property would show up in +// util.inspect, toEqual, and key enumeration. +const clientFns = new WeakMap>(); + +/** + * @example + * ```ts + * const client = createMockWorkspaceClient({ + * responses: { "jobs.getRun": { state: "TERMINATED" } }, + * }); + * ``` + */ +export function createMockWorkspaceClient( + options: CreateMockWorkspaceClientOptions = {}, +): MockWorkspaceClient { + const { responses = {}, config = {}, defaults = true } = options; + + // Caller entries win over the canned defaults for the same path. + const merged: Record = defaults + ? { ...DEFAULT_RESPONSES, ...responses } + : { ...responses }; + + // Shared with the legacy view and getMockFn, so both see the same functions. + const fns = new Map(); + + /** Mint once per path, so call assertions see a stable reference. */ + function mint(path: string): Mock { + const cached = fns.get(path); + if (cached) return cached; + + const response = merged[path]; + const fn = vi.fn(); + if (typeof response === "function") fn.mockImplementation(response); + else fn.mockResolvedValue(response); + + fns.set(path, fn); + return fn; + } + + /** Memoized, so `client.jobs === client.jobs`. */ + const services = new Map(); + function service(namespace: string): Any { + const cached = services.get(namespace); + if (cached) return cached; + const proxy = new Proxy({}, { get: neverCrashGet(namespace, mint) }); + services.set(namespace, proxy); + return proxy; + } + + /** Pull `"config.*"` / `"apiClient.*"` entries out so they seed real values. */ + function seededOverrides(namespace: string): Record { + const prefix = `${namespace}.`; + const out: Record = {}; + for (const [key, value] of Object.entries(merged)) { + if (key.startsWith(prefix)) out[key.slice(prefix.length)] = value; + } + return out; + } + + // `host` is read as a string and throws if falsy, so it cannot be a mock. + const configTarget: Record = { + host: "https://test.databricks.com", + authenticate: vi.fn((headers?: Headers) => { + headers?.set?.("Authorization", "Bearer test-token"); + }), + ensureResolved: vi.fn().mockResolvedValue(undefined), + ...config, + ...seededOverrides("config"), + }; + + // userAgent() must be synchronous (a Promise stringifies to "[object Promise]" + // inside a Headers value); request resolves {} so destructuring works. + const apiClientTarget: Record = { + userAgent: vi.fn().mockReturnValue("appkit-test/1.0"), + request: vi.fn().mockResolvedValue({}), + }; + for (const [key, value] of Object.entries(seededOverrides("apiClient"))) { + const fn = typeof value === "function" ? vi.fn(value) : vi.fn(); + if (typeof value !== "function") fn.mockResolvedValue(value); + apiClientTarget[key] = fn; + fns.set(`apiClient.${key}`, fn); + } + for (const key of ["userAgent", "request"]) { + if (!fns.has(`apiClient.${key}`)) { + fns.set(`apiClient.${key}`, apiClientTarget[key] as Mock); + } + } + for (const [key, value] of Object.entries(configTarget)) { + if (typeof value === "function" && !fns.has(`config.${key}`)) { + fns.set(`config.${key}`, value as Mock); + } + } + + const configProxy = new Proxy(configTarget, { + get: neverCrashGet("config", mint), + }); + const apiClientProxy = new Proxy(apiClientTarget, { + get: neverCrashGet("apiClient", mint), + }); + + /** Memoized; routes facade names onto the same objects, others onto the floor. */ + let legacy: LegacyClient | undefined; + function toLegacyWorkspaceClient(): LegacyClient { + legacy ??= new Proxy( + {}, + { + get: (target: Any, prop: Any): Any => { + if (typeof prop === "symbol") return Reflect.get(target, prop); + if (PASSTHROUGH_DENY.has(prop)) return undefined; + if (prop === "config") return configProxy; + if (prop === "apiClient") return apiClientProxy; + if (prop === "toLegacyWorkspaceClient") { + return toLegacyWorkspaceClient; + } + return service(String(prop)); + }, + }, + ) as LegacyClient; + return legacy; + } + + const client: WorkspaceClient = { + ...(Object.fromEntries( + FACADE_SERVICES.map((name) => [name, service(name)]), + ) as Pick), + config: configProxy as WorkspaceClient["config"], + apiClient: apiClientProxy as WorkspaceClient["apiClient"], + toLegacyWorkspaceClient, + }; + + clientFns.set(client, fns); + return client; +} + +/** + * The typed assertion path onto a mocked method — facade accessors are SDK-typed, + * so `expect(client.jobs.getRun).toHaveBeenCalled()` does not typecheck. + * + * Minting is idempotent, so this can be called before the code under test runs. + * Throws for a non-function member such as `"config.host"`. + */ +export function getMockFn(client: MockWorkspaceClient, path: string): Mock { + const fns = clientFns.get(client); + if (!fns) { + throw new Error( + "getMockFn: not a createMockWorkspaceClient() client. Pass the client " + + "the builder returned, not a hand-rolled object.", + ); + } + + const cached = fns.get(path); + if (cached) return cached; + + const dot = path.indexOf("."); + const namespace = dot === -1 ? path : path.slice(0, dot); + const member = dot === -1 ? "" : path.slice(dot + 1); + const resolved = member + ? (client as Any)[namespace]?.[member] + : (client as Any)[namespace]; + + if (typeof resolved !== "function") { + throw new Error( + `getMockFn: "${path}" is not a mocked function (got ${typeof resolved}). ` + + "Members seeded with a real value, such as config.host, have no mock.", + ); + } + return resolved as Mock; +} diff --git a/packages/appkit/src/testing/reset.ts b/packages/appkit/src/testing/reset.ts new file mode 100644 index 000000000..5124da5ea --- /dev/null +++ b/packages/appkit/src/testing/reset.ts @@ -0,0 +1,13 @@ +import { resetCoreSingletons } from "../core/reset-singletons"; + +/** + * Drop the process-wide singletons `createApp()` initializes, so a file can boot + * more than one app. + * + * Pointer drops, not teardown — always close first, or the old app's pools and + * exporters leak. `app.close()` already does both, so this is only for tests + * that hand-roll `createApp`. + */ +export function resetAppKitSingletons(): void { + resetCoreSingletons(); +} diff --git a/packages/appkit/src/testing/test-plugin-context.ts b/packages/appkit/src/testing/test-plugin-context.ts new file mode 100644 index 000000000..721e3cdbe --- /dev/null +++ b/packages/appkit/src/testing/test-plugin-context.ts @@ -0,0 +1,361 @@ +import type express from "express"; +import type { + AgentToolDefinition, + BasePlugin, + IAppRequest, + ToolProvider, +} from "shared"; + +import { CacheManager } from "../cache"; +import { InMemoryStorage } from "../cache/storage"; +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"; + +/** + * 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; + +/** + * 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 + * createTestPluginContext({ 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 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. + * + * The fake replicates the token precondition only, not the real dev-mode + * OTel `isDevOboFallback()` marker — assert OBO here, not via that flag. + */ + 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`. */ +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 createTestPluginContext}: the real `PluginContext` plus the + * seams a test needs to drive and inspect it. + */ +export interface TestPluginContext { + /** The real {@link PluginContext}, constructed with mock telemetry. */ + ctx: PluginContext; + /** + * 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 — + * 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 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`. + * + * Nothing about `PluginContext` is reimplemented. + * + * @param fakes - Canned tool responses keyed by plugin then tool name. + * + * @example + * ```ts + * 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 createTestPluginContext( + fakes: FakeProviders = {}, +): TestPluginContext { + 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, + userId: string | undefined, + ): Promise => { + 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( + `createTestPluginContext: 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, + signal, + ) + : response; + }; + + const base: ToolProvider = { + getAgentTools: () => record.tools, + executeAgentTool: (toolName, args, signal) => + resolve(toolName, args, signal, false, undefined), + }; + + // 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`. + // + // 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) + .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 { + ...base, + executeAgentTool: (toolName, args, signal) => + resolve(toolName, args, signal, true, userId), + }; + }; + + // `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 }); + + // 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; + } + + 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/create-test-app-http.test.ts b/packages/appkit/src/testing/tests/create-test-app-http.test.ts new file mode 100644 index 000000000..f4e79a39d --- /dev/null +++ b/packages/appkit/src/testing/tests/create-test-app-http.test.ts @@ -0,0 +1,247 @@ +import type { + IAppRequest, + IAppResponse, + IAppRouter, + PluginManifest, +} from "shared"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; + +import { getUserContext } from "../../context/execution-context"; +import { Plugin, toPlugin } from "../../plugin"; +import { createTestApp, type TestApp } from "../create-test-app"; +import { expectStream } from "../expect-stream"; + +/** + * The HTTP layer: `app.get/post/put/patch/delete` against a real Express stack. + * + * One app for the whole file — every assertion here is about the request, not + * about boot, so re-booting per test would only slow it down. + */ + +class HttpPlugin extends Plugin { + static manifest = { + name: "http", + displayName: "Http", + version: "0.0.0", + description: "HTTP layer probe", + resources: { required: [], optional: [] }, + } as unknown as PluginManifest; + + injectRoutes(router: IAppRouter): void { + // Registered through `this.route()`, the way real plugins do, rather than + // raw `router.get()`. That is what wraps each handler in + // forwardAsyncErrors, so a rejection reaches errorHandlerMiddleware instead + // of hanging the request — see the /boom test. + const get = ( + name: string, + path: string, + handler: (req: IAppRequest, res: IAppResponse) => Promise, + ) => this.route(router, { name, method: "get", path, handler }); + + get("json", "/json", async (_req, res) => { + res.status(201).json({ ok: true, method: "GET" }); + }); + + this.route(router, { + name: "echo", + method: "post", + path: "/echo", + handler: async (req, res) => { + res.json({ + body: req.body, + contentType: req.headers["content-type"] ?? null, + }); + }, + }); + + get("headers", "/headers", async (req, res) => { + res.json({ + custom: req.headers["x-custom"] ?? null, + user: req.headers["x-forwarded-user"] ?? null, + token: req.headers["x-forwarded-access-token"] ?? null, + email: req.headers["x-forwarded-email"] ?? null, + }); + }); + + // Uses the real asUser path, so the forwarded identity has to be genuine. + get("asUser", "/as-user", async (req, res) => { + const exports = this.asUser(req).exports() as { + whoami: () => { userId?: string }; + }; + res.json(exports.whoami()); + }); + + get("boom", "/boom", async () => { + throw new Error("handler exploded"); + }); + + this.route(router, { + name: "stream", + method: "post", + path: "/stream", + handler: async (_req, res) => { + res.setHeader("content-type", "text/event-stream"); + res.write(`event: status\ndata: ${JSON.stringify({ s: "start" })}\n\n`); + res.write(`event: result\ndata: ${JSON.stringify({ rows: [1] })}\n\n`); + res.end(); + }, + }); + + for (const method of ["put", "patch"] as const) { + this.route(router, { + name: `verb-${method}`, + method, + path: "/verb", + handler: async (req, res) => { + res.json({ m: method.toUpperCase(), b: req.body }); + }, + }); + } + this.route(router, { + name: "verb-delete", + method: "delete", + path: "/verb", + handler: async (_req, res) => { + res.json({ m: "DELETE" }); + }, + }); + } + + exports() { + return { + whoami: () => { + const ctx = getUserContext(); + return { userId: ctx?.userId }; + }, + }; + } +} +const http = toPlugin(HttpPlugin); + +describe("createTestApp HTTP layer", () => { + let app: TestApp<[ReturnType]>; + + beforeAll(async () => { + app = await createTestApp({ plugins: [http()] }); + }); + + afterAll(async () => { + await app?.close(); + }); + + test("GET returns the plugin's JSON body and status", async () => { + const res = await app.get("/api/http/json"); + expect(res.status).toBe(201); + await expect(res.json()).resolves.toEqual({ ok: true, method: "GET" }); + }); + + test("POST with an object body arrives JSON-parsed at the handler", async () => { + const res = await app.post("/api/http/echo", { + body: { q: 1, nested: [2] }, + }); + + // Proves the real express.json() middleware ran, not a shortcut. + await expect(res.json()).resolves.toEqual({ + body: { q: 1, nested: [2] }, + contentType: "application/json", + }); + }); + + test("POST with a string body and explicit content-type passes through unmodified", async () => { + const res = await app.post("/api/http/echo", { + body: "raw text, not JSON", + headers: { "content-type": "text/plain" }, + }); + + // express.json() ignores a non-JSON content-type, so the handler sees an + // empty body — the point is that the harness did not re-encode or override. + await expect(res.json()).resolves.toMatchObject({ + contentType: "text/plain", + }); + }); + + test("custom headers reach the handler and win over harness defaults", async () => { + const res = await app.get("/api/http/headers", { + obo: true, + headers: { "x-custom": "hello", "x-forwarded-user": "override" }, + }); + + await expect(res.json()).resolves.toMatchObject({ + custom: "hello", + // The explicit header beats the one `obo` generated. + user: "override", + token: "test-user-token", + }); + }); + + test("obo: true sets the forwarded identity headers", async () => { + const res = await app.get("/api/http/headers", { obo: true }); + await expect(res.json()).resolves.toMatchObject({ + user: "test-user", + token: "test-user-token", + }); + }); + + test("obo: { userId, email } overrides the identity", async () => { + const res = await app.get("/api/http/headers", { + obo: { userId: "alice", email: "alice@example.com" }, + }); + await expect(res.json()).resolves.toMatchObject({ + user: "alice", + email: "alice@example.com", + }); + }); + + test("a handler using asUser resolves the forwarded test user", async () => { + const res = await app.get("/api/http/as-user", { obo: { userId: "bob" } }); + // The real user-context path, driven entirely by the `obo` flag. + await expect(res.json()).resolves.toEqual({ userId: "bob" }); + }); + + test("an SSE route composes with expectStream directly", async () => { + // The dogfooding report's #1 friction, avoided by construction: the request + // methods return a native Response, which expectStream already accepts. + const res = await app.post("/api/http/stream"); + await expectStream(res).toEmit("status", "result"); + }); + + test("a throwing handler produces the real error-middleware response", async () => { + const res = await app.get("/api/http/boom"); + + // Handled by the real errorHandlerMiddleware rather than escaping as an + // unhandled rejection that would hang the request and fail the run. + expect(res.status).toBe(500); + + // The message is included because errorHandlerMiddleware redacts only when + // NODE_ENV === "production", and the harness pins "test". That is the + // useful behaviour for a test — an assertion can name the failure — but it + // does mean this response shape is the dev one, not what a deployed app + // returns to a client. + await expect(res.json()).resolves.toEqual({ error: "handler exploded" }); + }); + + test("an unmounted path is a 404", async () => { + const res = await app.get("/api/http/nope"); + expect(res.status).toBe(404); + }); + + test("put, patch, and delete reach their handlers", async () => { + await expect( + app.put("/api/http/verb", { body: { a: 1 } }).then((r) => r.json()), + ).resolves.toEqual({ m: "PUT", b: { a: 1 } }); + await expect( + app.patch("/api/http/verb", { body: { a: 2 } }).then((r) => r.json()), + ).resolves.toEqual({ m: "PATCH", b: { a: 2 } }); + await expect( + app.delete("/api/http/verb").then((r) => r.json()), + ).resolves.toEqual({ m: "DELETE" }); + }); + + test("a signal aborts an in-flight request", async () => { + const controller = new AbortController(); + controller.abort(); + await expect( + app.get("/api/http/json", { signal: controller.signal }), + ).rejects.toThrow(); + }); +}); diff --git a/packages/appkit/src/testing/tests/create-test-app.test.ts b/packages/appkit/src/testing/tests/create-test-app.test.ts new file mode 100644 index 000000000..db86580c2 --- /dev/null +++ b/packages/appkit/src/testing/tests/create-test-app.test.ts @@ -0,0 +1,448 @@ +import type { IAppRouter, PluginManifest } from "shared"; +import { describe, expect, test } from "vitest"; + +import { getWorkspaceClient } from "../../context"; +import { Plugin, toPlugin } from "../../plugin"; +import type { WorkspaceClient } from "../../workspace-client"; +import { createTestApp } from "../create-test-app"; +import { getMockFn } from "../mock-workspace-client"; + +/** + * Coverage for the harness itself. Nothing here is mocked beyond the workspace + * client the harness installs: these boots bind real sockets and run the real + * Express stack, because that is the claim being tested. + */ + +/** Builds a manifest with the fields the loader validates. */ +function manifest( + name: string, + extra: Record = {}, +): PluginManifest { + return { + name, + displayName: name, + version: "0.0.0", + description: `${name} test plugin`, + resources: { required: [] }, + ...extra, + } as unknown as PluginManifest; +} + +/** Serves JSON, echoes bodies, and reports what it saw of the client. */ +class EchoPlugin extends Plugin { + static manifest = manifest("echo"); + + /** The client this plugin resolved at request time. */ + seenClient: WorkspaceClient | undefined; + + injectRoutes(router: IAppRouter): void { + router.get("/ping", async (_req, res) => { + res.json({ pong: true }); + }); + + router.post("/echo", async (req, res) => { + res.json({ + received: req.body, + contentType: req.headers["content-type"], + }); + }); + + router.get("/whoami", async (req, res) => { + res.json({ + user: req.headers["x-forwarded-user"] ?? null, + token: req.headers["x-forwarded-access-token"] ?? null, + custom: req.headers["x-custom"] ?? null, + }); + }); + + router.get("/from-client", async (_req, res) => { + // Reaches the data plane exactly the way a real plugin does. + const client = getWorkspaceClient(); + this.seenClient = client; + const run = await client.jobs.getRun({ run_id: 1 } as never); + res.json({ run }); + }); + + router.get("/client-identity", async (_req, res) => { + this.seenClient = getWorkspaceClient(); + res.json({ ok: true }); + }); + + router.get("/boom", async () => { + throw new Error("handler exploded"); + }); + + router.put("/put", async (req, res) => res.json({ m: "PUT", b: req.body })); + router.patch("/patch", async (req, res) => + res.json({ m: "PATCH", b: req.body }), + ); + router.delete("/del", async (_req, res) => res.json({ m: "DELETE" })); + } + + exports() { + return { seenClient: () => this.seenClient }; + } +} +const echo = toPlugin(EchoPlugin); + +/** Declares a required env var, so resource validation has something to fail on. */ +class NeedsEnvPlugin extends Plugin { + static manifest = manifest("needsEnv", { + resources: { + required: [ + { + type: "sql_warehouse", + alias: "Harness Probe Warehouse", + resourceKey: "harness-probe", + description: "Exists only so validation has something to fail on", + permission: "CAN_USE", + fields: { + id: { + env: "MY_REQUIRED_SECRET", + description: "Stand-in for a required resource field", + }, + }, + }, + ], + optional: [], + }, + }); +} +const needsEnv = toPlugin(NeedsEnvPlugin); + +/** Fails during setup, to exercise the boot-failure teardown path. */ +class BadSetupPlugin extends Plugin { + static manifest = manifest("badSetup"); + async setup(): Promise { + throw new Error("setup went wrong"); + } +} +const badSetup = toPlugin(BadSetupPlugin); + +describe("createTestApp", () => { + test("boots with a single plugin and serves a real route", async () => { + const app = await createTestApp({ plugins: [echo()] }); + try { + expect(app.baseUrl).toMatch(/^http:\/\/127\.0\.0\.1:\d+$/); + expect(app.port).toBeGreaterThan(0); + + const res = await app.get("/api/echo/ping"); + expect(res.status).toBe(200); + await expect(res.json()).resolves.toEqual({ pong: true }); + } finally { + await app.close(); + } + }); + + test("two apps in one file get different ephemeral ports", async () => { + const a = await createTestApp({ plugins: [echo()] }); + const b = await createTestApp({ plugins: [echo()] }); + try { + // No EADDRINUSE, which is what makes the harness parallel-safe and is why + // hardcoded test ports are worth removing. + expect(a.port).not.toBe(b.port); + await expect(a.get("/api/echo/ping").then((r) => r.status)).resolves.toBe( + 200, + ); + await expect(b.get("/api/echo/ping").then((r) => r.status)).resolves.toBe( + 200, + ); + } finally { + await a.close(); + await b.close(); + } + }); + + test("boots with no credentials in the environment", async () => { + const saved = { ...process.env }; + for (const key of Object.keys(process.env)) { + if (key.startsWith("DATABRICKS_")) delete process.env[key]; + } + try { + const app = await createTestApp({ plugins: [echo()] }); + try { + await expect( + app.get("/api/echo/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await app.close(); + } + } finally { + process.env = saved; + } + }); + + test("the default mock client reaches the plugin instead of crashing", async () => { + const app = await createTestApp({ plugins: [echo()] }); + try { + const res = await app.get("/api/echo/from-client"); + expect(res.status).toBe(200); + // Undeclared path, so it resolves undefined rather than throwing — the + // never-crash floor, exercised through a real handler. + await expect(res.json()).resolves.toEqual({}); + } finally { + await app.close(); + } + }); + + test("caller-supplied responses reach the plugin's client calls", async () => { + const app = await createTestApp({ + plugins: [echo()], + responses: { "jobs.getRun": { state: "TERMINATED" } }, + }); + try { + const res = await app.get("/api/echo/from-client"); + await expect(res.json()).resolves.toEqual({ + run: { state: "TERMINATED" }, + }); + expect(getMockFn(app.client, "jobs.getRun")).toHaveBeenCalledWith({ + run_id: 1, + }); + } finally { + await app.close(); + } + }); + + test("app.client is the same object a handler resolves", async () => { + const app = await createTestApp({ plugins: [echo()] }); + try { + await app.get("/api/echo/client-identity"); + // Retires the "tribal seam knowledge" problem: no need to know that + // createApp({ client }) flows through ServiceContext to reach a handler. + expect(app.plugins.echo.seenClient()).toBe(app.client); + } finally { + await app.close(); + } + }); + + test("apiClient.request has zero calls after boot", async () => { + const app = await createTestApp({ plugins: [echo()] }); + try { + // A canary for two hazards at once: DATABRICKS_WORKSPACE_ID must + // short-circuit the SCIM probe in getWorkspaceId, and internal telemetry + // must stay off. If either regresses, request assertions get polluted and + // this fails loudly. + expect(getMockFn(app.client, "apiClient.request")).toHaveBeenCalledTimes( + 0, + ); + } finally { + await app.close(); + } + }); + + test("a caller-supplied server plugin is respected, and dedupes the injected one", async () => { + const { server: serverPlugin } = await import("../../plugins/server"); + const app = await createTestApp({ + plugins: [echo(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }); + try { + expect(app.port).toBeGreaterThan(0); + await expect( + app.get("/api/echo/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await app.close(); + } + }); + + test("server: false together with a server plugin is refused", async () => { + const { server: serverPlugin } = await import("../../plugins/server"); + await expect( + createTestApp({ + plugins: [echo(), serverPlugin({ port: 0, host: "127.0.0.1" })], + server: false, + }), + ).rejects.toThrow(/conflicts with the server plugin/); + }); + + test("server: false boots without a socket and request methods explain why", async () => { + const app = await createTestApp({ plugins: [echo()], server: false }); + try { + expect(app.server).toBeUndefined(); + expect(() => app.baseUrl).toThrow(/no HTTP server/); + await expect(app.get("/api/echo/ping")).rejects.toThrow(/no HTTP server/); + } finally { + await app.close(); + } + }); + + test("await using releases at scope exit", async () => { + let port: number | undefined; + { + await using app = await createTestApp({ plugins: [echo()] }); + port = app.port; + await expect( + app.get("/api/echo/ping").then((r) => r.status), + ).resolves.toBe(200); + } + await expect(fetch(`http://127.0.0.1:${port}/health`)).rejects.toThrow(); + }); + + describe("resource validation (the strict posture)", () => { + test("a missing required env var fails the boot", async () => { + delete process.env.MY_REQUIRED_SECRET; + await expect(createTestApp({ plugins: [needsEnv()] })).rejects.toThrow( + /MY_REQUIRED_SECRET/, + ); + }); + + test("supplying it through env makes the same boot pass", async () => { + const app = await createTestApp({ + plugins: [needsEnv(), echo()], + env: { MY_REQUIRED_SECRET: "s3cret" }, + }); + try { + await expect( + app.get("/api/echo/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await app.close(); + } + // Restored, not leaked into the next test. + expect(process.env.MY_REQUIRED_SECRET).toBeUndefined(); + }); + + test("validation always throws, because the harness pins NODE_ENV", async () => { + delete process.env.MY_REQUIRED_SECRET; + + // enforceValidation computes `shouldThrow = !isDevelopment || strict`, so + // pinning NODE_ENV away from "development" is what makes the throw + // unconditional. There is intentionally no option to soften this: the + // warning path exists only in dev mode, which the harness refuses. + await expect( + createTestApp({ plugins: [needsEnv()], nodeEnv: "production" }), + ).rejects.toThrow(/Missing required resources/); + await expect( + createTestApp({ plugins: [needsEnv()], nodeEnv: "test" }), + ).rejects.toThrow(/Missing required resources/); + }); + }); + + describe("environment hygiene", () => { + test("close() restores the snapshot, including pre-existing values", async () => { + process.env.DATABRICKS_HOST = "https://original.example.com"; + const before = { ...process.env }; + + const app = await createTestApp({ + plugins: [echo()], + env: { HARNESS_ADDED: "yes" }, + }); + // The harness overwrote DATABRICKS_HOST with its test default. + expect(process.env.DATABRICKS_HOST).not.toBe( + "https://original.example.com", + ); + await app.close(); + + // A pre-existing value is restored to *its* value, not the test default, + // and a key the harness added is deleted rather than left behind. + expect(process.env.DATABRICKS_HOST).toBe("https://original.example.com"); + expect(process.env.HARNESS_ADDED).toBeUndefined(); + expect(Object.keys(process.env).sort()).toEqual( + Object.keys(before).sort(), + ); + + delete process.env.DATABRICKS_HOST; + }); + + test("a boot failure still restores env and resets singletons", async () => { + const before = { ...process.env }; + + await expect( + createTestApp({ plugins: [badSetup()], env: { LEAKED: "no" } }), + ).rejects.toThrow(/setup went wrong/); + + // Teardown has to run from the setup-failure path, or every later test in + // the file inherits the mutated env. + expect(process.env.LEAKED).toBeUndefined(); + expect(Object.keys(process.env).sort()).toEqual( + Object.keys(before).sort(), + ); + + // And the next boot still works. + const app = await createTestApp({ plugins: [echo()] }); + await expect( + app.get("/api/echo/ping").then((r) => r.status), + ).resolves.toBe(200); + await app.close(); + }); + + test('nodeEnv: "development" is refused with an explanation', async () => { + // The get-port RangeError must never reach the user. + await expect( + createTestApp({ plugins: [echo()], nodeEnv: "development" }), + ).rejects.toThrow(/not supported/); + }); + + test("SIGTERM listener count is unchanged across boot and close", async () => { + const baseline = process.listenerCount("SIGTERM"); + const app = await createTestApp({ plugins: [echo()] }); + await app.close(); + // Guards the MaxListenersExceededWarning that shows up at ~6 un-closed + // boots in one file. + expect(process.listenerCount("SIGTERM")).toBe(baseline); + }); + + test("boot, close, boot again in one file", async () => { + const first = await createTestApp({ plugins: [echo()] }); + const firstPort = first.port; + await first.close(); + + const second = await createTestApp({ plugins: [echo()] }); + try { + expect(second.port).not.toBe(firstPort); + await expect( + second.get("/api/echo/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await second.close(); + } + }); + + test("overlapping boots restore env regardless of close order", async () => { + const before = { ...process.env }; + + // The second boot's view of "original" already contains the first boot's + // mutations. A per-app snapshot would let whichever closes last re-apply + // them, stranding harness keys and `A_ONLY` after both apps are gone. + const a = await createTestApp({ + plugins: [echo()], + env: { OVERLAP_A: "a" }, + }); + const b = await createTestApp({ + plugins: [echo()], + env: { OVERLAP_B: "b" }, + }); + + await a.close(); + await b.close(); + + const leaked = Object.keys(process.env).filter((k) => !(k in before)); + expect(leaked).toEqual([]); + expect(process.env.OVERLAP_A).toBeUndefined(); + expect(process.env.OVERLAP_B).toBeUndefined(); + expect(Object.keys(process.env).sort()).toEqual( + Object.keys(before).sort(), + ); + }); + + test("closing in reverse order also restores env", async () => { + const before = { ...process.env }; + const a = await createTestApp({ plugins: [echo()], env: { REV_A: "a" } }); + const b = await createTestApp({ plugins: [echo()], env: { REV_B: "b" } }); + + // Reverse of boot order — the outcome must not depend on it. + await b.close(); + await a.close(); + + expect(Object.keys(process.env).filter((k) => !(k in before))).toEqual( + [], + ); + }); + + test("close() is idempotent", async () => { + const app = await createTestApp({ plugins: [echo()] }); + await app.close(); + await expect(app.close()).resolves.toBeUndefined(); + }); + }); +}); diff --git a/packages/appkit/src/testing/tests/create-test-plugin.test.ts b/packages/appkit/src/testing/tests/create-test-plugin.test.ts new file mode 100644 index 000000000..e92731677 --- /dev/null +++ b/packages/appkit/src/testing/tests/create-test-plugin.test.ts @@ -0,0 +1,82 @@ +import type { BasePluginConfig, PluginManifest } from "shared"; +import { describe, expect, test } from "vitest"; + +import { Plugin, toPlugin } from "../../plugin"; +import { createTestPlugin } from "../create-test-plugin"; + +/** + * The behaviour that matters is the merge: an instance built by hand skips + * DEFAULT_CONFIG and forgets `name`, so a test against it can pass wrongly. + */ + +interface WidgetConfig extends BasePluginConfig { + size?: string; + colour?: string; +} + +class WidgetPlugin extends Plugin { + static manifest = { + name: "widget", + displayName: "Widget", + version: "0.0.0", + description: "config-merge probe", + resources: { required: [], optional: [] }, + } as unknown as PluginManifest; + + static DEFAULT_CONFIG = { size: "medium", colour: "blue" }; + + readonly received: WidgetConfig; + + constructor(config: WidgetConfig) { + super(config); + this.received = config; + } +} +// No cast: the class satisfies PluginConstructor, so the factory's config and +// instance types both infer — which is what lets createTestPlugin be typed. +const widget = toPlugin(WidgetPlugin); + +describe("createTestPlugin", () => { + test("returns an instance of the plugin class", () => { + const plugin = createTestPlugin(widget); + expect(plugin).toBeInstanceOf(WidgetPlugin); + }); + + test("applies DEFAULT_CONFIG", () => { + const plugin = createTestPlugin(widget); + // The hand-rolled `new (widget({}).plugin)({})` skips these entirely. + expect(plugin.received.size).toBe("medium"); + expect(plugin.received.colour).toBe("blue"); + }); + + test("explicit config wins over DEFAULT_CONFIG", () => { + const plugin = createTestPlugin(widget, { + size: "large", + }); + expect(plugin.received.size).toBe("large"); + // Unspecified keys still come from the defaults. + expect(plugin.received.colour).toBe("blue"); + }); + + test("sets the manifest name, which the hand-rolled form forgets", () => { + const plugin = createTestPlugin(widget); + expect(plugin.received.name).toBe("widget"); + expect(plugin.name).toBe("widget"); + }); + + test("a zero-argument call works", () => { + expect(() => createTestPlugin(widget)).not.toThrow(); + }); + + test("the merge order matches what registration produces", () => { + // Same order as AppKit.createAndRegisterPlugin: DEFAULT_CONFIG, then the + // factory's config, then `name`. A caller cannot override `name`, because + // the manifest owns it. + const plugin = createTestPlugin(widget, { + name: "not-this", + colour: "red", + }); + expect(plugin.received.name).toBe("widget"); + expect(plugin.received.colour).toBe("red"); + }); +}); 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..fe055e316 --- /dev/null +++ b/packages/appkit/src/testing/tests/expect-stream.test.ts @@ -0,0 +1,285 @@ +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) { + 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", + ]); + }); + + 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. + 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"]); + }); + + // 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", () => { + test("throws for a non-stream value", async () => { + await expect( + // Intentionally wrong type to exercise the runtime guard. + 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", () => { + 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( + `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/fixtures.test.ts b/packages/appkit/src/testing/tests/fixtures.test.ts new file mode 100644 index 000000000..201a21d71 --- /dev/null +++ b/packages/appkit/src/testing/tests/fixtures.test.ts @@ -0,0 +1,126 @@ +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/mock-workspace-client.test.ts b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts new file mode 100644 index 000000000..c0622b6b5 --- /dev/null +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -0,0 +1,471 @@ +import { inspect } from "node:util"; + +import { describe, expect, test, vi } from "vitest"; + +import { ServiceContext } from "../../context/service-context"; +import { mockServiceContext } from "../fixtures"; +import { createMockWorkspaceClient, getMockFn } from "../mock-workspace-client"; + +describe("createMockWorkspaceClient", () => { + describe("happy path", () => { + test("all 9 facade accessors are reachable and callable without throwing", async () => { + const client = createMockWorkspaceClient(); + + // Assert all 9 explicitly reachable. + expect(client.files).toBeDefined(); + expect(client.warehouses).toBeDefined(); + expect(client.genie).toBeDefined(); + expect(client.jobs).toBeDefined(); + expect(client.statementExecution).toBeDefined(); + expect(client.servingEndpoints).toBeDefined(); + expect(client.currentUser).toBeDefined(); + expect(client.config).toBeDefined(); + expect(client.apiClient).toBeDefined(); + + // And callable without throwing (using as any since these are Proxy mocks). + await expect( + (client.files as any).listDirectory({ path: "/x" }), + ).resolves.toBe(undefined); + await expect(client.warehouses.get({ id: "123" })).resolves.toEqual({ + state: "RUNNING", + }); + await expect( + (client.genie as any).getMessage({ message_id: "xyz" }), + ).resolves.toBe(undefined); + await expect(client.jobs.getRun({ run_id: 1 })).resolves.toBe(undefined); + await expect( + client.statementExecution.executeStatement({ + warehouse_id: "w1", + catalog: "c", + schema: "s", + statement: "SELECT 1", + }), + ).resolves.toEqual({ + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }); + await expect( + (client.servingEndpoints as any).get({ name: "ep" }), + ).resolves.toBe(undefined); + await expect(client.currentUser.me()).resolves.toEqual({ + id: "test-service-user", + userName: "test-service-user", + }); + expect(typeof client.config.host).toBe("string"); + expect(typeof client.apiClient.userAgent?.()).toBe("string"); + }); + + test("a declared path returns its value", async () => { + const response = { state: "TERMINATED", result_state: "SUCCESS" }; + const client = createMockWorkspaceClient({ + responses: { "jobs.getRun": response }, + }); + const result = await client.jobs.getRun({ run_id: 123 }); + expect(result).toEqual(response); + }); + + test("a function-valued response receives call arguments", async () => { + const fn = vi.fn().mockResolvedValue({ called: true }); + const client = createMockWorkspaceClient({ + responses: { "jobs.getRun": fn }, + }); + const args = { run_id: 456 }; + await client.jobs.getRun(args); + expect(fn).toHaveBeenCalledWith(args); + }); + + test("a rejecting function propagates the error", async () => { + const error = new Error("test error"); + const client = createMockWorkspaceClient({ + responses: { "jobs.getRun": () => Promise.reject(error) }, + }); + await expect(client.jobs.getRun({ run_id: 789 })).rejects.toBe(error); + }); + + test("an undeclared path resolves undefined and does not throw", async () => { + const client = createMockWorkspaceClient(); + const result = await (client.genie as any).getMessage({ + message_id: "missing", + }); + expect(result).toBe(undefined); + }); + + test("depth-2 apiClient.request resolves from the key", async () => { + const response = { results: [{ value: "x" }] }; + const client = createMockWorkspaceClient({ + responses: { "apiClient.request": response }, + }); + const result = await (client.apiClient.request as any)({ + path: "/api/2.0/something", + }); + expect(result).toEqual(response); + }); + + test("a caller-supplied response overrides the default", async () => { + const customResponse = { + status: { state: "RUNNING" }, + result: { data: ["custom"] }, + }; + const client = createMockWorkspaceClient({ + responses: { + "statementExecution.executeStatement": customResponse, + }, + }); + const result = await client.statementExecution.executeStatement({ + warehouse_id: "w", + catalog: "c", + schema: "s", + statement: "SELECT 1", + }); + expect(result).toEqual(customResponse); + }); + + test("currentUser.me() resolves an object with a non-empty id", async () => { + const client = createMockWorkspaceClient(); + const user = await client.currentUser.me(); + expect(user).toBeDefined(); + expect(user?.id).toBeTruthy(); + expect(user?.userName).toBeTruthy(); + }); + }); + + describe("stable identity (memoization)", () => { + test("client.jobs.getRun === client.jobs.getRun across accesses", async () => { + const client = createMockWorkspaceClient(); + const fn1 = client.jobs.getRun; + const fn2 = client.jobs.getRun; + expect(fn1).toBe(fn2); + // toHaveBeenCalledWith should also work with the stable reference. + await fn1({ run_id: 1 }); + expect(fn1).toHaveBeenCalledWith({ run_id: 1 }); + }); + + test("client.jobs === client.jobs (namespace memoization)", () => { + const client = createMockWorkspaceClient(); + const jobs1 = client.jobs; + const jobs2 = client.jobs; + expect(jobs1).toBe(jobs2); + }); + + test("client.toLegacyWorkspaceClient().jobs.getRun === client.jobs.getRun", async () => { + const client = createMockWorkspaceClient({ + responses: { "jobs.getRun": { state: "COMPLETED" } }, + }); + const legacy = client.toLegacyWorkspaceClient(); + const clientFn = client.jobs.getRun; + const legacyFn = legacy.jobs.getRun; + expect(clientFn).toBe(legacyFn); + // Call it and verify it's tracked on both references. + await clientFn({ run_id: 1 }); + expect(legacyFn).toHaveBeenCalledWith({ run_id: 1 }); + }); + + test("un-faceted legacy services work (e.g., legacy.clusters.list())", async () => { + const client = createMockWorkspaceClient({ + responses: { "clusters.list": { clusters: [] } }, + }); + const legacy = client.toLegacyWorkspaceClient(); + const result = await (legacy as any).clusters.list(); + expect(result).toEqual({ clusters: [] }); + }); + }); + + describe("footguns (the highest-value tests)", () => { + test("client.jobs.then is undefined; await client.jobs resolves to the service itself", async () => { + const client = createMockWorkspaceClient(); + // Should not hang or resolve to a mock's return value. + const resolved = await client.jobs; + expect(resolved).toBe(client.jobs); + }); + + test("util.inspect renders the client without throwing or recursing", () => { + const client = createMockWorkspaceClient(); + + // Because the traps leave `ownKeys`/`getOwnPropertyDescriptor` at their + // defaults, a service proxy has no enumerable keys and inspects as `{}` + // instead of recursing forever minting a mock per probed property. + expect(inspect(client.jobs)).toBe("{}"); + expect(inspect(client.toLegacyWorkspaceClient())).toBe("{}"); + + // The facade itself is a plain object, so its nine members are listed — + // and `config.host` shows through as the real string it is. + const whole = inspect(client); + expect(whole).toContain("jobs: {}"); + expect(whole).toContain("https://test.databricks.com"); + }); + + test("console.log('%O', client) works without hanging or recursing", () => { + const client = createMockWorkspaceClient(); + // Stubbed only to keep the formatted dump out of the test output — the + // formatting still runs, which is what could throw or recurse. + const log = vi.spyOn(console, "log").mockImplementation(() => {}); + try { + expect(() => { + console.log("%O", client); + }).not.toThrow(); + expect(log).toHaveBeenCalledTimes(1); + } finally { + log.mockRestore(); + } + }); + + test("expect(client.jobs).toEqual({}) does not blow the stack", () => { + const client = createMockWorkspaceClient(); + // This is the `asymmetricMatch` guard test. + expect(() => { + expect(client.jobs).toEqual({}); + }).not.toThrow(); + }); + + test("JSON.stringify(client.config) does not throw", () => { + const client = createMockWorkspaceClient(); + expect(() => { + JSON.stringify(client.config); + }).not.toThrow(); + }); + }); + + describe("special cases", () => { + test("typeof client.config.host === 'string' and truthy", () => { + const client = createMockWorkspaceClient(); + const host = client.config.host; + expect(typeof host).toBe("string"); + expect(host).toBeTruthy(); + // Can build a URL from it. + expect(() => { + new URL("/x", host as string); + }).not.toThrow(); + }); + + test("responses['config.host'] returns the raw string, not a mock", async () => { + const host = "https://custom.databricks.com"; + const client = createMockWorkspaceClient({ + responses: { "config.host": host }, + }); + expect(client.config.host).toBe(host); + }); + + test("client.config.authenticate(new Headers()) sets an Authorization header", async () => { + const client = createMockWorkspaceClient(); + const headers = new Headers(); + + // Unconditional: `authenticate` is always seeded, so a guard here would + // let the whole assertion be skipped if it ever stopped being. + await client.config.authenticate(headers); + + // The side effect is the point — asserting only "was called" would pass + // against a bare vi.fn() that does nothing, which is what the header- + // stamping paths in AppKit actually depend on. + expect(headers.get("Authorization")).toBe("Bearer test-token"); + expect(getMockFn(client, "config.authenticate")).toHaveBeenCalledWith( + headers, + ); + }); + + test("client.config.ensureResolved() resolves", async () => { + const client = createMockWorkspaceClient(); + await expect(client.config.ensureResolved()).resolves.toBe(undefined); + }); + + test("typeof client.apiClient.userAgent() === 'string' (synchronous)", () => { + const client = createMockWorkspaceClient(); + const result = client.apiClient.userAgent?.(); + expect(typeof result).toBe("string"); + // Not a Promise (shouldn't have a .then method). + expect(typeof (result as any)?.then).not.toBe("function"); + }); + + test("await client.apiClient.request({}) resolves to an object", async () => { + const client = createMockWorkspaceClient(); + const result = await (client.apiClient.request as any)({ + path: "/api/2.0/test", + }); + // Should be an object, not undefined. + expect(result).toEqual({}); + }); + + test("client.config.someUnknownField returns a mock", async () => { + const client = createMockWorkspaceClient(); + const unknownField = (client.config as any).someUnknownField; + // Should be a mock (vi.fn). + expect(typeof unknownField).toBe("function"); + if (typeof unknownField === "function" && (unknownField as any).mock) { + expect((unknownField as any).mock).toBeDefined(); + } + }); + + test("{ defaults: false } leaves statementExecution.executeStatement unresolved", async () => { + const client = createMockWorkspaceClient({ defaults: false }); + const result = await client.statementExecution.executeStatement({ + warehouse_id: "w", + catalog: "c", + schema: "s", + statement: "SELECT 1", + }); + expect(result).toBe(undefined); + }); + }); + + describe("getMockFn escape hatch", () => { + test("getMockFn retrieves the cached mock for a dotted path", async () => { + const client = createMockWorkspaceClient(); + await client.jobs.getRun({ run_id: 123 }); + + const mock = getMockFn(client, "jobs.getRun"); + expect(mock.mock).toBeDefined(); + expect(mock).toHaveBeenCalledWith({ run_id: 123 }); + }); + + test("getMockFn mints before first use, so it can be grabbed up front", async () => { + const client = createMockWorkspaceClient(); + + // Grabbing the handle before the code under test runs must yield the very + // function that code will call — otherwise every assertion would have to + // be written after the fact. + const getRun = getMockFn(client, "jobs.getRun"); + expect(getRun).toHaveBeenCalledTimes(0); + + await client.jobs.getRun({ run_id: 7 }); + + expect(getRun).toBe(getMockFn(client, "jobs.getRun")); + expect(getRun).toHaveBeenCalledWith({ run_id: 7 }); + }); + + test("getMockFn resolves seeded members and rejects non-function paths", () => { + const client = createMockWorkspaceClient(); + + // Seeded on the apiClient object rather than minted by the trap. + expect(getMockFn(client, "apiClient.request")).toBe( + client.apiClient.request, + ); + + // config.host is a real string, so there is no mock to hand back. + expect(() => getMockFn(client, "config.host")).toThrow( + /not a mocked function/, + ); + + expect(() => getMockFn({} as never, "jobs.getRun")).toThrow( + /not a createMockWorkspaceClient/, + ); + }); + }); + + describe("configuration override", () => { + test("the config option overrides defaults and adds members", () => { + const customAuth = vi.fn(); + const client = createMockWorkspaceClient({ + config: { + host: "https://custom.databricks.com", + authenticate: customAuth, + }, + }); + expect(client.config.host).toBe("https://custom.databricks.com"); + expect(client.config.authenticate).toBe(customAuth); + }); + }); +}); + +/** + * Compile-time contract. These assertions are enforced by `tsc --noEmit` + * (`pnpm --filter=@databricks/appkit typecheck`), not at runtime: a + * `@ts-expect-error` that stops being an error fails the typecheck, which is + * what guards the typed floor. The block runs as a test only so an accidental + * runtime throw is still caught. + */ +describe("compile-time contract", () => { + test("the typed facade rejects unknown members and keeps host a string", () => { + const client = createMockWorkspaceClient(); + + // A misspelled *service* is a compile error — this is what the typed + // 9-member floor buys over an untyped Proxy. + // @ts-expect-error - `jbos` is not a facade member + expect(client.jbos).toBeUndefined(); + + // And so is a misspelled *method*, because each accessor is typed against + // the SDK's own service class. The runtime floor is a fallback for calls + // that bypass the types, not the first line of defence — verified against a + // packed tarball from outside the monorepo. + // @ts-expect-error - `getRunz` is not a jobs method + void client.jobs.getRunz; + // @ts-expect-error - `getMessagez` is not a genie method + void client.genie.getMessagez; + // @ts-expect-error - `anything` is not a files method + void client.files.anything; + + // `config.host` is typed `string | undefined` by the SDK (production code + // guards it — see connectors/files/client.ts, which throws when falsy), so + // the honest compile-time claim is that it narrows to a *string*, not that + // it is non-optional. If the fake ever regressed to handing back a mock, + // this narrowing would not compile and `.startsWith` would not exist. + const host = client.config.host; + expect(typeof host).toBe("string"); + if (typeof host === "string") { + expect(host.startsWith("https://")).toBe(true); + } + + // The mock handle is a real `Mock`, so the mock API typechecks. + const getRun = getMockFn(client, "jobs.getRun"); + getRun.mockResolvedValue({ state: "TERMINATED" }); + expect(getRun.mock.calls).toEqual([]); + }); +}); + +/** + * The convergence guard. `mockServiceContext` hands this client to 13 test + * files that never name it — they just call `mockServiceContext()` and let the + * default client through. These assertions are what prove pointing that default + * at the new builder is a fix rather than a break. + */ +describe("convergence with mockServiceContext (D4)", () => { + test("the historical canned defaults are unchanged", async () => { + const client = createMockWorkspaceClient(); + + // Byte-identical to the shape the old two-service fixture returned. + await expect( + client.statementExecution.executeStatement({} as never), + ).resolves.toEqual({ + status: { state: "SUCCEEDED" }, + result: { data: [] }, + }); + await expect(client.warehouses.get({} as never)).resolves.toEqual({ + state: "RUNNING", + }); + await expect(client.warehouses.start({} as never)).resolves.toBeUndefined(); + }); + + test("the default client from mockServiceContext no longer crashes on jobs", async () => { + const mock = mockServiceContext(); + try { + const client = mock.serviceContext.client; + + // The whole point of U1+U2: before convergence this threw + // "Cannot read properties of undefined (reading 'getRun')", because the + // default client only had statementExecution and warehouses. + await expect(client.jobs.getRun({ run_id: 1 })).resolves.toBeUndefined(); + await expect( + ( + client.genie as never as Record Promise> + ).getMessage(), + ).resolves.toBeUndefined(); + + // ...while the SQL path 13 files depend on still succeeds. + await expect( + client.statementExecution.executeStatement({} as never), + ).resolves.toMatchObject({ status: { state: "SUCCEEDED" } }); + } finally { + mock.restore(); + } + }); + + test("the user-context client is faked too, not just the service one", async () => { + const mock = mockServiceContext(); + try { + const userCtx = ServiceContext.createUserContext("tok", "u-1", "alice"); + await expect( + userCtx.client.jobs.getRun({ run_id: 1 }), + ).resolves.toBeUndefined(); + } finally { + mock.restore(); + } + }); +}); diff --git a/packages/appkit/src/testing/tests/published-surface.integration.test.ts b/packages/appkit/src/testing/tests/published-surface.integration.test.ts new file mode 100644 index 000000000..baa65998d --- /dev/null +++ b/packages/appkit/src/testing/tests/published-surface.integration.test.ts @@ -0,0 +1,90 @@ +import { + createTestApp, + expectStream, + getMockFn, +} from "@databricks/appkit/testing"; +import { describe, expect, test } from "vitest"; + +import { Plugin, toPlugin } from "../../plugin"; + +/** + * Acceptance test for the published surface: everything the test needs comes from + * `@databricks/appkit/testing` — no `@tools` shim, no deep imports. + * `Plugin`/`toPlugin` come from the main entry because they are how you *write* a + * plugin, not how you test one. + */ + +class WidgetPlugin extends Plugin { + static manifest = { + name: "widget", + displayName: "Widget", + version: "0.0.0", + description: "A plugin an external author might write", + resources: { required: [], optional: [] }, + } as never; + + injectRoutes(router: never): void { + this.route(router, { + name: "run", + method: "post", + path: "/run", + handler: async (req, res) => { + // The data plane, faked by the harness with no workspace in sight. + const { getWorkspaceClient } = await import("../../context"); + const run = await getWorkspaceClient().jobs.getRun({ + run_id: (req.body as { id: number }).id, + } as never); + res.json({ run }); + }, + }); + + this.route(router, { + name: "stream", + method: "post", + path: "/stream", + handler: async (_req, res) => { + res.setHeader("content-type", "text/event-stream"); + res.write(`event: status\ndata: ${JSON.stringify({ s: "go" })}\n\n`); + res.write(`event: result\ndata: ${JSON.stringify({ n: 1 })}\n\n`); + res.end(); + }, + }); + } +} +const widget = toPlugin(WidgetPlugin); + +describe("@databricks/appkit/testing as a standalone surface", () => { + test("boot, request, assert a stream, and close — public imports only", async () => { + const app = await createTestApp({ + plugins: [widget()], + responses: { "jobs.getRun": { state: "TERMINATED" } }, + }); + + try { + const res = await app.post("/api/widget/run", { body: { id: 42 } }); + expect(res.status).toBe(200); + await expect(res.json()).resolves.toEqual({ + run: { state: "TERMINATED" }, + }); + expect(getMockFn(app.client, "jobs.getRun")).toHaveBeenCalledWith({ + run_id: 42, + }); + + const stream = await app.post("/api/widget/stream"); + await expectStream(stream).toEmit("status", "result"); + } finally { + await app.close(); + } + }); + + test("await using works from the public entry too", async () => { + let port: number | undefined; + { + await using app = await createTestApp({ plugins: [widget()] }); + port = app.port; + const res = await app.post("/api/widget/run", { body: { id: 1 } }); + expect(res.status).toBe(200); + } + await expect(fetch(`http://127.0.0.1:${port}/health`)).rejects.toThrow(); + }); +}); diff --git a/packages/appkit/src/testing/tests/test-plugin-context.test.ts b/packages/appkit/src/testing/tests/test-plugin-context.test.ts new file mode 100644 index 000000000..60b30b917 --- /dev/null +++ b/packages/appkit/src/testing/tests/test-plugin-context.test.ts @@ -0,0 +1,329 @@ +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 { createTestPluginContext } from "../test-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 `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. + */ + +// 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, + header: (name: string) => headers[name.toLowerCase()], + } as unknown as express.Request; +} + +describe("createTestPluginContext — construction", () => { + test("produces a real PluginContext instance", () => { + const { ctx } = createTestPluginContext(); + expect(ctx).toBeInstanceOf(PluginContext); + }); + + test("registers fake providers passed at construction", () => { + const { ctx } = createTestPluginContext({ + analytics: { query: [{ id: 1 }] }, + genie: { ask: "hi" }, + }); + const names = ctx.getToolProviders().map((p) => p.name); + expect(names).toContain("analytics"); + expect(names).toContain("genie"); + }); +}); + +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 = createTestPluginContext({ analytics: { top_users: rows } }); + + const result = await mock.ctx.executeTool( + mockReq(), + "analytics", + "top_users", + { limit: 10 }, + ); + + expect(result).toEqual(rows); + // 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 = createTestPluginContext({ 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 = createTestPluginContext({ 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 = createTestPluginContext({ 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 = createTestPluginContext({ + 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 = createTestPluginContext({ + 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 = 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 = createTestPluginContext({ analytics: { query: [] } }); + await expect( + 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 = createTestPluginContext({ analytics: { lookup: null } }); + const result = await mock.ctx.executeTool( + mockReq(), + "analytics", + "lookup", + {}, + ); + 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", () => { + test("records a span on the injected mock telemetry for each executeTool", async () => { + const mock = createTestPluginContext({ 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("createTestPluginContext — route recording", () => { + test("records addRoute calls with raw (pre-wrap) handlers", () => { + const mock = createTestPluginContext(); + 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 = createTestPluginContext(); + const mw: express.RequestHandler = (_req, _res, next) => next(); + mock.ctx.addMiddleware("/api", mw); + expect(mock.routes).toEqual([ + { method: "use", path: "/api", handlers: [mw] }, + ]); + }); +}); + +describe("createTestPluginContext — registerProvider after construction", () => { + test("adds a provider dynamically", async () => { + const mock = createTestPluginContext(); + mock.registerProvider("late", { ping: "pong" }); + const result = await mock.ctx.executeTool(mockReq(), "late", "ping", {}); + expect(result).toBe("pong"); + }); +}); + +describe("createTestPluginContext — attach()", () => { + test("seeds the cache, flips isReady, and registers the plugin", async () => { + const mock = createTestPluginContext(); + 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 = createTestPluginContext({ probe: { canned: "fake" } }); + const plugin = new ProbePlugin({}); + await mock.attach(plugin); + + const result = await mock.ctx.executeTool(mockReq(), "probe", "canned", {}); + expect(result).toBe("fake"); + }); +}); diff --git a/packages/appkit/tsconfig.json b/packages/appkit/tsconfig.json index 5265a6881..76212e07e 100644 --- a/packages/appkit/tsconfig.json +++ b/packages/appkit/tsconfig.json @@ -7,7 +7,9 @@ "@/*": ["src/*"], "@tools/*": ["../../tools/*"], "shared": ["../../packages/shared/src"], - "@databricks/lakebase": ["../../packages/lakebase/src"] + "@databricks/lakebase": ["../../packages/lakebase/src"], + "@databricks/appkit": ["src/index.ts"], + "@databricks/appkit/testing": ["src/testing/index.ts"] } }, "include": ["src/**/*"], 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/packages/shared/src/plugin.ts b/packages/shared/src/plugin.ts index 15148895e..f04278187 100644 --- a/packages/shared/src/plugin.ts +++ b/packages/shared/src/plugin.ts @@ -264,6 +264,29 @@ export type PluginMap< >; }; +/** + * What `createApp()` returns: every plugin's exports keyed by manifest name, + * plus the app's own teardown handle. + * + * `close()` releases what AppKit acquired — sockets, timers, pools, cache, and + * telemetry — without terminating the process, so a host can embed AppKit and a + * test can boot more than once in a file. + * + * `Symbol.asyncDispose` is exposed alongside it because a plugin's manifest name + * can never be a symbol: `await using app = await createApp(...)` is safe even + * if a plugin were somehow named `close`. + */ +export type AppHandle< + U extends readonly PluginData[], +> = PluginMap & { + /** + * @param options.timeoutMs - Overall teardown budget. Defaults to AppKit's + * programmatic budget, which is shorter than the signal path's. + */ + close(options?: { timeoutMs?: number }): Promise; + [Symbol.asyncDispose](): Promise; +}; + /** Tuple of plugin class, config, and name. Created by `toPlugin()` and passed to `createApp()`. */ export type PluginData = { plugin: T; config: U; name: N }; /** Factory function type returned by `toPlugin()`. Accepts optional config and returns a PluginData tuple. */ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d211162d6..aec4e33c1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -376,6 +376,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: @@ -7731,7 +7734,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: @@ -18311,6 +18314,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 @@ -25523,6 +25534,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 @@ -25568,6 +25600,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 @@ -25611,6 +25660,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: diff --git a/template/server/example.test.ts b/template/server/example.test.ts new file mode 100644 index 000000000..535f800fd --- /dev/null +++ b/template/server/example.test.ts @@ -0,0 +1,100 @@ +import { Plugin, type PluginManifest, toPlugin } from '@databricks/appkit'; +import { createTestApp, createTestPluginContext, expectStream } 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. + * + * Three headline helpers are shown below: + * - `createTestApp({ plugins })` — boot a real app (real Express, real routes, + * real validation) on an ephemeral port and call it over HTTP. Start here for + * a plugin's end-to-end behaviour. Every boot needs `close()`. + * - `createTestPluginContext()` — a real PluginContext with faked edges, attachable + * to a plugin so its real code paths (routes, tool dispatch, user scoping) + * run under test. No boot, no socket — the fastest option for unit tests. + * - `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. +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 createTestPluginContext().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}!` }; + } + + // A real HTTP route, so createTestApp has something to call. + injectRoutes(router: Parameters[0]) { + this.route(router, { + name: 'greet', + method: 'post', + path: '/greet', + handler: async (req, res) => { + const { name } = req.body as { name: string }; + res.json({ message: `Hello, ${name}!` }); + }, + }); + } +} + +// The factory form `createApp` (and `createTestApp`) take. `toPlugin` reads the +// plugin name from the static manifest. +const greeter = toPlugin(GreeterPlugin); + +describe('testing kit example', () => { + test('attaches a real PluginContext and records registered routes', async () => { + const mock = createTestPluginContext(); + const plugin = new GreeterPlugin({}); + + 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 = new GreeterPlugin({}); + + await expectStream(plugin.greet('world')).toEmit('greeting_start', 'greeting_end'); + }); + + test('boots a real app and calls the plugin over HTTP', async () => { + // No workspace, no credentials, no network. The harness fakes the whole + // Databricks data plane and binds an ephemeral port. + const app = await createTestApp({ plugins: [greeter()] }); + + try { + const res = await app.post('/api/greeter/greet', { body: { name: 'world' } }); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ message: 'Hello, world!' }); + } finally { + // Required: releases the socket and restores process.env. + await app.close(); + } + }); +}); diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index 7f75cff1c..191df9470 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -1,447 +1,42 @@ -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()}`, - }; -} + * The integration suites have already moved to the public entry point, which is + * what verifies the published surface is self-sufficient. The remaining + * importers are unit suites, migrated opportunistically. + * + * 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, + createMockTelemetry, + createMockWorkspaceClient, + createSuccessfulSQLResponse, + createTestPluginContext, + expectStream, + mockServiceContext, + createTestApp, + type CreateTestAppOptions, + getListeningPort, + getMockFn, + type MockWorkspaceClient, + parseSSEResponse, + resetAppKitSingletons, + runWithRequestContext, + setupDatabricksEnv, + type TestApp, + type TestContextOptions, + type TestRequestOptions, + useServiceContextMock, +} from "../packages/appkit/src/testing";