diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fc0fb6272..379d83951 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -213,6 +213,16 @@ jobs: working-directory: pr-template run: npm install + # The scaffold's own vitest suite, run against the tarballs built above + # rather than the published versions template/package.json pins. This is + # the only CI check that exercises the template's example test, and the + # only one that consumes @databricks/appkit/testing the way a customer + # does — through a real npm install of a packed tarball. Before the zip, + # so a broken example fails the build instead of shipping. + - name: Run the template's own tests + working-directory: pr-template + run: npm test + # npm install above runs under the JFrog .npmrc (setup-jfrog-npm), which # bakes internal registry URLs into the regenerated lock. Rewrite them back # to public npm and fail-closed if any non-public registry remains, so the diff --git a/.gitignore b/.gitignore index 835aa9ef3..b70e0c2e8 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,10 @@ tmp node_modules .env +# Staging dir and zip produced by tools/prepare-template-artifact.ts +pr-template +appkit-template-*.zip + coverage *.tsbuildinfo @@ -16,3 +20,5 @@ coverage internal .isaac/ + +.codex-tmp/ diff --git a/.oxlintrc.json b/.oxlintrc.json index e45ab67fc..3e66f570e 100644 --- a/.oxlintrc.json +++ b/.oxlintrc.json @@ -26,11 +26,8 @@ { "patterns": [ { - "group": [ - "@databricks/sdk-experimental", - "@databricks/sdk-experimental/**" - ], - "message": "Import the Databricks SDK only through the wrapper in packages/shared/src/workspace-client. Add a re-export there if you need a new symbol." + "group": ["@databricks/sdk-*", "@databricks/sdk-*/**"], + "message": "Import the Databricks SDK only through the wrapper in packages/shared/src/workspace-client (legacy.ts for @databricks/sdk-experimental, modular.ts for the modular @databricks/sdk-* packages). Add a re-export there if you need a new symbol." } ] } diff --git a/docs/docs/api/appkit/Function.createApp.md b/docs/docs/api/appkit/Function.createApp.md index bc656537d..b7cd2fc3a 100644 --- a/docs/docs/api/appkit/Function.createApp.md +++ b/docs/docs/api/appkit/Function.createApp.md @@ -29,15 +29,15 @@ with an `asUser(req)` method for user-scoped execution. ## Parameters -| Parameter | Type | -| ------ | ------ | -| `config` | \{ `cache?`: [`CacheConfig`](Interface.CacheConfig.md); `client?`: [`WorkspaceClient`](Interface.WorkspaceClient.md); `disableInternalTelemetry?`: `boolean`; `onPluginsReady?`: (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\>; `plugins?`: `T`; `telemetry?`: [`TelemetryConfig`](Interface.TelemetryConfig.md); \} | -| `config.cache?` | [`CacheConfig`](Interface.CacheConfig.md) | -| `config.client?` | [`WorkspaceClient`](Interface.WorkspaceClient.md) | -| `config.disableInternalTelemetry?` | `boolean` | -| `config.onPluginsReady?` | (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\> | -| `config.plugins?` | `T` | -| `config.telemetry?` | [`TelemetryConfig`](Interface.TelemetryConfig.md) | +| Parameter | Type | Description | +| ------ | ------ | ------ | +| `config` | \{ `cache?`: [`CacheConfig`](Interface.CacheConfig.md); `client?`: [`WorkspaceClient`](Interface.WorkspaceClient.md); `disableInternalTelemetry?`: `boolean`; `onPluginsReady?`: (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\>; `plugins?`: `T`; `telemetry?`: [`TelemetryConfig`](Interface.TelemetryConfig.md); \} | - | +| `config.cache?` | [`CacheConfig`](Interface.CacheConfig.md) | - | +| `config.client?` | [`WorkspaceClient`](Interface.WorkspaceClient.md) | - | +| `config.disableInternalTelemetry?` | `boolean` | - | +| `config.onPluginsReady?` | (`appkit`: `PluginMap`\<`T`\>) => `void` \| `Promise`\<`void`\> | Runs after plugin setup but **before** the server starts. | +| `config.plugins?` | `T` | - | +| `config.telemetry?` | [`TelemetryConfig`](Interface.TelemetryConfig.md) | - | ## Returns diff --git a/docs/docs/api/appkit/Interface.WorkspaceClient.md b/docs/docs/api/appkit/Interface.WorkspaceClient.md index bf508bbe4..26a680581 100644 --- a/docs/docs/api/appkit/Interface.WorkspaceClient.md +++ b/docs/docs/api/appkit/Interface.WorkspaceClient.md @@ -86,20 +86,20 @@ Serving Endpoints. ### statementExecution ```ts -readonly statementExecution: StatementExecutionService; +readonly statementExecution: StatementExecutionClient; ``` -Statement Execution. +Statement Execution (modular SDK). *** ### warehouses ```ts -readonly warehouses: WarehousesService; +readonly warehouses: WarehousesClient; ``` -SQL Warehouses. +SQL Warehouses (modular SDK). ## Methods diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md index 30fc66642..6d06d355b 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -1,26 +1,292 @@ --- -sidebar_position: 8 +sidebar_position: 10 --- # 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. +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. Plugin tests stay fast and 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. +Exercise a plugin's real code paths against a real `PluginContext` with only its outer edges faked. That covers route registration, cross-plugin tool dispatch, user-scoped (on-behalf-of) execution, and per-call timeouts. Nothing about the context is reimplemented, so a test can't drift from production behavior. -The kit has two entry points plus a set of fixture helpers: +The kit has three entry points plus a set of fixture helpers: -- **`createTestPluginContext()`** — build a real `PluginContext` with faked edges and attach it to a plugin. +- **`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`, `mockServiceContext`, and SQL response builders. +- **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. +`vitest` is an **optional peer dependency**: the kit's mocks use its `vi` and resolve against your installed copy. Apps that never import `@databricks/appkit/testing` don't install it, so production stays free of the test framework. Any Vitest v3 or v4 works. + +## Testing your plugin + +`createTestApp({ plugins })` boots a **real** AppKit app, with the real Express wiring, routes, and resource validation, then 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 behavior 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 behavior or reject to test an error path. `responses` configures the built-in mock, so passing it alongside your own `client` is rejected rather than silently ignored — configure the responses on that client instead. Any path you **don't** declare resolves `undefined` rather than crashing — see [Mocking Databricks services](#mocking-databricks-services). + +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. + +With one app open, `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 { getMock } from "@databricks/appkit/testing"; + +expect(getMock(app.client, "jobs.getRun")).toHaveBeenCalledWith({ run_id: 42 }); +``` + +Facade accessors are typed against the SDK, so `expect(app.client.jobs.getRun).toHaveBeenCalled()` won't typecheck — `getMock` reaches the underlying spy. + +### 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, so **every boot needs a `close()`**. It 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. + +Prefer `await using`, which closes the app at scope exit even if the test throws: + +```ts +await using app = await createTestApp({ plugins: [myPlugin()] }); +// released at scope exit +``` + +`try/finally` works too, and is what you need if the app has to outlive a block: + +```ts +const app = await createTestApp({ plugins: [myPlugin()] }); +try { + // ... +} finally { + await app.close(); +} +``` + +Miss the close and the app stays live — socket bound, singletons and `process.env` not restored — so the next `createTestApp` is refused (one app at a time). + +For a suite where **every** test needs its own app, `useTestApp()` wires both hooks for you — a fresh app before each test, closed after — so there is no `close()` to forget: + +```ts +import { useTestApp } from "@databricks/appkit/testing"; + +describe("my plugin over HTTP", () => { + const app = useTestApp({ plugins: [myPlugin()] }); + + test("answers a request", async () => { + const res = await app.current.post("/api/my-plugin/run", { body: { id: 1 } }); + expect(res.status).toBe(200); + }); +}); +``` + +Call it at the top of a `describe`, not inside a test — Vitest registers `beforeEach`/`afterEach` during collection. Read `.current` from within a test; outside one it throws rather than handing back a closed app. `await using` stays the shorter choice for a single test, but it cannot carry an app from a `beforeEach` into the test body. + +### 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. Override it only when reaching the network is the point of the test. + +## Testing your plugin + +`createTestApp({ plugins })` boots a **real** AppKit app, with the real Express wiring, routes, and resource validation, then 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 behavior 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 behavior or reject to test an error path. `responses` configures the built-in mock, so passing it alongside your own `client` is rejected rather than silently ignored — configure the responses on that client instead. Any path you **don't** declare resolves `undefined` rather than crashing — see [Mocking Databricks services](#mocking-databricks-services) for the trade-off it makes. + +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. + +With one app open, `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 { getMock } from "@databricks/appkit/testing"; + +expect(getMock(app.client, "jobs.getRun")).toHaveBeenCalledWith({ run_id: 42 }); +``` + +`getMock` 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, so **every boot needs a `close()`**. It 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. + +Prefer `await using`, which closes the app at scope exit even if the test throws: + +```ts +await using app = await createTestApp({ plugins: [myPlugin()] }); +// released at scope exit +``` + +`try/finally` works too, and is what you need if the app has to outlive a block: + +```ts +const app = await createTestApp({ plugins: [myPlugin()] }); +try { + // ... +} finally { + await app.close(); +} +``` + +Miss the close and the app stays live — socket bound, singletons and `process.env` not restored — so the next `createTestApp` is refused (one app at a time). + +### 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. ## `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: +`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 | | --- | --- | @@ -28,7 +294,7 @@ The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so `vitest` is a | 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. +The context is real, so `executeTool` runs the actual user-scope (`asUser(req)`) and timeout-composition paths — not stubs of them. ### Registering fake tool responses @@ -58,7 +324,34 @@ 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()`: +### Seeding with workspace responses and environment + +`createTestPluginContext` accepts a second `options` parameter to control the faked workspace client and environment: + +```ts +const mock = createTestPluginContext({}, { + responses: { + "jobs.getRun": { state: "TERMINATED" }, + "servingEndpoints.query": (args, signal) => runFakeQuery(args), + }, + env: { MY_VAR: "test-value" }, + strict: true, +}); + +// The factory call is synchronous; attach is the async part. +await mock.attach(plugin); +``` + +`options` is: +- `responses` — seed the mock workspace client with responses keyed by dotted path (`"jobs.getRun"`, `"genie.getMessage"`). A value can be static or a function of call arguments and the abort signal. +- `env` — set environment variables scoped to the test; they are restored on plugin detach. +- `strict` — throw if a handler calls an undeclared workspace-client path (instead of silently resolving `undefined`). The built-in defaults still count as declared. + +The context installs a test-scoped service context via `beforeEach` and restores it on `afterEach`, so it survives across tests in the same suite. Call the returned `.restore()` explicitly if you need to clear it mid-test. + +The workspace client and on-behalf-of stub are process-wide too: `ServiceContext` holds one client, and the `createUserContext` fake is a single spy. So **`createTestApp` allows one open app at a time** and throws if you boot a second before closing the first. Vitest isolates test *files* in separate workers, so this constrains only apps within one file — and a `describe` holding an app open in `beforeAll` can't contain a test that boots its own. + +The cache is a process-wide singleton too — initialized once per test process and shared by tests **within one file** (it never leaks across files). If one test populates it and a later one must not see that, clear between tests with `resetTestCache()`: ```ts import { resetTestCache } from "@databricks/appkit/testing"; @@ -70,6 +363,32 @@ beforeEach(async () => { It also helps *within* a single test — clear the cache to force a miss, then assert the following call is a hit. +### Asserting cache behaviour + +When a plugin caches its work (like `analytics` caching query results), test the caching *itself* — a second identical call is a hit, different users get different keys — with `useTestCache()`. It boots the real in-memory cache, clears it before each test, and hands back the real `CacheManager`, so you assert against production's own `getOrExecute` and `generateKey` rather than mocking the internal `cache` module: + +```ts +import { useTestCache } from "@databricks/appkit/testing"; + +describe("my plugin caches", () => { + const testCache = useTestCache(); + + test("a second identical request is served from cache", async () => { + const plugin = new MyPlugin(config); + // ...drive the same request twice against a mocked downstream call... + expect(downstreamMock).toHaveBeenCalledTimes(1); + }); + + test("scopes the cache key per user", () => { + const a = testCache.current.generateKey(["query", sql], "user-1"); + const b = testCache.current.generateKey(["query", sql], "user-2"); + expect(a).not.toBe(b); + }); +}); +``` + +Call it at the top of a `describe` (or module top-level), not inside a test — Vitest registers its `beforeEach`/`afterEach` at collection time. It boots the cache before each test, so a plugin you construct binds `this.cache` to the real cache and runs its actual caching path. Use `resetTestCache()` (above) instead when you only need to clear the cache, not a handle to it. + ### Inspecting what happened The returned object exposes live views you read after the action under test runs: @@ -96,7 +415,7 @@ 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. +Assert cross-plugin on-behalf-of through `RecordedToolCall.asUser`. The fake `asUser` enforces the real `Plugin.asUser`'s token precondition: a request carrying a forwarded token records `asUser: true` with the resolved `userId`, and one missing `x-forwarded-access-token` **rejects**. Assert both directions — a well-formed request records the expected `userId`, a token-less one throws. A silent `{ executeTool }` stub verifies neither. 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()`. @@ -131,7 +450,7 @@ await plugin._handleStream(createMockRequest({ obo: true }), res); 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. +`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. @@ -143,9 +462,14 @@ 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, handling routes, tool dispatch, and user scoping; `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 covers both. `createTestApp` fakes the data plane 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.) +- `createMockRouter()` — build a mock Express-style router for testing route-registration wiring. - `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 @@ -159,11 +483,82 @@ The kit re-exports the request/response/context fixtures AppKit uses internally: ``` - `createSuccessfulSQLResponse(rows, columns)` / `createFailedSQLResponse(message)` — build SQL Warehouse statement responses. - `setupDatabricksEnv(overrides?)` — set `DATABRICKS_HOST` / `DATABRICKS_WAREHOUSE_ID` to test values. +- `withEnv(vars, fn)` — set environment variables for the duration of a sync or async function, restoring each key's prior state (or deleting it if it was previously unset). Unlike a bare `process.env.X = ...` followed by `delete`, nested calls restore LIFO and don't accidentally leave prior values in place. + ```ts + // Before: process.env.X = "test"; try { /* code */ } finally { delete process.env.X } + // After: + await withEnv({ X: "test" }, async () => { /* code */ }); + ``` +- `createApiError({ statusCode, message, errorCode })` — create a genuine `ApiError` instance for testing error paths. Returns an instance where `error instanceof ApiError` holds, so your error handling resolves the right type. + ```ts + const error = createApiError({ statusCode: 404, message: "Not found", errorCode: "NOT_FOUND" }); + expect(error instanceof ApiError).toBe(true); + ``` - `resetTestCache()` — clear the shared cache singleton between (or within) tests; no-ops if the cache isn't initialized yet. +The kit uses both words deliberately: a **mock** records calls so you can assert on them (`createMockWorkspaceClient`, `mockServiceContext`), while a **fake** stands in and simply works (`FakeProvider`, `FakeToolResponse`). + +- `createTestPlugin(factory, config?)` — instantiate a plugin from its factory with the same config merge AppKit applies. See [Full example](#full-example). +- `getListeningPort(server)` — wait for a server to finish binding and return the port it landed on. `createTestApp` does this for you; reach for it when you start a server yourself with `port: 0`. + +## 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, getMock } 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 the 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 Undeclared methods return undefined +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. + +Pass `strict: true` to turn that silence into a failure: a call to a path with no declared response throws instead of resolving `undefined`, naming the path. The canned defaults still count as declared, so a harness boot works unchanged. + +```ts +const app = await createTestApp({ plugins: [myPlugin()], strict: true }); +// a handler calling an undeclared path now fails the request +``` + +TypeScript catches more than the obvious: each accessor is typed against the SDK's own service class, so 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: 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 in production. That's deliberate — reporting the keys would make `util.inspect` probe each one, minting a mock per probe. + +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 -Instantiate the plugin **class** directly with `new`. The `analytics()` / `agents()` factory functions you pass to `createApp` return a descriptor for the app to construct — for a unit test you want the instance itself. +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"; diff --git a/docs/static/appkit-ui/styles.gen.css b/docs/static/appkit-ui/styles.gen.css index 8fa4174c8..0cc301c72 100644 --- a/docs/static/appkit-ui/styles.gen.css +++ b/docs/static/appkit-ui/styles.gen.css @@ -401,6 +401,9 @@ .\!m-0 { margin: calc(var(--spacing) * 0) !important; } + .m-1 { + margin: calc(var(--spacing) * 1); + } .-mx-1 { margin-inline: calc(var(--spacing) * -1); } @@ -717,6 +720,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 e605829f5..5c5682bfb 100644 --- a/knip.json +++ b/knip.json @@ -8,7 +8,15 @@ ], "workspaces": { "packages/appkit": { - "ignoreDependencies": ["vitest", "@databricks/sdk-experimental"] + "ignoreDependencies": [ + "vitest", + "@databricks/sdk-auth", + "@databricks/sdk-core", + "@databricks/sdk-experimental", + "@databricks/sdk-options", + "@databricks/sdk-statementexecution", + "@databricks/sdk-warehouses" + ] }, "packages/appkit-ui": { "ignoreDependencies": ["tailwindcss", "tw-animate-css"] diff --git a/package.json b/package.json index 54c217804..0bb6b5b80 100644 --- a/package.json +++ b/package.json @@ -94,6 +94,9 @@ "protobufjs@<7.6.2": "7.6.2", "qs@<6.15.2": "6.15.2", "size-sensor": "1.0.3" + }, + "patchedDependencies": { + "@databricks/sdk-statementexecution@0.46.0": "patches/@databricks__sdk-statementexecution@0.46.0.patch" } } } diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 165211c4a..d7b06bfa2 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -71,7 +71,12 @@ "dependencies": { "@ast-grep/napi": "0.37.0", "@databricks/lakebase": "workspace:*", + "@databricks/sdk-auth": "0.46.0", + "@databricks/sdk-core": "0.46.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-options": "0.46.0", + "@databricks/sdk-statementexecution": "0.46.0", + "@databricks/sdk-warehouses": "0.46.0", "@mlflow/core": "0.4.0", "@opentelemetry/api": "1.9.0", "@opentelemetry/api-logs": "0.219.0", diff --git a/packages/appkit/src/cache/index.ts b/packages/appkit/src/cache/index.ts index a98b7dff4..486bd8133 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -128,6 +128,10 @@ export class CacheManager { if (!CacheManager.initPromise) { CacheManager.initPromise = CacheManager.create(userConfig).then( (instance) => { + // Publishes unconditionally: safe only because every getInstance() is + // awaited before any reset(), so a reset() can never land mid-init and + // this can never publish over it. A future unawaited-init caller would + // reintroduce that stale-publish race (the removed `generation` guard). CacheManager.instance = instance; return instance; }, @@ -557,6 +561,20 @@ export class CacheManager { await this.storage.close(); } + /** + * Drop the singleton so the next {@link getInstance} builds a fresh manager. + * + * Both fields must clear — `getInstance()` falls back to `initPromise` when + * `instance` is null. A pointer drop, not teardown: call {@link close} first + * or the old storage leaks (a `pg.Pool` under `PersistentStorage`). + * + * @internal + */ + static reset(): void { + CacheManager.instance = null; + CacheManager.initPromise = null; + } + /** * 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..a8804f0bb --- /dev/null +++ b/packages/appkit/src/cache/tests/cache-manager-reset.test.ts @@ -0,0 +1,100 @@ +import type { CacheEntry } from "shared"; +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("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("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. + class EndableStorage extends InMemoryStorage { + private ended = false; + override async close(): Promise { + this.ended = true; + } + override async set(key: string, entry: CacheEntry): Promise { + if (this.ended) + throw new Error("Cannot use a pool after calling end()"); + return super.set(key, entry); + } + } + const endable = () => + new EndableStorage({ enabled: true, maxSize: 100 } as never); + + const first = await CacheManager.getInstance({ storage: endable() }); + 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: endable() }); + 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: endable() }); + 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/connectors/files/tests/client.test.ts b/packages/appkit/src/connectors/files/tests/client.test.ts index 31ebeb489..09be16e7c 100644 --- a/packages/appkit/src/connectors/files/tests/client.test.ts +++ b/packages/appkit/src/connectors/files/tests/client.test.ts @@ -1,55 +1,38 @@ import { createMockTelemetry } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { createApiError } from "../../../testing"; import type { WorkspaceClient } from "../../../workspace-client"; +import { ApiError } from "../../../workspace-client"; import { FilesConnector } from "../client"; import { streamFromChunks, streamFromString } from "./utils"; -const { mockFilesApi, mockConfig, mockClient, MockApiError } = vi.hoisted( - () => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - - const mockConfig = { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }; - - const mockApiClient = { - userAgent: vi.fn(() => "@databricks/appkit/9.9.9"), - }; - const mockClient = { - files: mockFilesApi, - config: mockConfig, - apiClient: mockApiClient, - } as unknown as WorkspaceClient; - - class MockApiError extends Error { - errorCode: string; - statusCode: number; - constructor( - message: string, - errorCode: string, - statusCode: number, - _response?: any, - _details?: any[], - ) { - super(message); - this.name = "ApiError"; - this.errorCode = errorCode; - this.statusCode = statusCode; - } - } +const { mockFilesApi, mockConfig, mockClient } = vi.hoisted(() => { + const mockFilesApi = { + listDirectoryContents: vi.fn(), + download: vi.fn(), + getMetadata: vi.fn(), + upload: vi.fn(), + createDirectory: vi.fn(), + delete: vi.fn(), + }; - return { mockFilesApi, mockConfig, mockClient, MockApiError }; - }, -); + const mockConfig = { + host: "https://test.databricks.com", + authenticate: vi.fn(), + }; + + const mockApiClient = { + userAgent: vi.fn(() => "@databricks/appkit/9.9.9"), + }; + const mockClient = { + files: mockFilesApi, + config: mockConfig, + apiClient: mockApiClient, + } as unknown as WorkspaceClient; + + return { mockFilesApi, mockConfig, mockClient }; +}); vi.mock("../../../workspace-client", async (importOriginal) => { const actual = @@ -57,7 +40,6 @@ vi.mock("../../../workspace-client", async (importOriginal) => { return { ...actual, createWorkspaceClient: () => mockClient, - ApiError: MockApiError, }; }); @@ -375,7 +357,11 @@ describe("FilesConnector", () => { test("returns false on 404 ApiError", async () => { mockFilesApi.getMetadata.mockRejectedValue( - new MockApiError("Not found", "NOT_FOUND", 404), + createApiError({ + message: "Not found", + errorCode: "NOT_FOUND", + statusCode: 404, + }), ); const result = await connector.exists(mockClient, "missing.txt"); @@ -385,7 +371,11 @@ describe("FilesConnector", () => { test("rethrows non-404 ApiError", async () => { mockFilesApi.getMetadata.mockRejectedValue( - new MockApiError("Server error", "SERVER_ERROR", 500), + createApiError({ + message: "Server error", + errorCode: "SERVER_ERROR", + statusCode: 500, + }), ); await expect(connector.exists(mockClient, "file.txt")).rejects.toThrow( @@ -579,7 +569,7 @@ describe("FilesConnector", () => { try { await connector.upload(mockClient, "file.txt", "data"); } catch (error) { - expect(error).toBeInstanceOf(MockApiError); + expect(error).toBeInstanceOf(ApiError); expect((error as any).statusCode).toBe(403); } }); diff --git a/packages/appkit/src/connectors/lakebase/tests/pool-manager.test.ts b/packages/appkit/src/connectors/lakebase/tests/pool-manager.test.ts index 9def8f43e..8462d83fa 100644 --- a/packages/appkit/src/connectors/lakebase/tests/pool-manager.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/pool-manager.test.ts @@ -1,21 +1,6 @@ import type { Pool } from "pg"; import { afterEach, describe, expect, test, vi } from "vitest"; -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => ({ - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_k: unknown[], fn: (signal?: AbortSignal) => Promise) => - fn(), - ), - generateKey: vi.fn(() => "test-key"), - })), - }, -})); - const mockPools: Pool[] = []; vi.mock("../index", () => ({ diff --git a/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts b/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts index c19f7c15e..4ccccd86d 100644 --- a/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts @@ -3,21 +3,6 @@ import { describe, expect, test, vi } from "vitest"; import { RoutingPool } from "../routing-pool"; -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => ({ - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_k: unknown[], fn: (signal?: AbortSignal) => Promise) => - fn(), - ), - generateKey: vi.fn(() => "test-key"), - })), - }, -})); - function makeMockPool(label: string) { return { query: vi.fn(async () => ({ rows: [{ source: label }] })), diff --git a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts index 17d099e37..af5bfbb18 100644 --- a/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts +++ b/packages/appkit/src/connectors/sql-warehouse/arrow-schema.ts @@ -54,12 +54,12 @@ export function parseDatabricksType(typeText: string): DataType { export function buildEmptyArrowIPCBase64( columns: Array<{ name?: string; - type_text?: string; - type_name?: string; + typeText?: string; + typeName?: string; }>, ): string { const fields = columns.map((col, index) => { - const typeText = col.type_text ?? col.type_name ?? "STRING"; + const typeText = col.typeText ?? col.typeName ?? "STRING"; let dataType: DataType; try { dataType = parseDatabricksType(typeText); diff --git a/packages/appkit/src/connectors/sql-warehouse/client.ts b/packages/appkit/src/connectors/sql-warehouse/client.ts index 439658cbd..6e131c6e1 100644 --- a/packages/appkit/src/connectors/sql-warehouse/client.ts +++ b/packages/appkit/src/connectors/sql-warehouse/client.ts @@ -21,10 +21,14 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; -import { - Context, - type sql, - type WorkspaceClient, +import type { + EndpointState, + ExecuteStatementRequest, + ExternalLink, + ResultData, + StatementResponse, + StatementStatus, + WorkspaceClient, } from "../../workspace-client"; import { buildEmptyArrowIPCBase64 } from "./arrow-schema"; import { executeStatementDefaults } from "./defaults"; @@ -40,9 +44,7 @@ const logger = createLogger("connectors:sql-warehouse"); * Arrow result to match the JSON path. Returns `undefined` when the manifest * carries no columns. */ -function arrowColumnNames( - response: sql.StatementResponse, -): string[] | undefined { +function arrowColumnNames(response: StatementResponse): string[] | undefined { const cols = response.manifest?.schema?.columns; if (!cols || cols.length === 0) return undefined; return cols.map((c, i) => @@ -50,6 +52,46 @@ function arrowColumnNames( ); } +/** + * Coerce the modular SDK's `bigint` row/byte counts back to `number` (the type + * the legacy SDK used). AppKit never does arithmetic on these — they are purely + * informational — but a stray `bigint` makes `JSON.stringify` throw ("Do not + * know how to serialize a BigInt") the instant the result is cached or written + * to an SSE frame. Reyden's INLINE + ARROW_STREAM result — which the analytics + * arrow path caches — carries them on `result`/`manifest`, so normalize every + * statement response at the SDK boundary. Mutates in place (the response is a + * fresh unmarshalled object, owned by the caller). + */ +const BIGINT_COUNT_FIELDS = ["rowOffset", "rowCount", "byteCount"] as const; + +function normalizeResultCounts(result: unknown): void { + if (!result || typeof result !== "object") return; + const r = result as Record; + for (const key of BIGINT_COUNT_FIELDS) { + if (typeof r[key] === "bigint") r[key] = Number(r[key]); + } + // EXTERNAL_LINKS entries carry the same count fields. + if (Array.isArray(r.externalLinks)) { + for (const link of r.externalLinks) normalizeResultCounts(link); + } +} + +function normalizeStatementCounts(response: T): T { + const manifest = response?.manifest as Record | undefined; + if (manifest) { + for (const key of ["totalRowCount", "totalByteCount"] as const) { + if (typeof manifest[key] === "bigint") + manifest[key] = Number(manifest[key]); + } + // Per-chunk `BaseChunkInfo` entries carry the same bigint count fields. + if (Array.isArray(manifest.chunks)) { + for (const chunk of manifest.chunks) normalizeResultCounts(chunk); + } + } + normalizeResultCounts(response?.result); + return response; +} + /** * Maximum size for inline Arrow IPC attachments (25 MiB decoded — the * Databricks Statement Execution API hard cap on INLINE responses). @@ -64,8 +106,8 @@ const MAX_INLINE_ATTACHMENT_BYTES = 25 * 1024 * 1024; /** * Safety cap on how many additional EXTERNAL_LINKS chunks * {@link SQLWarehouseConnector._resolveAllExternalLinks} will follow when the - * manifest omits `total_chunk_count`. High enough to cover any real result; - * only bounds a misbehaving warehouse with a cyclic `next_chunk_index`. + * manifest omits `totalChunkCount`. High enough to cover any real result; + * only bounds a misbehaving warehouse with a cyclic `nextChunkIndex`. */ const MAX_EXTERNAL_CHUNK_FOLLOWS = 10_000; @@ -105,7 +147,7 @@ const WAREHOUSE_RUNNING_CACHE_TTL_MS = 30_000; */ export interface WarehouseStatusUpdate { /** Current state from the SDK (RUNNING | STARTING | STOPPED | STOPPING | DELETED | DELETING). */ - state: sql.State; + state: EndpointState; /** Milliseconds elapsed since `ensureWarehouseRunning` was called. */ elapsedMs: number; /** 1-based attempt counter — useful for tests and telemetry. */ @@ -203,7 +245,7 @@ export class SQLWarehouseConnector { async executeStatement( workspaceClient: WorkspaceClient, - input: sql.ExecuteStatementRequest, + input: ExecuteStatementRequest, signal?: AbortSignal, ) { const startTime = Date.now(); @@ -220,7 +262,7 @@ export class SQLWarehouseConnector { kind: SpanKind.CLIENT, attributes: { "db.system": "databricks", - "db.warehouse_id": input.warehouse_id || "", + "db.warehouse_id": input.warehouseId || "", "db.catalog": input.catalog ?? "", "db.schema": input.schema ?? "", "db.statement": input.statement?.substring(0, 500) || "", @@ -252,52 +294,52 @@ export class SQLWarehouseConnector { throw ValidationError.missingField("statement"); } - if (!input.warehouse_id) { + if (!input.warehouseId) { throw ValidationError.missingField("warehouse_id"); } - const body: sql.ExecuteStatementRequest = { + const body: ExecuteStatementRequest = { statement: input.statement, parameters: input.parameters, - warehouse_id: input.warehouse_id, + warehouseId: input.warehouseId, catalog: input.catalog, schema: input.schema, - wait_timeout: - input.wait_timeout || executeStatementDefaults.wait_timeout, + waitTimeout: + input.waitTimeout || executeStatementDefaults.waitTimeout, disposition: input.disposition || executeStatementDefaults.disposition, format: input.format || executeStatementDefaults.format, - byte_limit: input.byte_limit, - row_limit: input.row_limit, - on_wait_timeout: - input.on_wait_timeout || executeStatementDefaults.on_wait_timeout, + byteLimit: input.byteLimit, + rowLimit: input.rowLimit, + onWaitTimeout: + input.onWaitTimeout || executeStatementDefaults.onWaitTimeout, }; span.addEvent("statement.submitting", { - "db.warehouse_id": input.warehouse_id, + "db.warehouse_id": input.warehouseId, }); const response = - await workspaceClient.statementExecution.executeStatement( - body, - this._createContext(signal), - ); + await workspaceClient.statementExecution.executeStatement(body, { + signal, + }); if (!response) { throw ConnectionError.apiFailure("SQL Warehouse"); } + normalizeStatementCounts(response); const status = response.status; - const statementId = response.statement_id as string; + const statementId = response.statementId as string; span.setAttribute("db.statement_id", statementId); span.addEvent("statement.submitted", { - "db.statement_id": response.statement_id, + "db.statement_id": response.statementId, "db.status": status?.state, }); let result: - | sql.StatementResponse - | { result: { statement_id: string; status: sql.StatementStatus } }; + | StatementResponse + | { result: { statement_id: string; status: StatementStatus } }; switch (status?.state) { case "RUNNING": @@ -322,7 +364,7 @@ export class SQLWarehouseConnector { case "FAILED": throw ExecutionError.statementFailed( status.error?.message, - status.error?.error_code, + status.error?.errorCode, ); case "CANCELED": throw ExecutionError.canceled(); @@ -336,7 +378,7 @@ export class SQLWarehouseConnector { const resultData = result.result as any; const rowCount = - resultData?.data?.length ?? resultData?.data_array?.length ?? 0; + resultData?.data?.length ?? resultData?.dataArray?.length ?? 0; if (rowCount > 0) { span.setAttribute("db.result.row_count", rowCount); @@ -344,7 +386,7 @@ export class SQLWarehouseConnector { const duration = Date.now() - startTime; logger.event()?.setContext("sql-warehouse", { - warehouse_id: input.warehouse_id, + warehouse_id: input.warehouseId, rows_returned: rowCount, query_duration_ms: duration, }); @@ -385,7 +427,7 @@ export class SQLWarehouseConnector { } const attributes = { - "db.warehouse_id": input.warehouse_id, + "db.warehouse_id": input.warehouseId, "db.catalog": input.catalog ?? "", "db.schema": input.schema ?? "", "db.statement": input.statement?.substring(0, 500) || "", @@ -622,9 +664,9 @@ export class SQLWarehouseConnector { ); } - const info = await workspaceClient.warehouses.get( + const info = await workspaceClient.warehouses.getWarehouse( { id: warehouseId }, - this._createContext(signal), + { signal }, ); const state = info?.state; const summary = info?.health?.summary; @@ -650,9 +692,9 @@ export class SQLWarehouseConnector { if (!didStart) { emitter.emit("STARTING", summary); onWarehouseStartIssued?.(); - await workspaceClient.warehouses.start( + await workspaceClient.warehouses.startWarehouse( { id: warehouseId }, - this._createContext(signal), + { signal }, ); didStart = true; } else { @@ -799,15 +841,14 @@ export class SQLWarehouseConnector { }); const response = - await workspaceClient.statementExecution.getStatement( - { - statement_id: statementId, - }, - this._createContext(signal), + await workspaceClient.statementExecution.getStatementResult( + { statementId }, + { signal }, ); if (!response) { throw ConnectionError.apiFailure("SQL Warehouse"); } + normalizeStatementCounts(response); const status = response.status; @@ -837,7 +878,7 @@ export class SQLWarehouseConnector { case "FAILED": throw ExecutionError.statementFailed( status.error?.message, - status.error?.error_code, + status.error?.errorCode, ); case "CANCELED": throw ExecutionError.canceled(); @@ -871,13 +912,13 @@ export class SQLWarehouseConnector { } private async _transformDataArray( - response: sql.StatementResponse, + response: StatementResponse, workspaceClient: WorkspaceClient, signal?: AbortSignal, ) { if (response.manifest?.format === "ARROW_STREAM") { const result = response.result as - | (sql.ResultData & { attachment?: string }) + | (ResultData & { attachment?: string }) | undefined; // Inline Arrow: pass the base64 IPC attachment through unmodified so @@ -893,20 +934,20 @@ export class SQLWarehouseConnector { // rather than omitting it) — it must NOT go down the streaming path // (`streamChunks([])` rejects), so fall through to synthesize an empty // Arrow table below. - if (result?.external_links && result.external_links.length > 0) { + if (result?.externalLinks && result.externalLinks.length > 0) { return this.updateWithArrowStatus(response, workspaceClient, signal); } // Empty result with a known schema: synthesize a zero-row Arrow IPC // attachment so the client always receives an Arrow Table for // ARROW_STREAM, regardless of whether the warehouse returned data. - // Note: an empty array (`data_array: []`) is truthy, so length-check + // Note: an empty array (`dataArray: []`) is truthy, so length-check // explicitly — otherwise zero-row responses fall through to the JSON // row transform below and return `[]` JSON rows instead of an Arrow // table. const hasNoRows = - !result?.data_array || - (Array.isArray(result.data_array) && result.data_array.length === 0); + !result?.dataArray || + (Array.isArray(result.dataArray) && result.dataArray.length === 0); if (hasNoRows && response.manifest?.schema?.columns) { const synthesized = buildEmptyArrowIPCBase64( response.manifest.schema.columns, @@ -917,19 +958,19 @@ export class SQLWarehouseConnector { }; } - // Inline data_array under ARROW_STREAM (rare): fall through to the + // Inline dataArray under ARROW_STREAM (rare): fall through to the // row transform below. The hook will receive `type: "result"` rows; // callers asking for ARROW_STREAM should not hit this path with // current Databricks warehouses. } - if (!response.result?.data_array || !response.manifest?.schema?.columns) { + if (!response.result?.dataArray || !response.manifest?.schema?.columns) { return response; } const columns = response.manifest.schema.columns; - const transformedData = response.result.data_array.map((row) => { + const transformedData = response.result.dataArray.map((row) => { const obj: Record = {}; row.forEach((value, index) => { const column = columns[index]; @@ -937,7 +978,7 @@ export class SQLWarehouseConnector { // attempt to parse JSON strings for string columns if ( - column?.type_name === "STRING" && + column?.typeName === "STRING" && typeof value === "string" && value && (value[0] === "{" || value[0] === "[") @@ -955,8 +996,8 @@ export class SQLWarehouseConnector { return obj; }); - // remove data_array - const { data_array: _data_array, ...restResult } = response.result; + // remove dataArray + const { dataArray: _dataArray, ...restResult } = response.result; return { ...response, result: { @@ -978,7 +1019,7 @@ export class SQLWarehouseConnector { * mechanism used for both INLINE and EXTERNAL_LINKS. */ private _validateArrowAttachment( - response: sql.StatementResponse, + response: StatementResponse, attachment: string, ) { // Cap the size to protect against unbounded inline payloads from @@ -1006,14 +1047,16 @@ export class SQLWarehouseConnector { return { ...response, result: { - ...(response.result as sql.ResultData & { + ...(response.result as ResultData & { attachment?: string; columnNames?: string[]; }), - // `statement_id` is a top-level field, not on `ResultData` — carry it + // `statementId` is a top-level field, not on `ResultData` — carry it // onto the result (as the EXTERNAL_LINKS path does) so the route can // advertise it in `X-Appkit-Arrow-Columns-Ref` for wide inline schemas. - statement_id: response.statement_id, + // Kept as the synthetic `statement_id` key (the connector→route wire + // contract), sourced from the modular SDK's camelCase `statementId`. + statement_id: response.statementId, columnNames, }, }; @@ -1023,26 +1066,26 @@ export class SQLWarehouseConnector { } private async updateWithArrowStatus( - response: sql.StatementResponse, + response: StatementResponse, workspaceClient: WorkspaceClient, signal?: AbortSignal, ): Promise<{ result: { statement_id: string; - status: sql.StatementStatus; + status: StatementStatus; columnNames?: string[]; - external_links?: sql.ExternalLink[]; + external_links?: ExternalLink[]; refreshChunkLink?: RefreshChunkLink; }; }> { - const statementId = response.statement_id as string; + const statementId = response.statementId as string; return { result: { statement_id: statementId, status: { state: response.status?.state, error: response.status?.error, - } as sql.StatementStatus, + } as StatementStatus, columnNames: arrowColumnNames(response), // Resolve the pre-signed links for EVERY chunk in the caller's own // execution context. Streaming these directly (see @@ -1069,9 +1112,9 @@ export class SQLWarehouseConnector { /** * Resolve pre-signed links for EVERY chunk of an EXTERNAL_LINKS result. * - * The execute/getStatement response carries only the first chunk's links - * (each link, except the last, exposes `next_chunk_index`); the remaining - * chunks are fetched with `getStatementResultChunkN`. Runs in the caller's + * The execute/getStatementResult response carries only the first chunk's links + * (each link, except the last, exposes `nextChunkIndex`); the remaining + * chunks are fetched with `getResultData`. Runs in the caller's * identity context (user creds for `.obo.sql`), so there is no cross-identity * fetch. Only the tiny link metadata is resolved eagerly — the bytes still * stream one chunk at a time downstream. Without this a multi-chunk result @@ -1080,29 +1123,29 @@ export class SQLWarehouseConnector { private async _resolveAllExternalLinks( workspaceClient: WorkspaceClient, statementId: string, - response: sql.StatementResponse, + response: StatementResponse, signal?: AbortSignal, - ): Promise { - const first = response.result?.external_links; + ): Promise { + const first = response.result?.externalLinks; if (!first || first.length === 0) return first; - const links: sql.ExternalLink[] = [...first]; + const links: ExternalLink[] = [...first]; // Bound the follow loop so a warehouse returning a cyclic/never-ending // `next_chunk_index` can't spin forever. The manifest's chunk count is the // natural bound; fall back to a generous safety cap if it's absent (real // results still terminate earlier when `next_chunk_index` becomes null) so // a missing count doesn't silently truncate a genuine multi-chunk result. const maxFetches = - response.manifest?.total_chunk_count ?? MAX_EXTERNAL_CHUNK_FOLLOWS; + response.manifest?.totalChunkCount ?? MAX_EXTERNAL_CHUNK_FOLLOWS; let next = this._nextChunkIndex(first); for (let fetches = 0; next != null && fetches < maxFetches; fetches++) { if (signal?.aborted) throw ExecutionError.canceled(); - const chunk = - await workspaceClient.statementExecution.getStatementResultChunkN( - { statement_id: statementId, chunk_index: next }, - this._createContext(signal), - ); - const chunkLinks = chunk.external_links ?? []; + const chunk = await workspaceClient.statementExecution.getResultData( + { statementId, chunkIndex: next }, + { signal }, + ); + normalizeResultCounts(chunk); + const chunkLinks = chunk.externalLinks ?? []; if (chunkLinks.length === 0) break; links.push(...chunkLinks); next = this._nextChunkIndex(chunkLinks); @@ -1110,32 +1153,32 @@ export class SQLWarehouseConnector { return links; } - /** The `next_chunk_index` advertised by a chunk's links, if any. */ - private _nextChunkIndex(links: sql.ExternalLink[]): number | undefined { + /** The `nextChunkIndex` advertised by a chunk's links, if any. */ + private _nextChunkIndex(links: ExternalLink[]): number | undefined { for (const link of links) { - if (link.next_chunk_index != null) return link.next_chunk_index; + if (link.nextChunkIndex != null) return link.nextChunkIndex; } return undefined; } /** * A closure that re-mints a single chunk's pre-signed link via - * `getStatementResultChunkN`, bound to the caller's workspace client + + * `getResultData`, bound to the caller's workspace client + * statement id. Created here (in the caller's identity context) so the * streamer — which runs outside that context — can refresh an expired link - * for `.obo.sql` statements without a cross-identity `getStatement`. + * for `.obo.sql` statements without a cross-identity `getStatementResult`. */ private _makeChunkLinkRefresher( workspaceClient: WorkspaceClient, statementId: string, ): RefreshChunkLink { return async (chunkIndex, signal) => { - const chunk = - await workspaceClient.statementExecution.getStatementResultChunkN( - { statement_id: statementId, chunk_index: chunkIndex }, - this._createContext(signal), - ); - return chunk.external_links?.find((l) => l.chunk_index === chunkIndex); + const chunk = await workspaceClient.statementExecution.getResultData( + { statementId, chunkIndex }, + { signal }, + ); + normalizeResultCounts(chunk); + return chunk.externalLinks?.find((l) => l.chunkIndex === chunkIndex); }; } @@ -1147,7 +1190,7 @@ export class SQLWarehouseConnector { * the pre-signed URLs need no auth to download. */ streamExternalLinks( - chunks: sql.ExternalLink[], + chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, ): AsyncGenerator { @@ -1165,10 +1208,12 @@ export class SQLWarehouseConnector { jobId: string, signal?: AbortSignal, ): Promise { - const response = await workspaceClient.statementExecution.getStatement( - { statement_id: jobId }, - this._createContext(signal), - ); + const response = + await workspaceClient.statementExecution.getStatementResult( + { statementId: jobId }, + { signal }, + ); + normalizeStatementCounts(response); return arrowColumnNames(response); } @@ -1187,25 +1232,19 @@ export class SQLWarehouseConnector { if (error instanceof AppKitError) { throw error; } + // The legacy SDK exposed the Databricks error code as `errorCode`; the + // modular SDK's `ApiError` carries it as `code` (e.g. "INVALID_PARAMETER_VALUE"). + // Read either, so callers can still branch on the stable code — notably the + // analytics arrow disposition/format fallback, which keys on + // INVALID_PARAMETER_VALUE / NOT_IMPLEMENTED to switch INLINE↔EXTERNAL_LINKS. const sdkErrorCode = - error && typeof error === "object" && "errorCode" in error - ? (error as { errorCode?: unknown }).errorCode + error && typeof error === "object" + ? ((error as { errorCode?: unknown }).errorCode ?? + (error as { code?: unknown }).code) : undefined; throw ExecutionError.statementFailed( error instanceof Error ? error.message : String(error), typeof sdkErrorCode === "string" ? sdkErrorCode : undefined, ); } - - // create context for cancellation token - private _createContext(signal?: AbortSignal) { - return new Context({ - cancellationToken: { - isCancellationRequested: signal?.aborted ?? false, - onCancellationRequested: (cb: () => void) => { - signal?.addEventListener("abort", cb, { once: true }); - }, - }, - }); - } } diff --git a/packages/appkit/src/connectors/sql-warehouse/defaults.ts b/packages/appkit/src/connectors/sql-warehouse/defaults.ts index b046a5c4a..3a57c8058 100644 --- a/packages/appkit/src/connectors/sql-warehouse/defaults.ts +++ b/packages/appkit/src/connectors/sql-warehouse/defaults.ts @@ -1,18 +1,18 @@ -import type { sql } from "../../workspace-client"; +import type { ExecuteStatementRequest } from "../../workspace-client"; interface ExecuteStatementDefaults { - wait_timeout: string; - disposition: sql.ExecuteStatementRequest["disposition"]; - format: sql.ExecuteStatementRequest["format"]; - on_wait_timeout: sql.ExecuteStatementRequest["on_wait_timeout"]; + waitTimeout: string; + disposition: ExecuteStatementRequest["disposition"]; + format: ExecuteStatementRequest["format"]; + onWaitTimeout: ExecuteStatementRequest["onWaitTimeout"]; timeout: number; } // @TODO: Make these configurable globally and validate right values export const executeStatementDefaults: ExecuteStatementDefaults = { - wait_timeout: "30s", + waitTimeout: "30s", disposition: "INLINE", format: "JSON_ARRAY", - on_wait_timeout: "CONTINUE", + onWaitTimeout: "CONTINUE", timeout: 60000, }; diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts index d8f52f016..b7826e87e 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/arrow-schema.test.ts @@ -428,11 +428,11 @@ describe("parseDatabricksType — error / robustness", () => { describe("buildEmptyArrowIPCBase64", () => { test("produces a decodable empty Arrow Table with the right schema", () => { const columns = [ - { name: "user_id", type_text: "BIGINT" }, - { name: "name", type_text: "STRING" }, - { name: "created_at", type_text: "TIMESTAMP" }, - { name: "balance", type_text: "DECIMAL(10,2)" }, - { name: "active", type_text: "BOOLEAN" }, + { name: "user_id", typeText: "BIGINT" }, + { name: "name", typeText: "STRING" }, + { name: "created_at", typeText: "TIMESTAMP" }, + { name: "balance", typeText: "DECIMAL(10,2)" }, + { name: "active", typeText: "BOOLEAN" }, ]; const b64 = buildEmptyArrowIPCBase64(columns); const buf = Buffer.from(b64, "base64"); @@ -463,9 +463,9 @@ describe("buildEmptyArrowIPCBase64", () => { test("round-trips nested types end-to-end", () => { const columns = [ - { name: "tags", type_text: "ARRAY" }, - { name: "meta", type_text: "STRUCT" }, - { name: "counts", type_text: "MAP" }, + { name: "tags", typeText: "ARRAY" }, + { name: "meta", typeText: "STRUCT" }, + { name: "counts", typeText: "MAP" }, ]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); @@ -476,8 +476,8 @@ describe("buildEmptyArrowIPCBase64", () => { expect(table.schema.fields[2]?.type).toBeInstanceOf(Map_); }); - test("falls back from type_text to type_name when type_text missing", () => { - const columns = [{ name: "id", type_name: "BIGINT" }]; + test("falls back from typeText to typeName when typeText missing", () => { + const columns = [{ name: "id", typeName: "BIGINT" }]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); expect( @@ -487,8 +487,8 @@ describe("buildEmptyArrowIPCBase64", () => { test("unknown type degrades to Utf8 without throwing", () => { const columns = [ - { name: "id", type_text: "BIGINT" }, - { name: "weird", type_text: "FUTURE_TYPE_NOT_YET_SUPPORTED" }, + { name: "id", typeText: "BIGINT" }, + { name: "weird", typeText: "FUTURE_TYPE_NOT_YET_SUPPORTED" }, ]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); @@ -499,7 +499,7 @@ describe("buildEmptyArrowIPCBase64", () => { }); test("missing column name gets a synthesized placeholder", () => { - const columns = [{ type_text: "STRING" }, { name: "", type_text: "INT" }]; + const columns = [{ typeText: "STRING" }, { name: "", typeText: "INT" }]; const buf = Buffer.from(buildEmptyArrowIPCBase64(columns), "base64"); const table = tableFromIPC(buf); expect(table.schema.fields[0]?.name).toBe("column_0"); diff --git a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts index 5a945b0cc..8344df2fc 100644 --- a/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts +++ b/packages/appkit/src/connectors/sql-warehouse/tests/client.test.ts @@ -1,7 +1,10 @@ import { tableFromIPC } from "apache-arrow"; import { describe, expect, test, vi } from "vitest"; -import type { sql } from "../../../workspace-client"; +import type { + ExternalLink, + StatementResponse, +} from "../../../workspace-client"; vi.mock("../../../telemetry", () => { const mockMeter = { @@ -40,11 +43,11 @@ function createConnector() { // `_transformDataArray` is async — it paginates multi-chunk EXTERNAL_LINKS // results. The workspace client is only touched when following -// `next_chunk_index`, so a bare stub suffices for the inline / JSON / +// `nextChunkIndex`, so a bare stub suffices for the inline / JSON / // single-chunk cases; the multi-chunk tests pass a real mock. function transform( connector: SQLWarehouseConnector, - response: sql.StatementResponse, + response: StatementResponse, workspaceClient: unknown = {}, ) { return (connector as any)._transformDataArray(response, workspaceClient); @@ -58,11 +61,11 @@ const REAL_ARROW_ATTACHMENT = describe("SQLWarehouseConnector._transformDataArray", () => { describe("classic warehouse (JSON_ARRAY + INLINE)", () => { - test("transforms data_array rows into named objects", async () => { + test("transforms dataArray rows into named objects", async () => { const connector = createConnector(); // Real response shape from classic warehouse: INLINE + JSON_ARRAY const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY", @@ -71,14 +74,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { columns: [ { name: "test_col", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 0, }, { name: "test_col2", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 1, }, ], @@ -87,33 +90,33 @@ describe("SQLWarehouseConnector._transformDataArray", () => { truncated: false, }, result: { - data_array: [["1", "2"]], + dataArray: [["1", "2"]], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.data).toEqual([{ test_col: "1", test_col2: "2" }]); - expect(result.result.data_array).toBeUndefined(); + expect(result.result.dataArray).toBeUndefined(); }); test("parses JSON strings in STRING columns", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY", schema: { columns: [ - { name: "id", type_name: "INT" }, - { name: "metadata", type_name: "STRING" }, + { name: "id", typeName: "INT" }, + { name: "metadata", typeName: "STRING" }, ], }, }, result: { - data_array: [["1", '{"key":"value"}']], + dataArray: [["1", '{"key":"value"}']], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.data[0].metadata).toEqual({ key: "value" }); @@ -125,26 +128,26 @@ describe("SQLWarehouseConnector._transformDataArray", () => { const connector = createConnector(); // Real response shape from classic warehouse: EXTERNAL_LINKS + ARROW_STREAM const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "test_col", type_name: "INT" }, - { name: "test_col2", type_name: "INT" }, + { name: "test_col", typeName: "INT" }, + { name: "test_col2", typeName: "INT" }, ], }, }, result: { - external_links: [ + externalLinks: [ { - external_link: "https://storage.example.com/chunk0", + externalLink: "https://storage.example.com/chunk0", expiration: "2026-04-15T00:00:00Z", }, ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.statement_id).toBe("stmt-1"); @@ -156,9 +159,9 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("passes attachment through unchanged for client-side decoding", async () => { const connector = createConnector(); // Real response shape from serverless warehouse: INLINE + ARROW_STREAM - // Data arrives in result.attachment as base64-encoded Arrow IPC, not data_array. + // Data arrives in result.attachment as base64-encoded Arrow IPC, not dataArray. const response = { - statement_id: "00000001-test-stmt", + statementId: "00000001-test-stmt", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", @@ -167,30 +170,30 @@ describe("SQLWarehouseConnector._transformDataArray", () => { columns: [ { name: "test_col", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 0, }, { name: "test_col2", - type_text: "INT", - type_name: "INT", + typeText: "INT", + typeName: "INT", position: 1, }, ], - total_chunk_count: 1, - chunks: [{ chunk_index: 0, row_offset: 0, row_count: 1 }], + totalChunkCount: 1, + chunks: [{ chunkIndex: 0, row_offset: 0, row_count: 1 }], total_row_count: 1, }, truncated: false, }, result: { - chunk_index: 0, + chunkIndex: 0, row_offset: 0, row_count: 1, attachment: REAL_ARROW_ATTACHMENT, }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); @@ -206,56 +209,56 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("preserves manifest and status alongside attachment", async () => { const connector = createConnector(); const response = { - statement_id: "00000001-test-stmt", + statementId: "00000001-test-stmt", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "test_col", type_name: "INT" }, - { name: "test_col2", type_name: "INT" }, + { name: "test_col", typeName: "INT" }, + { name: "test_col2", typeName: "INT" }, ], }, }, result: { - chunk_index: 0, + chunkIndex: 0, row_count: 1, attachment: REAL_ARROW_ATTACHMENT, }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); // Manifest, statement_id, and attachment are all preserved expect(result.manifest.format).toBe("ARROW_STREAM"); - expect(result.statement_id).toBe("00000001-test-stmt"); + expect(result.statementId).toBe("00000001-test-stmt"); expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); }); test("synthesizes an empty Arrow IPC attachment for empty results so the client always gets a Table", async () => { const connector = createConnector(); - // Empty result: no attachment, no data_array, no external_links — but + // Empty result: no attachment, no dataArray, no external_links — but // the manifest still describes the schema. The connector should fill in // `attachment` with a zero-row Arrow IPC matching the schema. const response = { - statement_id: "stmt-empty", + statementId: "stmt-empty", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "user_id", type_text: "BIGINT", type_name: "BIGINT" }, - { name: "name", type_text: "STRING", type_name: "STRING" }, + { name: "user_id", typeText: "BIGINT", typeName: "BIGINT" }, + { name: "name", typeText: "STRING", typeName: "STRING" }, { name: "balance", - type_text: "DECIMAL(10,2)", - type_name: "DECIMAL", + typeText: "DECIMAL(10,2)", + typeName: "DECIMAL", }, ], }, total_row_count: 0, }, result: {}, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const transformed = await transform(connector, response); const attachment: string = transformed.result.attachment; @@ -275,18 +278,18 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("does NOT synthesize an attachment when external_links are present", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-ext", + statementId: "stmt-ext", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", - schema: { columns: [{ name: "x", type_text: "INT" }] }, + schema: { columns: [{ name: "x", typeText: "INT" }] }, }, result: { - external_links: [ - { external_link: "https://example.com/x", expiration: "9999" }, + externalLinks: [ + { externalLink: "https://example.com/x", expiration: "9999" }, ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const transformed = await transform(connector, response); // External-links path returns the statement_id projection — no attachment. @@ -296,21 +299,21 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("empty external_links array is a zero-row result → synthesizes an empty table (not the streaming path)", async () => { const connector = createConnector(); - // Some warehouses emit `external_links: []` for a zero-row result rather + // Some warehouses emit `externalLinks: []` for a zero-row result rather // than omitting it. An empty array must NOT go down the streaming path // (streamChunks([]) rejects) — synthesize an empty Arrow table instead. const response = { - statement_id: "stmt-empty-ext", + statementId: "stmt-empty-ext", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { - columns: [{ name: "x", type_text: "INT", type_name: "INT" }], + columns: [{ name: "x", typeText: "INT", typeName: "INT" }], }, total_row_count: 0, }, - result: { external_links: [] }, - } as unknown as sql.StatementResponse; + result: { externalLinks: [] }, + } as unknown as StatementResponse; const transformed = await transform(connector, response); const attachment: string = transformed.result.attachment; @@ -323,11 +326,11 @@ describe("SQLWarehouseConnector._transformDataArray", () => { test("does NOT synthesize an attachment when schema is missing", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-no-schema", + statementId: "stmt-no-schema", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: {}, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const transformed = await transform(connector, response); // Without a schema we cannot build a Table — pass through unchanged. @@ -340,11 +343,11 @@ describe("SQLWarehouseConnector._transformDataArray", () => { // base64 chars decodes to ~27 MiB, comfortably above the limit. const oversized = "A".repeat(36 * 1024 * 1024); const response = { - statement_id: "stmt-oversized", + statementId: "stmt-oversized", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: oversized }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; await expect(transform(connector, response)).rejects.toThrow( /exceeds maximum size/, @@ -352,28 +355,28 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }); }); - describe("ARROW_STREAM with data_array (hypothetical inline variant)", () => { - test("transforms data_array like JSON_ARRAY path", async () => { + describe("ARROW_STREAM with dataArray (hypothetical inline variant)", () => { + test("transforms dataArray like JSON_ARRAY path", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "id", type_name: "INT" }, - { name: "value", type_name: "STRING" }, + { name: "id", typeName: "INT" }, + { name: "value", typeName: "STRING" }, ], }, }, result: { - data_array: [ + dataArray: [ ["1", "hello"], ["2", "world"], ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result.result.data).toEqual([ @@ -384,88 +387,88 @@ describe("SQLWarehouseConnector._transformDataArray", () => { }); describe("edge cases", () => { - test("returns response unchanged when no data_array, attachment, or schema", async () => { + test("returns response unchanged when no dataArray, attachment, or schema", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY" }, result: {}, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); expect(result).toBe(response); }); - test("attachment takes priority over data_array when both present", async () => { + test("attachment takes priority over dataArray when both present", async () => { const connector = createConnector(); const response = { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", schema: { columns: [ - { name: "test_col", type_name: "INT" }, - { name: "test_col2", type_name: "INT" }, + { name: "test_col", typeName: "INT" }, + { name: "test_col2", typeName: "INT" }, ], }, }, result: { attachment: REAL_ARROW_ATTACHMENT, - data_array: [["999", "999"]], + dataArray: [["999", "999"]], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; const result = await transform(connector, response); - // Should pass attachment through (client decodes), not transform data_array + // Should pass attachment through (client decodes), not transform dataArray expect(result.result.attachment).toBe(REAL_ARROW_ATTACHMENT); expect(result.result.data).toBeUndefined(); }); }); describe("multi-chunk EXTERNAL_LINKS pagination", () => { - function multiChunkResponse(totalChunks: number): sql.StatementResponse { + function multiChunkResponse(totalChunks: number): StatementResponse { return { - statement_id: "stmt-multi", + statementId: "stmt-multi", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM", - total_chunk_count: totalChunks, - schema: { columns: [{ name: "x", type_name: "INT" }] }, + totalChunkCount: totalChunks, + schema: { columns: [{ name: "x", typeName: "INT" }] }, }, result: { - external_links: [ + externalLinks: [ { - chunk_index: 0, - external_link: "https://example.com/chunk0", - next_chunk_index: 1, + chunkIndex: 0, + externalLink: "https://example.com/chunk0", + nextChunkIndex: 1, }, ], }, - } as unknown as sql.StatementResponse; + } as unknown as StatementResponse; } - test("follows next_chunk_index to resolve every chunk's links", async () => { + test("follows nextChunkIndex to resolve every chunk's links", async () => { const connector = createConnector(); - const getStatementResultChunkN = vi + const getResultData = vi .fn() .mockResolvedValueOnce({ - external_links: [ + externalLinks: [ { - chunk_index: 1, - external_link: "https://example.com/chunk1", - next_chunk_index: 2, + chunkIndex: 1, + externalLink: "https://example.com/chunk1", + nextChunkIndex: 2, }, ], }) .mockResolvedValueOnce({ - external_links: [ - { chunk_index: 2, external_link: "https://example.com/chunk2" }, + externalLinks: [ + { chunkIndex: 2, externalLink: "https://example.com/chunk2" }, ], }); const workspaceClient = { - statementExecution: { getStatementResultChunkN }, + statementExecution: { getResultData }, }; const result = await transform( @@ -474,16 +477,14 @@ describe("SQLWarehouseConnector._transformDataArray", () => { workspaceClient, ); - expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); - expect(getStatementResultChunkN).toHaveBeenNthCalledWith( + expect(getResultData).toHaveBeenCalledTimes(2); + expect(getResultData).toHaveBeenNthCalledWith( 1, - expect.objectContaining({ statement_id: "stmt-multi", chunk_index: 1 }), + expect.objectContaining({ statementId: "stmt-multi", chunkIndex: 1 }), expect.anything(), ); expect( - result.result.external_links.map( - (l: sql.ExternalLink) => l.external_link, - ), + result.result.external_links.map((l: ExternalLink) => l.externalLink), ).toEqual([ "https://example.com/chunk0", "https://example.com/chunk1", @@ -491,20 +492,20 @@ describe("SQLWarehouseConnector._transformDataArray", () => { ]); }); - test("is bounded by total_chunk_count when next_chunk_index never terminates", async () => { + test("is bounded by totalChunkCount when nextChunkIndex never terminates", async () => { const connector = createConnector(); // Misbehaving warehouse: always advertises another chunk. - const getStatementResultChunkN = vi.fn().mockResolvedValue({ - external_links: [ + const getResultData = vi.fn().mockResolvedValue({ + externalLinks: [ { - chunk_index: 1, - external_link: "https://example.com/loop", - next_chunk_index: 99, + chunkIndex: 1, + externalLink: "https://example.com/loop", + nextChunkIndex: 99, }, ], }); const workspaceClient = { - statementExecution: { getStatementResultChunkN }, + statementExecution: { getResultData }, }; const result = await transform( @@ -513,8 +514,8 @@ describe("SQLWarehouseConnector._transformDataArray", () => { workspaceClient, ); - // Terminates (no hang) — capped at total_chunk_count fetches. - expect(getStatementResultChunkN).toHaveBeenCalledTimes(2); + // Terminates (no hang) — capped at totalChunkCount fetches. + expect(getResultData).toHaveBeenCalledTimes(2); expect(result.result.external_links.length).toBeGreaterThan(0); }); }); diff --git a/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts b/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts index aba8488c3..a9061f61d 100644 --- a/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts +++ b/packages/appkit/src/connectors/sql-warehouse/warehouse-status-emitter.ts @@ -1,5 +1,5 @@ import type { Span } from "../../telemetry"; -import type { sql } from "../../workspace-client"; +import type { EndpointState } from "../../workspace-client"; import type { WarehouseStatusUpdate } from "./client"; /** @@ -10,7 +10,7 @@ import type { WarehouseStatusUpdate } from "./client"; */ export class WarehouseStatusEmitter { attempt = 0; - private lastEmittedState: sql.State | null = null; + private lastEmittedState: EndpointState | null = null; constructor( private readonly span: Span, @@ -18,7 +18,7 @@ export class WarehouseStatusEmitter { private readonly onStatus: (update: WarehouseStatusUpdate) => void, ) {} - emit(state: sql.State, summary: string | undefined): void { + emit(state: EndpointState, summary: string | undefined): void { this.attempt += 1; this.span.addEvent("warehouse.status", { "db.warehouse.state": state, diff --git a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts index 285fa8d0b..af1480aa1 100644 --- a/packages/appkit/src/connectors/tests/sql-warehouse.test.ts +++ b/packages/appkit/src/connectors/tests/sql-warehouse.test.ts @@ -61,7 +61,7 @@ describe("SQLWarehouseConnector", () => { await expect( connector.executeStatement(mockWorkspaceClient as any, { statement: sensitiveStatement, - warehouse_id: "test-warehouse", + warehouseId: "test-warehouse", }), ).rejects.toThrow(); @@ -89,7 +89,9 @@ describe("SQLWarehouseConnector", () => { statement_id: "stmt-123", status: { state: "RUNNING" }, }), - getStatement: vi.fn().mockRejectedValue(new Error("polling timeout")), + getStatementResult: vi + .fn() + .mockRejectedValue(new Error("polling timeout")), }, config: { host: "https://test.databricks.com" }, }; @@ -97,7 +99,7 @@ describe("SQLWarehouseConnector", () => { await expect( connector.executeStatement(mockWorkspaceClient as any, { statement: "SELECT secret_data FROM vault", - warehouse_id: "test-warehouse", + warehouseId: "test-warehouse", }), ).rejects.toThrow(); @@ -118,6 +120,141 @@ describe("SQLWarehouseConnector", () => { }); }); + describe("statement error-code propagation", () => { + let connector: SQLWarehouseConnector; + + beforeEach(() => { + vi.clearAllMocks(); + connector = new SQLWarehouseConnector({ timeout: 5000 }); + }); + + // Regression: the modular `@databricks/sdk-core` `ApiError` carries the + // Databricks error code on `.code`, whereas the legacy SDK used + // `.errorCode`. The analytics arrow disposition/format fallback keys on + // this code ("INVALID_PARAMETER_VALUE" / "NOT_IMPLEMENTED") to switch + // INLINE→EXTERNAL_LINKS, so the connector MUST surface either field as + // `ExecutionError.errorCode` — reading only `.errorCode` broke every arrow + // query (the INLINE+ARROW_STREAM probe rejection went unrecognized). + test("surfaces the modular SDK ApiError.code as ExecutionError.errorCode", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + class FakeApiError extends Error { + readonly code = "INVALID_PARAMETER_VALUE"; + } + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi + .fn() + .mockRejectedValue( + new FakeApiError( + "Incompatible parameters: The format field must be JSON_ARRAY when the disposition field is INLINE.", + ), + ), + }, + config: { host: "https://test.databricks.com" }, + }; + + await expect( + connector.executeStatement(mockWorkspaceClient as any, { + statement: "SELECT 1", + warehouseId: "test-warehouse", + disposition: "INLINE", + format: "ARROW_STREAM", + }), + ).rejects.toMatchObject({ errorCode: "INVALID_PARAMETER_VALUE" }); + + errorSpy.mockRestore(); + }); + + // A failed statement STATUS (not a thrown ApiError) still carries the code + // on `status.error.errorCode` — the SDK unmarshals `error_code` there, so + // that path was already correct and must stay so. + test("surfaces status.error.errorCode from a FAILED statement status", async () => { + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + statementId: "stmt-123", + status: { + state: "FAILED", + error: { + errorCode: "INVALID_PARAMETER_VALUE", + message: "bad parameter", + }, + }, + }), + }, + config: { host: "https://test.databricks.com" }, + }; + + await expect( + connector.executeStatement(mockWorkspaceClient as any, { + statement: "SELECT 1", + warehouseId: "test-warehouse", + }), + ).rejects.toMatchObject({ errorCode: "INVALID_PARAMETER_VALUE" }); + + errorSpy.mockRestore(); + }); + }); + + describe("bigint count normalization", () => { + let connector: SQLWarehouseConnector; + + beforeEach(() => { + vi.clearAllMocks(); + connector = new SQLWarehouseConnector({ timeout: 5000 }); + }); + + // Regression: the modular SDK types rowCount/byteCount/rowOffset (and the + // per-chunk BaseChunkInfo counts) as `bigint`, whereas the legacy SDK used + // `number`. Reyden's INLINE+ARROW_STREAM result is cached by the analytics + // arrow path, and `JSON.stringify` throws ("Do not know how to serialize a + // BigInt") on any surviving bigint — which broke EVERY query on Reyden. The + // connector must coerce these to `number` at the SDK boundary so the result + // stays serializable for the cache / SSE frames. + test("coerces bigint manifest/result/chunk counts so the result is JSON-serializable", async () => { + const mockWorkspaceClient = { + statementExecution: { + executeStatement: vi.fn().mockResolvedValue({ + statementId: "stmt-1", + status: { state: "SUCCEEDED" }, + manifest: { + format: "JSON_ARRAY", + totalRowCount: 2n, + totalByteCount: 100n, + chunks: [ + { chunkIndex: 0, rowOffset: 0n, rowCount: 2n, byteCount: 100n }, + ], + schema: { columns: [{ name: "id", typeName: "INT" }] }, + }, + result: { + dataArray: [["1"], ["2"]], + rowOffset: 0n, + rowCount: 2n, + byteCount: 100n, + }, + }), + }, + config: { host: "https://test.databricks.com" }, + }; + + const out: any = await connector.executeStatement( + mockWorkspaceClient as any, + { statement: "SELECT id FROM t", warehouseId: "reyden" }, + ); + + // The arrow cache serializes exactly this — it must not throw. + expect(() => JSON.stringify(out)).not.toThrow(); + // Counts are coerced to number (legacy parity), including per-chunk ones. + expect(typeof out.manifest.totalRowCount).toBe("number"); + expect(typeof out.manifest.totalByteCount).toBe("number"); + expect(typeof out.manifest.chunks[0].byteCount).toBe("number"); + expect(typeof out.result.rowCount).toBe("number"); + }); + }); + describe("ensureWarehouseRunning", () => { let connector: SQLWarehouseConnector; @@ -137,7 +274,9 @@ describe("SQLWarehouseConnector", () => { test("emits a single RUNNING update and returns when warehouse is already running", async () => { const get = vi.fn().mockResolvedValue({ state: "RUNNING" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await connector.ensureWarehouseRunning(wsClient as any, "wh-1", { @@ -158,7 +297,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn().mockResolvedValue(undefined); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; const promise = connector.ensureWarehouseRunning( @@ -191,7 +332,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; const promise = connector.ensureWarehouseRunning( @@ -212,7 +355,9 @@ describe("SQLWarehouseConnector", () => { test("rejects when warehouse is DELETED", async () => { const get = vi.fn().mockResolvedValue({ state: "DELETED" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await expect( @@ -228,7 +373,9 @@ describe("SQLWarehouseConnector", () => { test("rejects when warehouse is DELETING", async () => { const get = vi.fn().mockResolvedValue({ state: "DELETING" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const updates: any[] = []; await expect( @@ -243,7 +390,9 @@ describe("SQLWarehouseConnector", () => { test("aborts immediately when signal is already aborted", async () => { const get = vi.fn(); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const controller = new AbortController(); controller.abort(); @@ -258,7 +407,9 @@ describe("SQLWarehouseConnector", () => { test("times out if warehouse never reaches RUNNING", async () => { const get = vi.fn().mockResolvedValue({ state: "STARTING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const promise = connector.ensureWarehouseRunning( wsClient as any, @@ -281,7 +432,7 @@ describe("SQLWarehouseConnector", () => { test("rejects when warehouse_id is empty", async () => { const wsClient = { - warehouses: { get: vi.fn(), start: vi.fn() }, + warehouses: { getWarehouse: vi.fn(), startWarehouse: vi.fn() }, }; await expect( @@ -293,7 +444,9 @@ describe("SQLWarehouseConnector", () => { test("skips the SDK round-trip on a subsequent call within the recently-running TTL", async () => { const get = vi.fn().mockResolvedValue({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const updates1: any[] = []; await connector.ensureWarehouseRunning(wsClient as any, "wh-cache", { @@ -314,7 +467,9 @@ describe("SQLWarehouseConnector", () => { test("rejects with ConfigurationError when STOPPED and autoStart is false", async () => { const get = vi.fn().mockResolvedValue({ state: "STOPPED" }); const start = vi.fn(); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; await expect( connector.ensureWarehouseRunning(wsClient as any, "wh-no-auto", { @@ -332,7 +487,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const updates: any[] = []; const promise = connector.ensureWarehouseRunning( @@ -356,7 +513,9 @@ describe("SQLWarehouseConnector", () => { const sensitive = "getaddrinfo ENOTFOUND adb-1234567890.10.azuredatabricks.net"; const get = vi.fn().mockRejectedValue(new Error(sensitive)); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; await expect( connector.ensureWarehouseRunning(wsClient as any, "wh-leak", { @@ -381,7 +540,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn().mockResolvedValue(undefined); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const allUpdates = [0, 1, 2].map(() => [] as { state: string }[]); const waits = allUpdates.map((updates) => @@ -406,7 +567,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const controller = new AbortController(); const aborted = connector.ensureWarehouseRunning( @@ -438,7 +601,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const mount1 = new AbortController(); const first = connector.ensureWarehouseRunning( @@ -467,7 +632,9 @@ describe("SQLWarehouseConnector", () => { test("orphan before warehouses.start is aborted on the next microtask", async () => { const get = vi.fn().mockResolvedValue({ state: "STARTING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; const controller = new AbortController(); const only = connector.ensureWarehouseRunning( @@ -485,7 +652,7 @@ describe("SQLWarehouseConnector", () => { await Promise.resolve(); expect(get).toHaveBeenCalledTimes(1); - expect(wsClient.warehouses.start).not.toHaveBeenCalled(); + expect(wsClient.warehouses.startWarehouse).not.toHaveBeenCalled(); }); test("orphan after warehouses.start runs poll to completion", async () => { @@ -495,7 +662,9 @@ describe("SQLWarehouseConnector", () => { .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); const start = vi.fn().mockResolvedValue(undefined); - const wsClient = { warehouses: { get, start } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: start }, + }; const controller = new AbortController(); const only = connector.ensureWarehouseRunning( @@ -532,7 +701,9 @@ describe("SQLWarehouseConnector", () => { .fn() .mockResolvedValueOnce({ state: "STARTING" }) .mockResolvedValueOnce({ state: "RUNNING" }); - const wsClient = { warehouses: { get, start: vi.fn() } }; + const wsClient = { + warehouses: { getWarehouse: get, startWarehouse: vi.fn() }, + }; let callCount = 0; const promise = connector.ensureWarehouseRunning( diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index f5f4f2725..42227b992 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -27,10 +27,20 @@ import { isToolProvider, PluginContext } from "./plugin-context"; const logger = createLogger("appkit"); +/** + * Internal teardown entry for the test harness (see `createTestApp`). + * Symbol-keyed so it cannot collide with — or be shadowed by — a plugin + * manifest name, and so it stays off the public `PluginMap` surface. + * @internal + */ +export const disposeApp = Symbol("appkit.internal.dispose"); + export class AppKit { #pluginInstances: Record = {}; #setupPromises: Promise[] = []; #context: PluginContext; + /** Owns the shutdown sequence; assigned once every plugin has started. */ + #lifecycle: LifecycleManager | undefined; private constructor(config: { plugins: TPlugins }) { const { plugins, ...globalConfig } = config; @@ -190,6 +200,13 @@ export class AppKit { client?: WorkspaceClient; onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; + /** + * Skip installing the SIGTERM/SIGINT handlers. Internal, and not exposed + * on {@link createApp}: only the test harness sets it, because it boots + * repeatedly in one process and manages its own teardown, so + * accumulating signal handlers would be a leak. + */ + installSignalHandlers?: boolean; } = {}, ): Promise> { // Initialize core services @@ -231,11 +248,11 @@ export class AppKit { await instance.#context.emitLifecycle("setup:complete"); - const handle = instance as unknown as PluginMap; + const app = instance as unknown as PluginMap; if (config.onPluginsReady) { logger.debug("Running onPluginsReady hook"); - await config.onPluginsReady(handle); + await config.onPluginsReady(app); logger.debug("onPluginsReady hook completed"); } @@ -252,9 +269,24 @@ 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(); + instance.#lifecycle = new LifecycleManager(instance.#context); + if (config.installSignalHandlers !== false) { + instance.#lifecycle.installSignalHandlers(); + } + + return app; + } - return handle; + /** + * Internal teardown entry point for the test harness: delegates to the + * lifecycle's non-exiting `shutdown({ exit: false })` (the phases are + * canonical there). Not public API — AppKit does not support re-booting or + * embedding, so real apps tear down only through the signal path. The harness + * drops the core singletons and restores env after this resolves. + * @internal + */ + async [disposeApp](): Promise { + await this.#lifecycle?.shutdown({ exit: false }); } private static bootstrapInternalTelemetry(): void { @@ -388,6 +420,7 @@ export async function createApp< telemetry?: TelemetryConfig; cache?: CacheConfig; client?: WorkspaceClient; + /** Runs after plugin setup but **before** the server starts. */ onPluginsReady?: (appkit: PluginMap) => void | Promise; disableInternalTelemetry?: boolean; } = {}, diff --git a/packages/appkit/src/core/lifecycle-manager.ts b/packages/appkit/src/core/lifecycle-manager.ts index 84dcb4a9e..cfe4147b6 100644 --- a/packages/appkit/src/core/lifecycle-manager.ts +++ b/packages/appkit/src/core/lifecycle-manager.ts @@ -46,34 +46,31 @@ export class LifecycleManager { private static readonly PHASE_SHUTDOWN_TIMEOUT_MS = 2_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 the + * harness path (`{ exit: false }`), 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"; constructor(private readonly context: PluginContext) {} /** - * 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}. + * Install the SIGTERM/SIGINT handlers that trigger {@link shutdown}. Never + * removed: the signal path exits the process, and the harness opts out of + * installing them, so nothing accumulates across boots. */ installSignalHandlers(): void { - process.once("SIGTERM", () => this.shutdown()); - process.once("SIGINT", () => this.shutdown()); + process.once("SIGTERM", () => void this.shutdown()); + process.once("SIGINT", () => void this.shutdown()); } /** - * Run the graceful-shutdown sequence and exit the process. + * Run the graceful-shutdown sequence. Exits the process unless + * `exit: false` — the flag the test harness passes so it can tear a booted + * app down between tests without killing the vitest process. * * Phases: * 1. stop the internal-telemetry reporter @@ -83,25 +80,28 @@ export class LifecycleManager { * 4. emit the `"shutdown"` lifecycle event, bounded * 5. close the cache storage and flush telemetry concurrently, each bounded * + * Every phase is individually bounded, so the sequence always completes — + * `{ exit: false }` therefore needs no outer timeout and an `afterEach` + * cannot hang on it. A second call joins the first teardown. + * * 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. */ - 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; + async shutdown(options: { exit?: boolean } = {}): Promise { + const exit = options.exit ?? true; - logger.info("Starting graceful shutdown..."); - - let exitCode = 0; + if (!exit) { + // Harness path: no backstop, no process.exit. The phases are internally + // bounded, and this is fully awaited before the harness drops the + // singletons — so phase 5 always acts on this app's own cache/telemetry. + await this.runPhasesOnce(); + return; + } - // 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. Belt-and-suspenders over the per-phase budgets. 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 +110,28 @@ 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.runPhasesOnce(); + + clearTimeout(forceExitTimer); + process.exit(exitCode); + } + + /** No `await` between read and assign — that gap is the re-entrancy window. */ + private runPhasesOnce(): 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..."); + + let exitCode = 0; + try { const plugins = Array.from(this.context.getPlugins().values()); @@ -184,17 +199,16 @@ export class LifecycleManager { exitCode = 1; } - clearTimeout(forceExitTimer); - process.exit(exitCode); + return exitCode; } - /** Close the cache storage, bounded and error-isolated. */ + /** Bounded and error-isolated. Reads the cache manager at phase-5 time. */ private async closeCacheStorage(): Promise { - let cache: CacheManager; + let cache: CacheManager | undefined; try { cache = CacheManager.getInstanceSync(); } catch { - // Cache was never initialized — nothing to close. + // Never initialized — nothing to close. return; } try { @@ -208,11 +222,19 @@ export class LifecycleManager { } } - /** Flush and shut down the telemetry SDK, bounded and error-isolated. */ + /** Bounded and error-isolated. Reads the telemetry manager at phase-5 time. */ private async flushTelemetry(): Promise { + let telemetry: TelemetryManager | undefined; + try { + telemetry = TelemetryManager.getInstance(); + } catch { + // Unavailable or mocked away — nothing to flush. + return; + } + 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/tests/lifecycle-manager.test.ts b/packages/appkit/src/core/tests/lifecycle-manager.test.ts index 121e7eb5c..461881910 100644 --- a/packages/appkit/src/core/tests/lifecycle-manager.test.ts +++ b/packages/appkit/src/core/tests/lifecycle-manager.test.ts @@ -380,5 +380,206 @@ describe("LifecycleManager", () => { expect(signals).toContain("SIGINT"); onceSpy.mockRestore(); }); + + test("a manager that never installs them adds no listener", () => { + // The harness boots this way (installSignalHandlers: false), so repeated + // boots in one process must not accumulate handlers — there is no removal + // path any more. + const termBaseline = process.listenerCount("SIGTERM"); + const intBaseline = process.listenerCount("SIGINT"); + + new LifecycleManager(contextWithPlugins({})); + + expect(process.listenerCount("SIGTERM")).toBe(termBaseline); + expect(process.listenerCount("SIGINT")).toBe(intBaseline); + }); + }); + + describe("shutdown({ exit: false }) (the harness 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.shutdown({ exit: false }); + + expect(stop).toHaveBeenCalledTimes(1); + expect(abortActiveOperations).toHaveBeenCalledTimes(1); + expect(shutdown).toHaveBeenCalledTimes(1); + expect(emit).toHaveBeenCalledWith("shutdown"); + expect(cacheClose).toHaveBeenCalledTimes(1); + expect(telemetryShutdown).toHaveBeenCalledTimes(1); + + // shutdown({ exit: false }) never exits — that is the whole point of the + // harness path. (Dropping the core singletons is the harness's job, not + // the manager's.) + 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 + .shutdown({ exit: false }) + .then(() => observed.push(`first:${teardownFinished}`)); + const second = manager + .shutdown({ exit: false }) + .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 a harness shutdown 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.shutdown({ exit: false }); + // The signal path after a harness shutdown: teardown is memoized, so the + // phases do not run twice even though the exiting shutdown() is callable. + await manager.shutdown(); + + expect(shutdown).toHaveBeenCalledTimes(1); + }); + + test("a harness shutdown after a signal-initiated teardown awaits the in-flight one", async () => { + let releaseShutdown: (() => void) | undefined; + // Sentinel rather than a tick count: the harness shutdown joins the + // memoized teardown, 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 harnessSawFinishedTeardown: boolean | undefined; + const harnessPath = manager.shutdown({ exit: false }).then(() => { + harnessSawFinishedTeardown = teardownFinished; + }); + + // A full macrotask turn, so a harness shutdown that resolved early would. + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(harnessSawFinishedTeardown).toBeUndefined(); + + releaseShutdown?.(); + await Promise.all([signalPath, harnessPath]); + + expect(shutdown).toHaveBeenCalledTimes(1); + // It joined the in-flight teardown rather than resolving alongside it. + expect(harnessSawFinishedTeardown).toBe(true); + // The signal wanted the process dead, and still gets it. + expect(exitSpy).toHaveBeenCalledWith(0); + }); + + test("a rejecting plugin shutdown() is isolated and the harness shutdown 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.shutdown({ exit: false })).resolves.toBeUndefined(); + expect(mockLoggerError).toHaveBeenCalled(); + expect(exitSpy).not.toHaveBeenCalled(); + }); + + test("the harness path is bounded by the internal per-plugin timeout, not an outer one", async () => { + vi.useFakeTimers(); + // Never resolves on its own: the only bound is now the internal + // PLUGIN_SHUTDOWN_TIMEOUT_MS (10s) — the same one the signal path uses — + // since the harness path has no outer budget of its own any more. + const hanging = vi.fn(() => new Promise(() => {})); + const ctx = contextWithPlugins({ + stuck: { name: "stuck", shutdown: hanging } as never, + }); + const manager = new LifecycleManager(ctx); + + const done = manager.shutdown({ exit: false }); + await vi.advanceTimersByTimeAsync(10_000); + await expect(done).resolves.toBeUndefined(); + + expect(hanging).toHaveBeenCalledTimes(1); + expect( + mockLoggerError.mock.calls.some( + (c) => + String(c[0]).includes("Error shutting down plugin") && + c[1] === "stuck" && + String(c[2]).includes("timed out"), + ), + ).toBe(true); + // The harness path never exits, even when a phase times out internally. + expect(exitSpy).not.toHaveBeenCalled(); + }); }); }); diff --git a/packages/appkit/src/evals/dataset.ts b/packages/appkit/src/evals/dataset.ts index b6cdc0dac..f49270fcc 100644 --- a/packages/appkit/src/evals/dataset.ts +++ b/packages/appkit/src/evals/dataset.ts @@ -83,7 +83,7 @@ export async function readEvalDataset( : ""; const connector = new SQLWarehouseConnector({}); const response = await connector.executeStatement(client, { - warehouse_id: options.warehouseId, + warehouseId: options.warehouseId, statement: `SELECT inputs, expectations FROM ${options.table}${limit}`, }); diff --git a/packages/appkit/src/evals/tests/dataset.test.ts b/packages/appkit/src/evals/tests/dataset.test.ts index 12214d1ac..76d113b4d 100644 --- a/packages/appkit/src/evals/tests/dataset.test.ts +++ b/packages/appkit/src/evals/tests/dataset.test.ts @@ -45,7 +45,7 @@ describe("readEvalDataset", () => { // SELECT targets the table; no LIMIT when unset. const [, input] = executeStatement.mock.calls[0]; - expect(input.warehouse_id).toBe("wh1"); + expect(input.warehouseId).toBe("wh1"); expect(input.statement).toBe( "SELECT inputs, expectations FROM main.default.eval_ds", ); diff --git a/packages/appkit/src/plugin/tests/asUser-proxy.test.ts b/packages/appkit/src/plugin/tests/asUser-proxy.test.ts index 733cdc3b7..36f9e4f49 100644 --- a/packages/appkit/src/plugin/tests/asUser-proxy.test.ts +++ b/packages/appkit/src/plugin/tests/asUser-proxy.test.ts @@ -42,17 +42,6 @@ import type { ITelemetry, TelemetryProvider } from "../../telemetry"; import { TelemetryManager } from "../../telemetry"; import { isDevOboFallback, Plugin } from "../plugin"; -vi.mock("../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - ApiError: class extends Error { - statusCode = 500; - }, - }; -}); - vi.mock("../../app"); vi.mock("../../cache", () => ({ CacheManager: { getInstanceSync: vi.fn() }, diff --git a/packages/appkit/src/plugin/tests/plugin.test.ts b/packages/appkit/src/plugin/tests/plugin.test.ts index d42da93ad..64484fb0c 100644 --- a/packages/appkit/src/plugin/tests/plugin.test.ts +++ b/packages/appkit/src/plugin/tests/plugin.test.ts @@ -35,30 +35,10 @@ import { import { StreamManager } from "../../stream"; import type { ITelemetry, TelemetryProvider } from "../../telemetry"; import { TelemetryManager } from "../../telemetry"; +import { createApiError } from "../../testing"; import type { InterceptorContext } from "../interceptors/types"; import { isDevOboFallback, Plugin } from "../plugin"; -const { MockApiError } = vi.hoisted(() => { - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - return { MockApiError }; -}); - -vi.mock("../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - ApiError: MockApiError, - }; -}); - // Mock all dependencies vi.mock("../../app"); vi.mock("../../cache", () => ({ @@ -441,7 +421,11 @@ describe("Plugin", () => { test("should preserve 404 statusCode from ApiError (non-AppKitError)", async () => { const plugin = new TestPlugin(config); - const apiError = new MockApiError("Not found", 404); + const apiError = createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "ERROR", + }); const mockFn = vi.fn().mockRejectedValue(apiError); const result = await (plugin as any).execute( @@ -458,7 +442,11 @@ describe("Plugin", () => { test("should preserve 401 statusCode from ApiError (non-AppKitError)", async () => { const plugin = new TestPlugin(config); - const apiError = new MockApiError("Unauthorized", 401); + const apiError = createApiError({ + statusCode: 401, + message: "Unauthorized", + errorCode: "ERROR", + }); const mockFn = vi.fn().mockRejectedValue(apiError); const result = await (plugin as any).execute( @@ -475,7 +463,11 @@ describe("Plugin", () => { test("should preserve 403 statusCode from ApiError (non-AppKitError)", async () => { const plugin = new TestPlugin(config); - const apiError = new MockApiError("Forbidden", 403); + const apiError = createApiError({ + statusCode: 403, + message: "Forbidden", + errorCode: "ERROR", + }); const mockFn = vi.fn().mockRejectedValue(apiError); const result = await (plugin as any).execute( @@ -492,7 +484,11 @@ describe("Plugin", () => { test("should preserve 502 statusCode from non-AppKitError", async () => { const plugin = new TestPlugin(config); - const apiError = new MockApiError("Bad gateway", 502); + const apiError = createApiError({ + statusCode: 502, + message: "Bad gateway", + errorCode: "ERROR", + }); const mockFn = vi.fn().mockRejectedValue(apiError); const result = await (plugin as any).execute( @@ -512,7 +508,11 @@ describe("Plugin", () => { process.env.NODE_ENV = "production"; try { const plugin = new TestPlugin(config); - const apiError = new MockApiError("Internal upstream detail", 502); + const apiError = createApiError({ + statusCode: 502, + message: "Internal upstream detail", + errorCode: "ERROR", + }); const mockFn = vi.fn().mockRejectedValue(apiError); const result = await (plugin as any).execute( @@ -535,7 +535,11 @@ describe("Plugin", () => { process.env.NODE_ENV = "production"; try { const plugin = new TestPlugin(config); - const apiError = new MockApiError("Forbidden", 403); + const apiError = createApiError({ + statusCode: 403, + message: "Forbidden", + errorCode: "ERROR", + }); const mockFn = vi.fn().mockRejectedValue(apiError); const result = await (plugin as any).execute( 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 92c44a90a..0fe6fa67c 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 @@ -9,7 +9,7 @@ import { CacheManager } from "../../../cache"; import { resolveSkillCatalog } from "../../../core/agent/skills/resolve-catalog"; import type { SkillDefinition } from "../../../core/agent/skills/types"; import type { ResolvedToolEntry } from "../../../core/agent/types"; -import { createTestPluginContext } from "../../../testing"; +import { createMockRequest, createTestPluginContext } from "../../../testing"; import { AgentsPlugin } from "../agents"; import { dispatchToolCall, @@ -49,25 +49,15 @@ 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: (name: string) => headers[name.toLowerCase()], - } as unknown as express.Request; -} - function makeRunState(plugin: AgentsPlugin) { const abortController = new AbortController(); const pushed: unknown[] = []; const runState = { - req: mockReq(), + // `obo` carries the forwarded identity headers so executeTool's asUser(req) + // resolves a user scope (the mock context enforces the token precondition). + req: createMockRequest({ + obo: { token: "user-token", userId: "alice" }, + }) as unknown as express.Request, userId: "alice", requestId: "stream-1", abortController, 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 0eed1f884..db60da0cd 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,9 +2,29 @@ import type express from "express"; import { beforeEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../../cache"; -import { createTestPluginContext } from "../../../testing"; +import { createMockRequest, 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, @@ -21,6 +41,8 @@ import { AgentsPlugin } from "../agents"; */ beforeEach(() => { + linkTraceToRun.mockClear(); + mockTraceId = undefined; (CacheManager as any).instance = { get: vi.fn(), set: vi.fn(), @@ -33,15 +55,10 @@ beforeEach(() => { }); function mockReq(body: unknown, userId = "alice"): express.Request { - const headers: Record = { - "x-forwarded-user": userId, - "x-forwarded-access-token": "fake-token", - }; - return { + return createMockRequest({ body, - headers, - header: (name: string) => headers[name.toLowerCase()], - } as unknown as express.Request; + obo: { token: "fake-token", userId }, + }) as unknown as express.Request; } function mockRes() { @@ -61,12 +78,12 @@ function mockRes() { }; } -function seedPlugin(): AgentsPlugin { +function seedPlugin(adapter: unknown = { async *run() {} }): AgentsPlugin { const plugin = new AgentsPlugin({}); (plugin as any).agents.set("default", { name: "default", instructions: "hi", - adapter: { async *run() {} }, + adapter, toolIndex: new Map(), }); (plugin as any).defaultAgentName = "default"; @@ -374,6 +391,70 @@ 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("POST /invocations — tool failures surfaced", () => { diff --git a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts index 25ad528c9..9f8cf3d10 100644 --- a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -5,6 +5,8 @@ import { } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { withEnv } from "../../../testing"; +import { useTestCache } from "../../../testing/test-cache"; import { Context } from "../../../workspace-client"; vi.mock("../../../context", () => ({ @@ -52,8 +54,16 @@ vi.mock("../../../telemetry", () => ({ ) => fn({ setAttribute: vi.fn(), + setAttributes: vi.fn(), setStatus: vi.fn(), recordException: vi.fn(), + addEvent: vi.fn(), + addLink: vi.fn(), + addLinks: vi.fn(), + updateName: vi.fn(), + isRecording: vi.fn().mockReturnValue(false), + spanContext: vi.fn(), + end: vi.fn(), }), ), }), @@ -63,37 +73,9 @@ vi.mock("../../../telemetry", () => ({ normalizeTelemetryOptions: () => ({ traces: false, metrics: false }), })); -// In-memory cache keyed like the real CacheManager.generateKey, so tests -// exercise real key composition. Never stores rejections. -const { mockCacheStore } = vi.hoisted(() => ({ - mockCacheStore: new Map(), -})); - -vi.mock("../../../cache", () => { - const keyOf = (parts: unknown[], userKey: string) => - JSON.stringify([userKey, ...parts]); - return { - CacheManager: { - getInstanceSync: () => ({ - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - generateKey: keyOf, - getOrExecute: async ( - key: unknown[], - fn: (signal?: AbortSignal) => Promise, - userKey: string, - ) => { - const k = keyOf(key, userKey); - if (mockCacheStore.has(k)) return mockCacheStore.get(k); - const result = await fn(); - mockCacheStore.set(k, result); - return result; - }, - }), - }, - }; -}); +// Real in-memory cache so the plugin's caching path runs under test — no mock +// of the internal cache module. Boots and clears the cache before each test. +useTestCache(); vi.mock("../../../app", () => ({ AppManager: vi.fn().mockImplementation(() => ({})), @@ -139,7 +121,6 @@ describe("AiSearchPlugin", () => { beforeEach(() => { mockRequest.mockClear(); mockRequest.mockResolvedValue(validVsResponse); - mockCacheStore.clear(); }); describe("setup()", () => { @@ -206,31 +187,23 @@ describe("AiSearchPlugin", () => { }); it("throws outside development when an index has no columns", async () => { - const originalNodeEnv = process.env.NODE_ENV; - process.env.NODE_ENV = "production"; - try { + await withEnv({ NODE_ENV: "production" }, async () => { const plugin = new AiSearchPlugin({ indexes: { docs: { indexName: "cat.sch.idx" } }, }); await expect(plugin.setup()).rejects.toThrow( 'Index "docs" has no columns configured', ); - } finally { - process.env.NODE_ENV = originalNodeEnv; - } + }); }); it("does not throw outside development when columns are configured", async () => { - const originalNodeEnv = process.env.NODE_ENV; - process.env.NODE_ENV = "production"; - try { + await withEnv({ NODE_ENV: "production" }, async () => { const plugin = new AiSearchPlugin({ indexes: { docs: { indexName: "cat.sch.idx", columns: ["id"] } }, }); await expect(plugin.setup()).resolves.not.toThrow(); - } finally { - process.env.NODE_ENV = originalNodeEnv; - } + }); }); }); diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index dc3543be4..f95ea335d 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -28,7 +28,6 @@ import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; -import type { WorkspaceClient } from "../../workspace-client"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; import { @@ -1069,7 +1068,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { workspaceClient, { statement, - warehouse_id: warehouseId, + warehouseId, parameters: sqlParameters, ...formatParameters, }, diff --git a/packages/appkit/src/plugins/analytics/query.ts b/packages/appkit/src/plugins/analytics/query.ts index bcd77a817..4b57b051b 100644 --- a/packages/appkit/src/plugins/analytics/query.ts +++ b/packages/appkit/src/plugins/analytics/query.ts @@ -4,7 +4,7 @@ import { isSQLTypeMarker, type SQLTypeMarker, sql as sqlHelpers } from "shared"; import { getWorkspaceId } from "../../context"; import { ValidationError } from "../../errors"; -import type { sql } from "../../workspace-client"; +import type { StatementParameter } from "../../workspace-client"; type SQLParameterValue = SQLTypeMarker | null | undefined; @@ -37,8 +37,8 @@ export class QueryProcessor { convertToSQLParameters( query: string, parameters?: Record, - ): { statement: string; parameters: sql.StatementParameterListItem[] } { - const sqlParameters: sql.StatementParameterListItem[] = []; + ): { statement: string; parameters: StatementParameter[] } { + const sqlParameters: StatementParameter[] = []; if (parameters) { // extract all params from the query @@ -72,7 +72,7 @@ export class QueryProcessor { private _createParameter( key: string, value: SQLParameterValue, - ): sql.StatementParameterListItem | null { + ): StatementParameter | null { if (value === null || value === undefined) { return null; } diff --git a/packages/appkit/src/plugins/analytics/result-delivery.ts b/packages/appkit/src/plugins/analytics/result-delivery.ts index a0435513c..3f40439d2 100644 --- a/packages/appkit/src/plugins/analytics/result-delivery.ts +++ b/packages/appkit/src/plugins/analytics/result-delivery.ts @@ -4,7 +4,7 @@ import type { SQLTypeMarker } from "shared"; import { ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import type { RefreshChunkLink } from "../../stream/arrow-stream-processor"; -import type { sql } from "../../workspace-client"; +import type { ExternalLink } from "../../workspace-client"; /** * Centralized disposition/format fallback for analytics result delivery. @@ -39,7 +39,7 @@ export interface QueryExecutor { | { attachment?: string; data?: Record[]; - external_links?: sql.ExternalLink[]; + external_links?: ExternalLink[]; columnNames?: string[]; statement_id?: string; status?: unknown; @@ -52,7 +52,7 @@ export interface QueryExecutor { /** Streams already-resolved EXTERNAL_LINKS chunks; the connector provides it. */ export interface ArrowChunkStreamer { streamExternalLinks( - chunks: sql.ExternalLink[], + chunks: ExternalLink[], signal?: AbortSignal, refresh?: RefreshChunkLink, ): AsyncGenerator; 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..0265feeea 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, + getMock, parseSSEResponse, - setupDatabricksEnv, -} from "@tools/test-helpers"; + type TestApp, +} from "@databricks/appkit/testing"; import { sql } from "shared"; import { afterAll, @@ -20,85 +18,40 @@ 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 getStatementResult: 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. + app = await createTestApp({ plugins: [analytics({})] }); + executeStatement = getMock( + app.client, + "statementExecution.executeStatement", + ); + getStatementResult = getMock( + app.client, + "statementExecution.getStatementResult", + ); }); 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(); + getStatementResult.mockReset(); getAppQuerySpy.mockReset(); }); @@ -110,8 +63,8 @@ describe("Analytics Plugin Integration", () => { ["Bob", "25"], ]; const mockColumns = [ - { name: "name", type_name: "STRING" }, - { name: "age", type_name: "STRING" }, + { name: "name", typeName: "STRING" }, + { name: "age", typeName: "STRING" }, ]; getAppQuerySpy.mockResolvedValueOnce({ @@ -119,18 +72,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,11 +92,11 @@ 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", + warehouseId: "test-warehouse-id", }), expect.anything(), ); @@ -162,26 +110,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 +137,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 +156,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 +175,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 +196,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.readonly.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts index 68b6b94d7..758ac62f0 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.readonly.test.ts @@ -1,22 +1,11 @@ import { describe, expect, test, vi } from "vitest"; -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => ({ - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_k: unknown[], fn: (signal?: AbortSignal) => Promise) => - fn(), - ), - generateKey: vi.fn(() => "test-key"), - })), - }, -})); - +import { useTestCache } from "../../../testing/test-cache"; import { AnalyticsPlugin } from "../analytics"; +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); + /** * Tests the read-only SQL enforcement on the analytics agent tool. * diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 5101f9424..a30e4be8e 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -22,48 +22,13 @@ import { sql } from "shared"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ServiceContext } from "../../../context/service-context"; +import { useTestCache } from "../../../testing/test-cache"; import { AnalyticsPlugin, analytics, writeChunk } from "../analytics"; import type { IAnalyticsConfig } from "../types"; -// Mock CacheManager singleton with actual caching behavior -const { mockCacheStore, mockCacheInstance } = vi.hoisted(() => { - const store = new Map(); - - const generateKey = (parts: unknown[], userKey: string): string => { - const { createHash } = require("node:crypto"); - const allParts = [userKey, ...parts]; - const serialized = JSON.stringify(allParts); - return createHash("sha256").update(serialized).digest("hex"); - }; - - const instance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (key: unknown[], fn: () => Promise, userKey: string) => { - const cacheKey = generateKey(key, userKey); - if (store.has(cacheKey)) { - return store.get(cacheKey); - } - const result = await fn(); - store.set(cacheKey, result); - return result; - }, - ), - generateKey: vi.fn((parts: unknown[], userKey: string) => - generateKey(parts, userKey), - ), - }; - - return { mockCacheStore: store, mockCacheInstance: instance }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Real in-memory cache so the plugin's caching path runs under test — no mock +// of the internal cache module. Boots and clears the cache before each test. +useTestCache(); describe("Analytics Plugin", () => { let config: IAnalyticsConfig; @@ -72,7 +37,6 @@ describe("Analytics Plugin", () => { beforeEach(async () => { config = { timeout: 5000 }; setupDatabricksEnv(); - mockCacheStore.clear(); ServiceContext.reset(); serviceContextMock = await mockServiceContext(); }); @@ -173,7 +137,7 @@ describe("Analytics Plugin", () => { expect.anything(), expect.objectContaining({ statement: "SELECT * FROM test", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -241,7 +205,7 @@ describe("Analytics Plugin", () => { expect.anything(), expect.objectContaining({ statement: "SELECT * FROM users WHERE id = :user_id", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -638,7 +602,7 @@ describe("Analytics Plugin", () => { expect.objectContaining({ statement: "SELECT * FROM test", parameters: [], - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -673,7 +637,7 @@ describe("Analytics Plugin", () => { expect.anything(), expect.objectContaining({ statement: "SELECT * FROM test", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", disposition: "INLINE", format: "ARROW_STREAM", }), @@ -1577,14 +1541,17 @@ describe("Analytics Plugin", () => { isAsUser: false, }); - const executeMock = vi.fn().mockImplementation((_wc, _opts, signal) => { - // Simulate a signal that becomes aborted before the failure surfaces — - // e.g. the client cancelled the SSE stream mid-query. Use vitest's - // getter spy rather than Object.defineProperty so we don't try to - // override the native non-configurable AbortSignal.aborted getter. - if (signal) { - vi.spyOn(signal, "aborted", "get").mockReturnValue(true); - } + const mockRes = createMockResponse(); + + const executeMock = vi.fn().mockImplementation(() => { + // Simulate the client cancelling mid-query: firing the response's + // "close" event aborts the handler's own AbortController (see + // `onClose` in `_handleArrowStreamQuery`), exactly as a real disconnect + // would. Modelling the abort on the handler signal — the one the + // fallback guard actually checks — rather than on whatever signal + // reaches executeStatement keeps this independent of how the cache + // threads its shared signal into the inner fn. + mockRes.end(); return Promise.reject( new Error( "INVALID_PARAMETER_VALUE: ARROW_STREAM not supported with INLINE disposition", @@ -1600,12 +1567,11 @@ describe("Analytics Plugin", () => { params: { query_key: "test_query" }, body: { parameters: {}, format: "ARROW_STREAM" }, }); - const mockRes = createMockResponse(); await handler(mockReq, mockRes); - // Even though the error message would normally trigger fallback, the - // aborted signal should short-circuit and prevent a second statement. + // The aborted request short-circuits the INLINE→EXTERNAL_LINKS fallback, + // so exactly one statement runs. expect(executeMock).toHaveBeenCalledTimes(1); }); @@ -1682,7 +1648,7 @@ describe("Analytics Plugin", () => { result: { data: [] }, }), }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ @@ -1708,7 +1674,6 @@ describe("Analytics Plugin", () => { expect(resultIdx).toBeGreaterThanOrEqual(0); expect(warehouseIdx).toBeLessThan(resultIdx); - // The status payload should include the RUNNING state. expect(mockRes.write).toHaveBeenCalledWith( expect.stringMatching(/"type":"warehouse_status".*"state":"RUNNING"/), ); diff --git a/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts index 934e83c28..404719440 100644 --- a/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/arrow-delivery.integration.test.ts @@ -38,9 +38,9 @@ describe.runIf(!!warehouseId)("arrow delivery (live warehouse)", () => { client, { statement, - warehouse_id: warehouseId as string, - wait_timeout: "50s", - on_wait_timeout: "CONTINUE", + warehouseId: warehouseId as string, + waitTimeout: "50s", + onWaitTimeout: "CONTINUE", disposition: fp.disposition as never, format: fp.format as never, }, diff --git a/packages/appkit/src/plugins/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index bf721feee..63cb7ea27 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -15,6 +15,7 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AppManager } from "../../../app"; import { ServiceContext } from "../../../context/service-context"; import { AuthenticationError } from "../../../errors"; +import { useTestCache } from "../../../testing/test-cache"; import { AnalyticsPlugin } from "../analytics"; import { buildMetricSql, @@ -31,40 +32,9 @@ import type { MetricRegistration, } from "../types"; -// Mirror the analytics.test.ts CacheManager mock so the inner `execute`'s -// cache interceptor is a no-op pass-through (each request re-executes). -const { mockCacheStore, mockCacheInstance } = vi.hoisted(() => { - const store = new Map(); - const generateKey = (parts: unknown[], userKey: string): string => { - const { createHash } = require("node:crypto"); - const serialized = JSON.stringify([userKey, ...parts]); - return createHash("sha256").update(serialized).digest("hex"); - }; - const instance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (key: unknown[], fn: () => Promise, userKey: string) => { - const cacheKey = generateKey(key, userKey); - if (store.has(cacheKey)) return store.get(cacheKey); - const result = await fn(); - store.set(cacheKey, result); - return result; - }, - ), - generateKey: vi.fn((parts: unknown[], userKey: string) => - generateKey(parts, userKey), - ), - }; - return { mockCacheStore: store, mockCacheInstance: instance }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Real in-memory cache so the metric route's cache interceptor runs under test +// — no mock of the internal cache module. Boots and clears it before each test. +const testCache = useTestCache(); // Temp dirs created by `registryDir` / `writeRegistry`, cleaned up after each // test. Using real files (pointing the plugin's `AppManager` at the dir, see @@ -146,7 +116,6 @@ describe("analytics metric route", () => { beforeEach(async () => { config = { timeout: 5000 }; setupDatabricksEnv(); - mockCacheStore.clear(); ServiceContext.reset(); serviceContextMock = await mockServiceContext(); }); @@ -731,7 +700,7 @@ describe("analytics metric route", () => { expect.objectContaining({ statement: "SELECT MEASURE(`arr`) AS `arr` FROM `cat`.`sch`.`revenue_metrics`", - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.any(AbortSignal), ); @@ -779,7 +748,7 @@ describe("analytics metric route", () => { result: { data: [] }, }), }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ @@ -1064,8 +1033,12 @@ describe("analytics metric route", () => { // Capture the composed cache key the inner `execute` hands to the shared // CacheManager mock — the same key whether or not metadata is injected. + // Spy the real cache's getOrExecute to capture the composed key parts the + // metric route's cache interceptor passes — the same whether or not + // metadata is injected. + const getOrExecuteSpy = vi.spyOn(testCache.current, "getOrExecute"); const cacheKeyFor = async (mvMeta?: MetricViewsMetadata) => { - mockCacheInstance.getOrExecute.mockClear(); + getOrExecuteSpy.mockClear(); const plugin = pluginForDir( { ...config, metricViewsMetadata: mvMeta }, registryDir(registry), @@ -1079,12 +1052,13 @@ describe("analytics metric route", () => { createMockResponse(), ); // First getOrExecute call is the SQL execution's cache interceptor. - const call = mockCacheInstance.getOrExecute.mock.calls[0]; + const call = getOrExecuteSpy.mock.calls[0]; return { cacheKey: call[0], userKey: call[2] }; }; const withMeta = await cacheKeyFor(REVENUE_METADATA); const withoutMeta = await cacheKeyFor(undefined); + getOrExecuteSpy.mockRestore(); expect(withMeta.cacheKey).toEqual(withoutMeta.cacheKey); expect(withMeta.userKey).toEqual(withoutMeta.userKey); @@ -2489,7 +2463,6 @@ describe("metric — filter translator", () => { beforeEach(async () => { config = { timeout: 5000 }; setupDatabricksEnv(); - mockCacheStore.clear(); ServiceContext.reset(); serviceContextMock = await mockServiceContext(); }); @@ -2960,7 +2933,6 @@ describe("metric route — lane dispatch", () => { beforeEach(async () => { config = { timeout: 5000 }; setupDatabricksEnv(); - mockCacheStore.clear(); ServiceContext.reset(); serviceContextMock = await mockServiceContext(); }); diff --git a/packages/appkit/src/plugins/files/tests/_test-helpers.ts b/packages/appkit/src/plugins/files/tests/_test-helpers.ts index 1531ce1a4..b559e2722 100644 --- a/packages/appkit/src/plugins/files/tests/_test-helpers.ts +++ b/packages/appkit/src/plugins/files/tests/_test-helpers.ts @@ -134,13 +134,19 @@ export function makeStreamResponse(content: string) { return { contents: stream }; } -export async function setupTestEnv() { +export async function setupTestEnv(client?: unknown) { vi.clearAllMocks(); setupDatabricksEnv(); ServiceContext.reset(); process.env.DATABRICKS_VOLUME_UPLOADS = "/Volumes/catalog/schema/uploads"; process.env.DATABRICKS_VOLUME_EXPORTS = "/Volumes/catalog/schema/exports"; - return mockServiceContext(); + // Injecting a client makes `getWorkspaceClient()` resolve to it through the + // real ServiceContext, so a suite needs no vi.mock of `../../../context`. + return mockServiceContext( + client + ? { serviceDatabricksClient: client, userDatabricksClient: client } + : {}, + ); } export function teardownTestEnv( diff --git a/packages/appkit/src/plugins/files/tests/delete.test.ts b/packages/appkit/src/plugins/files/tests/delete.test.ts index 1cdaa7e5d..9ef00f942 100644 --- a/packages/appkit/src/plugins/files/tests/delete.test.ts +++ b/packages/appkit/src/plugins/files/tests/delete.test.ts @@ -1,5 +1,11 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + createApiError, + createMockWorkspaceClient, + getMock, + useTestCache, +} from "../../../testing"; import { FilesPlugin } from "../plugin"; import { getRouteHandler, @@ -10,73 +16,28 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Real in-memory cache; spy on `testCache.current` to assert the plugin's +// cache-invalidation calls. +const testCache = useTestCache(); describe("FilesPlugin delete", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -88,7 +49,10 @@ describe("FilesPlugin delete", () => { const handler = getRouteHandler(plugin, "delete", ""); const res = mockRes(); - mockClient.files.delete.mockResolvedValue(undefined); + const generateKey = vi.spyOn(testCache.current, "generateKey"); + const del = vi.spyOn(testCache.current, "delete"); + + getMock(client, "files.delete").mockResolvedValue(undefined); await handler( mockReq("uploads", { @@ -100,8 +64,8 @@ describe("FilesPlugin delete", () => { expect(res.json).toHaveBeenCalledWith( expect.objectContaining({ success: true }), ); - expect(mockCacheInstance.generateKey).toHaveBeenCalled(); - expect(mockCacheInstance.delete).toHaveBeenCalled(); + expect(generateKey).toHaveBeenCalled(); + expect(del).toHaveBeenCalled(); }); test("delete without path returns 400", async () => { @@ -122,8 +86,12 @@ describe("FilesPlugin delete", () => { const handler = getRouteHandler(plugin, "delete", ""); const res = mockRes(); - mockClient.files.delete.mockRejectedValue( - new MockApiError("Not found", 404), + getMock(client, "files.delete").mockRejectedValue( + createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "ERROR", + }), ); await handler( diff --git a/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts b/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts index e96470e1b..622aa702d 100644 --- a/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts +++ b/packages/appkit/src/plugins/files/tests/download-endpoint.test.ts @@ -1,5 +1,10 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import { + createMockWorkspaceClient, + getMock, + useTestCache, +} from "../../../testing"; import { FilesPlugin } from "../plugin"; import { getRouteHandler, @@ -11,73 +16,27 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("FilesPlugin download endpoint Content-Disposition", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -89,7 +48,7 @@ describe("FilesPlugin download endpoint Content-Disposition", () => { const handler = getRouteHandler(plugin, "get", "/download"); const res = mockRes(); - mockClient.files.download.mockResolvedValue( + getMock(client, "files.download").mockResolvedValue( makeStreamResponse("file data"), ); @@ -111,7 +70,9 @@ describe("FilesPlugin download endpoint Content-Disposition", () => { const handler = getRouteHandler(plugin, "get", "/download"); const res = mockRes(); - mockClient.files.download.mockResolvedValue(makeStreamResponse("data")); + getMock(client, "files.download").mockResolvedValue( + makeStreamResponse("data"), + ); await handler( mockReq("uploads", { @@ -131,7 +92,9 @@ describe("FilesPlugin download endpoint Content-Disposition", () => { const handler = getRouteHandler(plugin, "get", "/download"); const res = mockRes(); - mockClient.files.download.mockResolvedValue(makeStreamResponse("{}")); + getMock(client, "files.download").mockResolvedValue( + makeStreamResponse("{}"), + ); await handler( mockReq("uploads", { @@ -169,7 +132,7 @@ describe("FilesPlugin download endpoint Content-Disposition", () => { const res = mockRes(); // Response with no contents field (empty file) - mockClient.files.download.mockResolvedValue({}); + getMock(client, "files.download").mockResolvedValue({}); await handler( mockReq("uploads", { diff --git a/packages/appkit/src/plugins/files/tests/error-handling.test.ts b/packages/appkit/src/plugins/files/tests/error-handling.test.ts index f02fb8683..296b6f2b8 100644 --- a/packages/appkit/src/plugins/files/tests/error-handling.test.ts +++ b/packages/appkit/src/plugins/files/tests/error-handling.test.ts @@ -1,6 +1,12 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test } from "vitest"; import { AuthenticationError } from "../../../errors"; +import { + createApiError, + createMockWorkspaceClient, + useTestCache, +} from "../../../testing"; +import { withEnv } from "../../../testing"; import { FilesPlugin } from "../plugin"; import { getRouteHandler, @@ -10,73 +16,27 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("FilesPlugin error handling", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -107,7 +67,11 @@ describe("FilesPlugin error handling", () => { (plugin as any)._handleApiError( res, - new MockApiError("Forbidden", 403), + createApiError({ + statusCode: 403, + message: "Forbidden", + errorCode: "ERROR", + }), "fallback msg", ); @@ -125,7 +89,11 @@ describe("FilesPlugin error handling", () => { (plugin as any)._handleApiError( res, - new MockApiError("Not found", 404), + createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "ERROR", + }), "fallback msg", ); @@ -143,7 +111,11 @@ describe("FilesPlugin error handling", () => { (plugin as any)._handleApiError( res, - new MockApiError("Conflict", 409), + createApiError({ + statusCode: 409, + message: "Conflict", + errorCode: "ERROR", + }), "fallback msg", ); @@ -161,7 +133,11 @@ describe("FilesPlugin error handling", () => { (plugin as any)._handleApiError( res, - new MockApiError("Bad Gateway", 502), + createApiError({ + statusCode: 502, + message: "Bad Gateway", + errorCode: "ERROR", + }), "Operation failed", ); @@ -178,7 +154,11 @@ describe("FilesPlugin error handling", () => { (plugin as any)._handleApiError( res, - new MockApiError("Internal error", 500), + createApiError({ + statusCode: 500, + message: "Internal error", + errorCode: "ERROR", + }), "Fallback", ); @@ -220,38 +200,37 @@ describe("FilesPlugin error handling", () => { }); test("AuthenticationError via route returns generic 401 on OBO volume without token", async () => { - process.env.DATABRICKS_VOLUME_OBO = "/Volumes/catalog/schema/obo"; - const plugin = new FilesPlugin({ - volumes: { - obo: { auth: "on-behalf-of-user", policy: () => true }, + await withEnv( + { + DATABRICKS_VOLUME_OBO: "/Volumes/catalog/schema/obo", + NODE_ENV: "production", }, - }); - const handler = getRouteHandler(plugin, "get", "/list"); - const res = mockRes(); - - const originalEnv = process.env.NODE_ENV; - process.env.NODE_ENV = "production"; - - try { - await handler( - { - params: { volumeKey: "obo" }, - query: {}, - headers: {}, - header: () => undefined, - }, - res, - ); - - expect(res.status).toHaveBeenCalledWith(401); - expect(res.json).toHaveBeenCalledWith({ - error: "Unauthorized", - plugin: "files", - }); - } finally { - process.env.NODE_ENV = originalEnv; - delete process.env.DATABRICKS_VOLUME_OBO; - } + async () => { + const plugin = new FilesPlugin({ + volumes: { + obo: { auth: "on-behalf-of-user", policy: () => true }, + }, + }); + const handler = getRouteHandler(plugin, "get", "/list"); + const res = mockRes(); + + await handler( + { + params: { volumeKey: "obo" }, + query: {}, + headers: {}, + header: () => undefined, + }, + res, + ); + + expect(res.status).toHaveBeenCalledWith(401); + expect(res.json).toHaveBeenCalledWith({ + error: "Unauthorized", + plugin: "files", + }); + }, + ); }); }); diff --git a/packages/appkit/src/plugins/files/tests/mkdir.test.ts b/packages/appkit/src/plugins/files/tests/mkdir.test.ts index 00623bef5..e2ecdee10 100644 --- a/packages/appkit/src/plugins/files/tests/mkdir.test.ts +++ b/packages/appkit/src/plugins/files/tests/mkdir.test.ts @@ -1,5 +1,11 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { + createApiError, + createMockWorkspaceClient, + getMock, + useTestCache, +} from "../../../testing"; import { FilesPlugin } from "../plugin"; import { getRouteHandler, @@ -10,73 +16,21 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Real in-memory cache; spy on `testCache.current` to assert invalidation. +const testCache = useTestCache(); describe("FilesPlugin mkdir", () => { let serviceContextMock: Awaited>; + let client: ReturnType; beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + // strict: true keeps the loudness the hand-rolled literal had by accident — + // an undeclared data-plane call throws instead of resolving undefined. + client = createMockWorkspaceClient({ + strict: true, + responses: { "files.createDirectory": undefined }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -88,7 +42,8 @@ describe("FilesPlugin mkdir", () => { const handler = getRouteHandler(plugin, "post", "/mkdir"); const res = mockRes(); - mockClient.files.createDirectory.mockResolvedValue(undefined); + const generateKey = vi.spyOn(testCache.current, "generateKey"); + const del = vi.spyOn(testCache.current, "delete"); await handler( mockReq("uploads", { @@ -100,8 +55,9 @@ describe("FilesPlugin mkdir", () => { expect(res.json).toHaveBeenCalledWith( expect.objectContaining({ success: true }), ); - expect(mockCacheInstance.generateKey).toHaveBeenCalled(); - expect(mockCacheInstance.delete).toHaveBeenCalled(); + expect(getMock(client, "files.createDirectory")).toHaveBeenCalled(); + expect(generateKey).toHaveBeenCalled(); + expect(del).toHaveBeenCalled(); }); test("mkdir without path returns 400", async () => { @@ -122,8 +78,12 @@ describe("FilesPlugin mkdir", () => { const handler = getRouteHandler(plugin, "post", "/mkdir"); const res = mockRes(); - mockClient.files.createDirectory.mockRejectedValue( - new MockApiError("Conflict", 409), + getMock(client, "files.createDirectory").mockRejectedValue( + createApiError({ + statusCode: 409, + message: "Conflict", + errorCode: "ALREADY_EXISTS", + }), ); await handler( diff --git a/packages/appkit/src/plugins/files/tests/path-validation.test.ts b/packages/appkit/src/plugins/files/tests/path-validation.test.ts index 7705b4778..42e21b883 100644 --- a/packages/appkit/src/plugins/files/tests/path-validation.test.ts +++ b/packages/appkit/src/plugins/files/tests/path-validation.test.ts @@ -1,5 +1,10 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import { + createMockWorkspaceClient, + getMock, + useTestCache, +} from "../../../testing"; import { FilesPlugin } from "../plugin"; import { getRouteHandler, @@ -11,73 +16,27 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("FilesPlugin path validation", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -87,12 +46,14 @@ describe("FilesPlugin path validation", () => { // Defends against regressions where the handler calls the SDK and *also* // returns 400 — the status assertion alone wouldn't catch that. function expectNoSdkCall() { - expect(mockClient.files.download).not.toHaveBeenCalled(); - expect(mockClient.files.upload).not.toHaveBeenCalled(); - expect(mockClient.files.delete).not.toHaveBeenCalled(); - expect(mockClient.files.createDirectory).not.toHaveBeenCalled(); - expect(mockClient.files.getMetadata).not.toHaveBeenCalled(); - expect(mockClient.files.listDirectoryContents).not.toHaveBeenCalled(); + expect(getMock(client, "files.download")).not.toHaveBeenCalled(); + expect(getMock(client, "files.upload")).not.toHaveBeenCalled(); + expect(getMock(client, "files.delete")).not.toHaveBeenCalled(); + expect(getMock(client, "files.createDirectory")).not.toHaveBeenCalled(); + expect(getMock(client, "files.getMetadata")).not.toHaveBeenCalled(); + expect( + getMock(client, "files.listDirectoryContents"), + ).not.toHaveBeenCalled(); } test("path with null bytes returns 400", async () => { 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..2fd955297 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, @@ -13,11 +17,12 @@ import { import { ServiceContext } from "../../../context/service-context"; import { createApp } from "../../../core"; +import { createApiError } from "../../../testing"; import { server as serverPlugin } from "../../server"; import { files } from "../index"; import { streamFromString } from "./utils"; -const { mockFilesApi, mockSdkClient, MockApiError } = vi.hoisted(() => { +const { mockFilesApi, mockSdkClient } = vi.hoisted(() => { const mockFilesApi = { listDirectoryContents: vi.fn(), download: vi.fn(), @@ -38,25 +43,7 @@ const { mockFilesApi, mockSdkClient, MockApiError } = vi.hoisted(() => { }, }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - - return { mockFilesApi, mockSdkClient, MockApiError }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - ApiError: MockApiError, - }; + return { mockFilesApi, mockSdkClient }; }); const MOCK_AUTH_HEADERS = { @@ -67,29 +54,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; @@ -265,7 +229,11 @@ describe("Files Plugin Integration", () => { test(`GET /api/files/${VOL}/exists returns { exists: false } on 404`, async () => { mockFilesApi.getMetadata.mockRejectedValue( - new MockApiError("Not found", 404), + createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "ERROR", + }), ); const response = await fetch( @@ -696,7 +664,11 @@ describe("Files Plugin Integration", () => { test("ApiError 404 preserves upstream status code", async () => { mockFilesApi.getMetadata.mockRejectedValue( - new MockApiError("Not found", 404), + createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "ERROR", + }), ); const response = await fetch( @@ -715,7 +687,11 @@ describe("Files Plugin Integration", () => { test("ApiError 409 preserves upstream status code", async () => { mockFilesApi.getMetadata.mockRejectedValue( - new MockApiError("Conflict", 409), + createApiError({ + statusCode: 409, + message: "Conflict", + errorCode: "ERROR", + }), ); const response = await fetch( diff --git a/packages/appkit/src/plugins/files/tests/plugin.test.ts b/packages/appkit/src/plugins/files/tests/plugin.test.ts index a9612d47e..a6736d307 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.test.ts @@ -7,6 +7,7 @@ import { ServiceContext } from "../../../context/service-context"; import { createApp } from "../../../core"; import { AuthenticationError } from "../../../errors"; import { ResourceType } from "../../../registry"; +import { useTestCache, withEnv } from "../../../testing"; import { FILES_DOWNLOAD_DEFAULTS, FILES_READ_DEFAULTS, @@ -15,7 +16,7 @@ import { import { FilesPlugin, files } from "../plugin"; import { PolicyDeniedError, policy } from "../policy"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { +const { mockClient, MockApiError } = await vi.hoisted(async () => { const mockFilesApi = { listDirectoryContents: vi.fn(), download: vi.fn(), @@ -42,18 +43,7 @@ const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { } } - 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 { mockFilesApi, mockClient, MockApiError, mockCacheInstance }; + return { mockFilesApi, mockClient, MockApiError }; }); vi.mock("../../../workspace-client", async (importOriginal) => { @@ -75,13 +65,6 @@ vi.mock("../../../context", async (importOriginal) => { }; }); -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - getInstance: vi.fn(async () => mockCacheInstance), - }, -})); - const VOLUMES_CONFIG = { volumes: { uploads: { maxUploadSize: 100_000_000, policy: policy.allowAll() }, @@ -89,6 +72,10 @@ const VOLUMES_CONFIG = { }, }; +// Boots AppKit's real in-memory cache; spy on `testCache.current` to assert +// cache behaviour. No mock of the internal cache module. +const testCache = useTestCache(); + describe("FilesPlugin", () => { let serviceContextMock: Awaited>; @@ -137,33 +124,27 @@ describe("FilesPlugin", () => { }); test("skips bare DATABRICKS_VOLUME_ prefix (no suffix)", () => { - process.env.DATABRICKS_VOLUME_ = "/Volumes/bare"; - try { + withEnv({ DATABRICKS_VOLUME_: "/Volumes/bare" }, () => { const volumes = FilesPlugin.discoverVolumes({}); expect(Object.keys(volumes)).not.toContain(""); - } finally { - delete process.env.DATABRICKS_VOLUME_; - } + }); }); test("skips empty env var values", () => { - process.env.DATABRICKS_VOLUME_EMPTY = ""; - try { + withEnv({ DATABRICKS_VOLUME_EMPTY: "" }, () => { const volumes = FilesPlugin.discoverVolumes({}); expect(volumes).not.toHaveProperty("empty"); - } finally { - delete process.env.DATABRICKS_VOLUME_EMPTY; - } + }); }); test("lowercases env var suffix", () => { - process.env.DATABRICKS_VOLUME_MY_DATA = "/Volumes/catalog/schema/data"; - try { - const volumes = FilesPlugin.discoverVolumes({}); - expect(volumes).toHaveProperty("my_data"); - } finally { - delete process.env.DATABRICKS_VOLUME_MY_DATA; - } + withEnv( + { DATABRICKS_VOLUME_MY_DATA: "/Volumes/catalog/schema/data" }, + () => { + const volumes = FilesPlugin.discoverVolumes({}); + expect(volumes).toHaveProperty("my_data"); + }, + ); }); test("returns only explicit volumes when no env vars match", () => { @@ -2448,6 +2429,8 @@ describe("FilesPlugin", () => { }, ); + const getOrExecute = vi.spyOn(testCache.current, "getOrExecute"); + // Alice's request. await handler( mockReq("obo_vol", { @@ -2468,7 +2451,7 @@ describe("FilesPlugin", () => { // Cache is disabled on OBO: `getOrExecute` is bypassed. The SDK // must execute on every request — no cross-user staleness possible. - expect(mockCacheInstance.getOrExecute).not.toHaveBeenCalled(); + expect(getOrExecute).not.toHaveBeenCalled(); expect(mockClient.files.listDirectoryContents).toHaveBeenCalledTimes(2); }); @@ -2494,6 +2477,8 @@ describe("FilesPlugin", () => { }, ); + const getOrExecute = vi.spyOn(testCache.current, "getOrExecute"); + // SP volume request — must consult the cache (cache enabled). await listHandler( mockReq("uploads", { @@ -2512,7 +2497,7 @@ describe("FilesPlugin", () => { mockRes(), ); - const calls = mockCacheInstance.getOrExecute.mock.calls; + const calls = getOrExecute.mock.calls; // Exactly one cache consultation — the SP volume's. The OBO request // bypassed the cache entirely. expect(calls).toHaveLength(1); @@ -2994,18 +2979,9 @@ describe("FilesPlugin", () => { mockClient.files.createDirectory.mockResolvedValue(undefined); - // Track which (parts, userKey) pairs go through generateKey so we - // can match the invalidation segment exactly. - const generateKeyCalls: Array<{ - parts: (string | number | object)[]; - userKey: string; - }> = []; - mockCacheInstance.generateKey.mockImplementation( - (parts: (string | number | object)[], userKey: string) => { - generateKeyCalls.push({ parts, userKey }); - return "stub-key"; - }, - ); + // Track which (parts, userKey) pairs go through the real generateKey + // so we can match the invalidation segment exactly. + const generateKey = vi.spyOn(testCache.current, "generateKey"); await mkdirHandler( mockReq("sp_vol", {}, { body: { path: "/Volumes/c/s/sp/foo/bar" } }), @@ -3013,9 +2989,11 @@ describe("FilesPlugin", () => { ); // Exactly one list-cache invalidation key was constructed. - const listInvalidations = generateKeyCalls.filter( - (c) => Array.isArray(c.parts) && c.parts[0] === "files:sp_vol:list", - ); + const listInvalidations = generateKey.mock.calls + .map((c) => ({ parts: c[0], userKey: c[1] })) + .filter( + (c) => Array.isArray(c.parts) && c.parts[0] === "files:sp_vol:list", + ); expect(listInvalidations).toHaveLength(1); // The path-segment is the PARENT directory (resolved), not the @@ -3067,25 +3045,19 @@ describe("FilesPlugin", () => { mockClient.files.createDirectory.mockResolvedValue(undefined); - const generateKeyCalls: Array<{ - parts: (string | number | object)[]; - userKey: string; - }> = []; - mockCacheInstance.generateKey.mockImplementation( - (parts: (string | number | object)[], userKey: string) => { - generateKeyCalls.push({ parts, userKey }); - return "stub-key"; - }, - ); + const generateKey = vi.spyOn(testCache.current, "generateKey"); await mkdirHandler( mockReq("uploads", {}, { body: { path: writePath } }), mockRes(), ); - const listInvalidations = generateKeyCalls.filter( - (c) => Array.isArray(c.parts) && c.parts[0] === "files:uploads:list", - ); + const listInvalidations = generateKey.mock.calls + .map((c) => ({ parts: c[0], userKey: c[1] })) + .filter( + (c) => + Array.isArray(c.parts) && c.parts[0] === "files:uploads:list", + ); const segments = listInvalidations.map((c) => c.parts[1]); expect(segments).toEqual( expect.arrayContaining([ @@ -3144,7 +3116,6 @@ describe("FilesPlugin", () => { const mkdirHandler = getRouteHandler(plugin, "post", "/mkdir"); mockClient.files.createDirectory.mockResolvedValue(undefined); - mockCacheInstance.generateKey.mockReturnValue("stub-key"); // Deferred promise that gates the cache delete. The handler must // await this before writing the success response. @@ -3152,9 +3123,9 @@ describe("FilesPlugin", () => { const deletePending = new Promise((resolve) => { releaseDelete = resolve; }); - mockCacheInstance.delete.mockImplementation( - async () => await deletePending, - ); + const del = vi + .spyOn(testCache.current, "delete") + .mockImplementation(async () => await deletePending); const res = mockRes(); @@ -3172,15 +3143,12 @@ describe("FilesPlugin", () => { // Use setImmediate to also drain macrotask queue items (telemetry/ // timeout interceptors may use setTimeout under the hood). const deadline = Date.now() + 1000; - while ( - mockCacheInstance.delete.mock.calls.length === 0 && - Date.now() < deadline - ) { + while (del.mock.calls.length === 0 && Date.now() < deadline) { await new Promise((resolve) => setImmediate(resolve)); } expect(mockClient.files.createDirectory).toHaveBeenCalledTimes(1); - expect(mockCacheInstance.delete).toHaveBeenCalledTimes(1); + expect(del).toHaveBeenCalledTimes(1); // Critical assertion: drain plenty of microtasks AND macrotasks // while `cache.delete` is still parked on the deferred. If the diff --git a/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts b/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts index 5614e8b37..9a2e1587a 100644 --- a/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts +++ b/packages/appkit/src/plugins/files/tests/raw-endpoint.test.ts @@ -1,5 +1,10 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import { + createMockWorkspaceClient, + getMock, + useTestCache, +} from "../../../testing"; import { FilesPlugin } from "../plugin"; import { getRouteHandler, @@ -11,73 +16,27 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("FilesPlugin raw endpoint security headers", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -89,7 +48,9 @@ describe("FilesPlugin raw endpoint security headers", () => { const handler = getRouteHandler(plugin, "get", "/raw"); const res = mockRes(); - mockClient.files.download.mockResolvedValue(makeStreamResponse("data")); + getMock(client, "files.download").mockResolvedValue( + makeStreamResponse("data"), + ); await handler( mockReq("uploads", { @@ -109,7 +70,9 @@ describe("FilesPlugin raw endpoint security headers", () => { const handler = getRouteHandler(plugin, "get", "/raw"); const res = mockRes(); - mockClient.files.download.mockResolvedValue(makeStreamResponse("PNG data")); + getMock(client, "files.download").mockResolvedValue( + makeStreamResponse("PNG data"), + ); await handler( mockReq("uploads", { @@ -135,7 +98,7 @@ describe("FilesPlugin raw endpoint security headers", () => { const handler = getRouteHandler(plugin, "get", "/raw"); const res = mockRes(); - mockClient.files.download.mockResolvedValue( + getMock(client, "files.download").mockResolvedValue( makeStreamResponse(""), ); @@ -162,7 +125,7 @@ describe("FilesPlugin raw endpoint security headers", () => { const handler = getRouteHandler(plugin, "get", "/raw"); const res = mockRes(); - mockClient.files.download.mockResolvedValue( + getMock(client, "files.download").mockResolvedValue( makeStreamResponse(""), ); @@ -184,7 +147,9 @@ describe("FilesPlugin raw endpoint security headers", () => { const handler = getRouteHandler(plugin, "get", "/raw"); const res = mockRes(); - mockClient.files.download.mockResolvedValue(makeStreamResponse("content")); + getMock(client, "files.download").mockResolvedValue( + makeStreamResponse("content"), + ); await handler( mockReq("uploads", { diff --git a/packages/appkit/src/plugins/files/tests/shutdown.test.ts b/packages/appkit/src/plugins/files/tests/shutdown.test.ts index 239d92cff..31edbb8d0 100644 --- a/packages/appkit/src/plugins/files/tests/shutdown.test.ts +++ b/packages/appkit/src/plugins/files/tests/shutdown.test.ts @@ -1,75 +1,30 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { createMockWorkspaceClient, useTestCache } from "../../../testing"; import { FilesPlugin } from "../plugin"; import { setupTestEnv, teardownTestEnv, VOLUMES_CONFIG } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("FilesPlugin shutdown and trackWrite", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); vi.useFakeTimers(); }); diff --git a/packages/appkit/src/plugins/files/tests/upload.test.ts b/packages/appkit/src/plugins/files/tests/upload.test.ts index 0e893cd59..9aa0c6357 100644 --- a/packages/appkit/src/plugins/files/tests/upload.test.ts +++ b/packages/appkit/src/plugins/files/tests/upload.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { createMockWorkspaceClient, useTestCache } from "../../../testing"; import { FilesPlugin } from "../plugin"; import { policy } from "../policy"; import { @@ -11,73 +12,27 @@ import { VOLUMES_CONFIG, } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Real in-memory cache; spy on `testCache.current` to assert invalidation. +const testCache = useTestCache(); describe("FilesPlugin upload", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -202,6 +157,9 @@ describe("FilesPlugin upload", () => { const handler = getRouteHandler(plugin, "post", "/upload"); const res = mockRes(); + const generateKey = vi.spyOn(testCache.current, "generateKey"); + const del = vi.spyOn(testCache.current, "delete"); + const req = mockUploadReq("uploads", [Buffer.from("file content")], { query: { path: "/Volumes/catalog/schema/uploads/dir/file.txt" }, }); @@ -223,8 +181,8 @@ describe("FilesPlugin upload", () => { expect(res.json).toHaveBeenCalledWith( expect.objectContaining({ success: true }), ); - expect(mockCacheInstance.generateKey).toHaveBeenCalled(); - expect(mockCacheInstance.delete).toHaveBeenCalled(); + expect(generateKey).toHaveBeenCalled(); + expect(del).toHaveBeenCalled(); }); }); }); diff --git a/packages/appkit/src/plugins/files/tests/volume-config.test.ts b/packages/appkit/src/plugins/files/tests/volume-config.test.ts index 121bdbe39..8ddba4b88 100644 --- a/packages/appkit/src/plugins/files/tests/volume-config.test.ts +++ b/packages/appkit/src/plugins/files/tests/volume-config.test.ts @@ -1,75 +1,31 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import { createMockWorkspaceClient, useTestCache } from "../../../testing"; +import { withEnv } from "../../../testing"; import { FilesPlugin } from "../plugin"; import { setupTestEnv, teardownTestEnv, VOLUMES_CONFIG } from "./_test-helpers"; -const { mockClient, MockApiError, mockCacheInstance } = vi.hoisted(() => { - const mockFilesApi = { - listDirectoryContents: vi.fn(), - download: vi.fn(), - getMetadata: vi.fn(), - upload: vi.fn(), - createDirectory: vi.fn(), - delete: vi.fn(), - }; - const mockClient = { - files: mockFilesApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; - class MockApiError extends Error { - statusCode: number; - constructor(message: string, statusCode: number) { - super(message); - this.name = "ApiError"; - this.statusCode = statusCode; - } - } - 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((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockClient, MockApiError, mockCacheInstance }; -}); - -vi.mock("../../../workspace-client", async (importOriginal) => { - const actual = - await importOriginal(); - return { - ...actual, - createWorkspaceClient: (..._args: unknown[]) => mockClient, - ApiError: MockApiError, - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("FilesPlugin volume config surface", () => { let serviceContextMock: Awaited>; + let client: ReturnType; + beforeEach(async () => { - serviceContextMock = await setupTestEnv(); + client = createMockWorkspaceClient({ + strict: true, + responses: { + "files.listDirectoryContents": undefined, + "files.download": undefined, + "files.getMetadata": undefined, + "files.upload": undefined, + "files.createDirectory": undefined, + "files.delete": undefined, + }, + }); + serviceContextMock = await setupTestEnv(client); }); afterEach(() => { @@ -94,14 +50,13 @@ describe("FilesPlugin volume config surface", () => { }); test("discovered volumes get empty config objects", () => { - process.env.DATABRICKS_VOLUME_DATA = "/Volumes/catalog/schema/data"; - - try { - const volumes = FilesPlugin.discoverVolumes({}); - expect(volumes.data).toEqual({}); - } finally { - delete process.env.DATABRICKS_VOLUME_DATA; - } + withEnv( + { DATABRICKS_VOLUME_DATA: "/Volumes/catalog/schema/data" }, + () => { + const volumes = FilesPlugin.discoverVolumes({}); + expect(volumes.data).toEqual({}); + }, + ); }); test("explicit volumes without env vars still appear", () => { @@ -119,20 +74,19 @@ describe("FilesPlugin volume config surface", () => { }); test("env var volume is not added when explicit config has the same key", () => { - process.env.DATABRICKS_VOLUME_SPECIAL = "/Volumes/catalog/schema/special"; - - try { - const volumes = FilesPlugin.discoverVolumes({ - volumes: { - special: { maxUploadSize: 500 }, - }, - }); - - // Explicit wins; should not be overwritten with {} - expect(volumes.special).toEqual({ maxUploadSize: 500 }); - } finally { - delete process.env.DATABRICKS_VOLUME_SPECIAL; - } + withEnv( + { DATABRICKS_VOLUME_SPECIAL: "/Volumes/catalog/schema/special" }, + () => { + const volumes = FilesPlugin.discoverVolumes({ + volumes: { + special: { maxUploadSize: 500 }, + }, + }); + + // Explicit wins; should not be overwritten with {} + expect(volumes.special).toEqual({ maxUploadSize: 500 }); + }, + ); }); }); diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 2d867d335..5cd03a348 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -11,36 +11,12 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { genieConnectorDefaults } from "../../../connectors/genie/defaults"; import { ServiceContext } from "../../../context/service-context"; import { Plugin } from "../../../plugin"; +import { useTestCache } from "../../../testing/test-cache"; import { GeniePlugin, genie } from "../genie"; import type { IGenieConfig } from "../types"; -// Mock CacheManager singleton -const { mockCacheInstance } = vi.hoisted(() => { - const instance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi - .fn() - .mockImplementation( - async ( - _key: unknown[], - fn: (signal?: AbortSignal) => Promise, - ) => { - return await fn(); - }, - ), - generateKey: vi.fn((...args: unknown[]) => JSON.stringify(args)), - }; - - return { mockCacheInstance: instance }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); function createMockGenieService() { const getMessageAttachmentQueryResult = vi.fn(); @@ -339,13 +315,14 @@ describe("Genie Plugin", () => { // 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( + const stream = expectStream(mockRes); + await stream.toEmit( "message_start", "status", "message_result", "query_result", ); - const events = await expectStream(mockRes).collect(); + const events = await stream.collect(); expect(events.find((e) => e.type === "message_start")).toMatchObject({ conversationId: "new-conv-id", }); diff --git a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts index 783debc8a..84b44ed0c 100644 --- a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts @@ -4,6 +4,7 @@ import { z } from "zod"; import { ServiceContext } from "../../../context/service-context"; import { ResourceType } from "../../../registry"; +import { createApiError, useTestCache, withEnv } from "../../../testing"; import { JOBS_READ_DEFAULTS, JOBS_STREAM_DEFAULTS, @@ -12,63 +13,42 @@ 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, jobsApi } = 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, getMock } = + await import("../../../testing/mock-workspace-client"); + + const mockClient = createMockWorkspaceClient(); + + // Facade accessors are typed against the legacy SDK, so `.mockResolvedValue` + // on them would not typecheck. `getMock` is the typed handle; it mints + // idempotently, so these are the very functions the plugin will call. + const jobsApi = { + runNow: getMock(mockClient, "jobs.runNow"), + submit: getMock(mockClient, "jobs.submit"), + getRun: getMock(mockClient, "jobs.getRun"), + getRunOutput: getMock(mockClient, "jobs.getRunOutput"), + cancelRun: getMock(mockClient, "jobs.cancelRun"), + listRuns: getMock(mockClient, "jobs.listRuns"), + get: getMock(mockClient, "jobs.get"), }; - const mockClient = { - jobs: mockJobsApi, - config: { - host: "https://test.databricks.com", - authenticate: 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 }; }); vi.mock("../../../workspace-client", async (importOriginal) => { const actual = await importOriginal(); - return { - ...actual, - createWorkspaceClient: () => mockClient, - Context: vi.fn(), - }; + // Only `Context` — the client itself is injected through ServiceContext. + return { ...actual, Context: vi.fn() }; }); -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); describe("JobsPlugin", () => { let serviceContextMock: Awaited>; @@ -77,7 +57,10 @@ describe("JobsPlugin", () => { vi.clearAllMocks(); setupDatabricksEnv(); ServiceContext.reset(); - serviceContextMock = await mockServiceContext(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: mockClient, + userDatabricksClient: mockClient, + }); }); afterEach(() => { @@ -154,33 +137,24 @@ describe("JobsPlugin", () => { }); test("skips bare DATABRICKS_JOB_ prefix (no suffix)", () => { - process.env.DATABRICKS_JOB_ = "999"; - try { + withEnv({ DATABRICKS_JOB_: "999" }, () => { const jobs = JobsPlugin.discoverJobs({}); expect(Object.keys(jobs)).not.toContain(""); - } finally { - delete process.env.DATABRICKS_JOB_; - } + }); }); test("skips empty env var values", () => { - process.env.DATABRICKS_JOB_EMPTY = ""; - try { + withEnv({ DATABRICKS_JOB_EMPTY: "" }, () => { const jobs = JobsPlugin.discoverJobs({}); expect(jobs).not.toHaveProperty("empty"); - } finally { - delete process.env.DATABRICKS_JOB_EMPTY; - } + }); }); test("lowercases env var suffix", () => { - process.env.DATABRICKS_JOB_MY_PIPELINE = "111"; - try { + withEnv({ DATABRICKS_JOB_MY_PIPELINE: "111" }, () => { const jobs = JobsPlugin.discoverJobs({}); expect(jobs).toHaveProperty("my_pipeline"); - } finally { - delete process.env.DATABRICKS_JOB_MY_PIPELINE; - } + }); }); test("returns only explicit jobs when no env vars match", () => { @@ -290,7 +264,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 +272,7 @@ describe("JobsPlugin", () => { await handle.runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123 }), expect.anything(), ); @@ -307,7 +281,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 +291,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 +323,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 +337,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 +349,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 +362,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 +389,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 +413,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 +421,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 +431,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 +452,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 +479,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 +495,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 +512,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 +528,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 +544,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"); }); @@ -592,9 +562,15 @@ describe("JobsPlugin", () => { test("error result preserves upstream HTTP status code", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - const error = new Error("Detailed internal failure: db connection reset"); - (error as any).statusCode = 403; - mockClient.jobs.getRun.mockRejectedValue(error); + // A genuine ApiError (as the SDK throws): the real cache preserves an + // ApiError's status but wraps a plain Error into a 500. + jobsApi.getRun.mockRejectedValue( + createApiError({ + statusCode: 403, + message: "Detailed internal failure: db connection reset", + errorCode: "PERMISSION_DENIED", + }), + ); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -611,7 +587,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 +604,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 +617,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 +627,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 +642,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 +672,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 +694,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 +805,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(), ); @@ -929,7 +905,10 @@ describe("injectRoutes", () => { vi.clearAllMocks(); setupDatabricksEnv(); ServiceContext.reset(); - serviceContextMock = await mockServiceContext(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: mockClient, + userDatabricksClient: mockClient, + }); }); afterEach(() => { @@ -1081,7 +1060,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 +1182,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 +1220,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 +1247,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 +1263,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 +1330,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 +1372,7 @@ describe("injectRoutes", () => { run_id: 42, state: { life_cycle_state: "TERMINATED" }, }; - mockClient.jobs.listRuns.mockReturnValue( + jobsApi.listRuns.mockReturnValue( (async function* () { yield mockRun; })(), @@ -1432,7 +1411,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 +1448,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 +1519,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 +1549,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 +1701,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 +1739,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 +1786,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"); @@ -1847,10 +1826,12 @@ describe("injectRoutes", () => { test("GET /:jobKey/runs returns upstream status on failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - const error = new Error("Unauthorized"); - (error as any).statusCode = 401; - mockClient.jobs.listRuns.mockImplementation(() => { - throw error; + jobsApi.listRuns.mockImplementation(() => { + throw createApiError({ + statusCode: 401, + message: "Unauthorized", + errorCode: "UNAUTHENTICATED", + }); }); const plugin = new JobsPlugin({}); @@ -1884,10 +1865,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/lakebase/tests/lakebase-agent-tool.test.ts b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts index 7e035bdea..89aaab9c7 100644 --- a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts +++ b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts @@ -1,5 +1,4 @@ import { beforeEach, describe, expect, test, vi } from "vitest"; - /** * Tests the agent-tool surface of the Lakebase plugin. * @@ -9,20 +8,11 @@ import { beforeEach, describe, expect, test, vi } from "vitest"; * (SP or per-user via RoutingPool). */ -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => ({ - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_k: unknown[], fn: (signal?: AbortSignal) => Promise) => - fn(), - ), - generateKey: vi.fn(() => "test-key"), - })), - }, -})); +import { useTestCache } from "../../../testing/test-cache"; + +// Boot AppKit's real in-memory cache so the base Plugin constructor's +// getInstanceSync() resolves instead of throwing. +useTestCache(); // Client calls recorded by the read-only-statement test. The `connect()` // mock returns a fresh client whose `query` pushes to this array so tests 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 bca2f091a..2906b13ca 100644 --- a/packages/appkit/src/plugins/serving/tests/serving.test.ts +++ b/packages/appkit/src/plugins/serving/tests/serving.test.ts @@ -10,35 +10,12 @@ import { import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { ServiceContext } from "../../../context/service-context"; +import { useTestCache } from "../../../testing/test-cache"; import { ServingPlugin, serving } from "../serving"; import type { IServingConfig } from "../types"; -// Mock CacheManager singleton -const { mockCacheInstance } = vi.hoisted(() => { - const instance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi - .fn() - .mockImplementation( - async ( - _key: unknown[], - fn: (signal?: AbortSignal) => Promise, - ) => { - return await fn(); - }, - ), - generateKey: vi.fn((...args: unknown[]) => JSON.stringify(args)), - }; - return { mockCacheInstance: instance }; -}); - -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: vi.fn(() => mockCacheInstance), - }, -})); +// Boots AppKit's real in-memory cache (no cache-module mock needed). +useTestCache(); // Mock the serving connector const mockInvoke = vi.fn(); diff --git a/packages/appkit/src/stream/arrow-stream-processor.ts b/packages/appkit/src/stream/arrow-stream-processor.ts index 62cab4df4..e063b6abb 100644 --- a/packages/appkit/src/stream/arrow-stream-processor.ts +++ b/packages/appkit/src/stream/arrow-stream-processor.ts @@ -1,11 +1,9 @@ import { ExecutionError, ValidationError } from "../errors"; import { createLogger } from "../logging/logger"; -import type { sql } from "../workspace-client"; +import type { ExternalLink } from "../workspace-client"; const logger = createLogger("stream:arrow"); -type ExternalLink = sql.ExternalLink; - /** * Re-mint a chunk's pre-signed URL. DBSQL external links expire in <= 15 min, * so a large result whose tail chunks are reached after the earlier chunks @@ -83,11 +81,11 @@ export class ArrowStreamProcessor { signal?: AbortSignal, refresh?: RefreshChunkLink, ): AsyncGenerator { - let externalLink = chunk.external_link; + let externalLink = chunk.externalLink; if (!externalLink) { // A missing link cannot be fixed by retrying — fail loudly. throw ExecutionError.statementFailed( - `External link missing for chunk ${chunk.chunk_index}`, + `External link missing for chunk ${chunk.chunkIndex}`, ); } @@ -114,7 +112,7 @@ export class ArrowStreamProcessor { clearTimeout(timer); if (!r.ok) { throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index}: ${r.status} ${r.statusText}`, + `Failed to download chunk ${chunk.chunkIndex}: ${r.status} ${r.statusText}`, ); } // Keep this attempt's controller alive to drive the body read + idle @@ -134,16 +132,16 @@ export class ArrowStreamProcessor { // chunk's link — a stale URL would just 403 again on the same address. // Only meaningful before any bytes are yielded (below), which is why // this lives in the establish-response loop. - if (refresh && chunk.chunk_index != null) { + if (refresh && chunk.chunkIndex != null) { try { - const fresh = await refresh(chunk.chunk_index, signal); - if (fresh?.external_link) externalLink = fresh.external_link; + const fresh = await refresh(chunk.chunkIndex, signal); + if (fresh?.externalLink) externalLink = fresh.externalLink; } catch (refreshError) { // Keep retrying the current URL; surface the original error if // all attempts fail. logger.warn( "Failed to re-resolve link for chunk %s: %O", - chunk.chunk_index, + chunk.chunkIndex, refreshError, ); } @@ -154,7 +152,7 @@ export class ArrowStreamProcessor { if (!response || !controller) { throw ExecutionError.statementFailed( - `Failed to download chunk ${chunk.chunk_index} after ${this.options.retries} attempts: ${ + `Failed to download chunk ${chunk.chunkIndex} after ${this.options.retries} attempts: ${ lastError instanceof Error ? lastError.message : String(lastError) }`, ); @@ -194,13 +192,13 @@ export class ArrowStreamProcessor { if (signal?.aborted) throw ExecutionError.canceled(); logger.error( "Failed streaming chunk %s body: %O", - chunk.chunk_index, + chunk.chunkIndex, error, ); throw error instanceof ExecutionError ? error : ExecutionError.statementFailed( - `Failed streaming chunk ${chunk.chunk_index}: ${ + `Failed streaming chunk ${chunk.chunkIndex}: ${ error instanceof Error ? error.message : String(error) }`, ); diff --git a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts index 555f87339..84d1ee031 100644 --- a/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts +++ b/packages/appkit/src/stream/tests/arrow-stream-processor.test.ts @@ -1,6 +1,5 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; -import type { sql } from "../../workspace-client"; import { ArrowStreamProcessor } from "../arrow-stream-processor"; /** A ReadableStream that emits the given pieces in order, then closes. */ @@ -15,10 +14,10 @@ function streamOf(...pieces: Uint8Array[]): ReadableStream { function mockChunks(count: number) { return Array.from({ length: count }, (_, i) => ({ - chunk_index: i, - external_link: `https://example.com/chunk-${i}`, - row_offset: i * 100, - row_count: 100, + chunkIndex: i, + externalLink: `https://example.com/chunk-${i}`, + rowOffset: BigInt(i * 100), + rowCount: 100n, })); } @@ -166,7 +165,7 @@ describe("ArrowStreamProcessor.streamChunks", () => { }); test("throws immediately when a chunk has no external_link", async () => { - const chunks = [{ chunk_index: 0 }] as any; + const chunks = [{ chunkIndex: 0 }] as any; await expect(drain(processor.streamChunks(chunks))).rejects.toThrow( /External link missing/, ); @@ -218,8 +217,8 @@ describe("ArrowStreamProcessor.streamChunks", () => { globalThis.fetch = fetchMock; const refresh = vi.fn(async (chunkIndex: number) => ({ - chunk_index: chunkIndex, - external_link: "https://example.com/fresh-link", + chunkIndex, + externalLink: "https://example.com/fresh-link", })); const p = new ArrowStreamProcessor({ timeout: 5000, retries: 3 }); diff --git a/packages/appkit/src/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index 3d3815610..a2b3c9d5f 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -281,6 +281,12 @@ export class TelemetryManager { * 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 reassigned + * for whatever providers are currently live, so a stale resolved promise can + * only be returned when there is nothing to flush. The covering test asserts + * every provider set across repeated initialize/shutdown cycles is flushed. */ async shutdown(): Promise { const providers = [ @@ -308,4 +314,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.shutdown()` 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..97b10e33f --- /dev/null +++ b/packages/appkit/src/telemetry/tests/telemetry-manager-reset.test.ts @@ -0,0 +1,184 @@ +import { context, metrics, trace } from "@opentelemetry/api"; +import { logs } from "@opentelemetry/api-logs"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +/** + * Telemetry now builds three providers — the meter and logger in `initialize()`, + * the tracer in `start()` — instead of a single `NodeSDK`. These mock the three + * provider constructors so the shutdown/flush path is observable, and set an OTLP + * endpoint so `initialize()` actually builds the meter/logger providers. + * + * The never-cleared `shutdownPromise` was suspected of skipping a re-booted + * provider set's flush. It does not — the memo is reassigned whenever providers + * are live. Re-boot goes through `reset()` (what `dropCoreSingletons()` does + * between harness boots), because `initialize()`/`start()` are idempotent within + * one manager: they guard on `resource`/`started`, which `shutdown()` deliberately + * does not clear. A bare re-`initialize()` after `shutdown()` is therefore a no-op. + */ + +const { + meterProviderShutdown, + loggerProviderShutdown, + tracerProviderShutdown, + MeterProviderMock, + LoggerProviderMock, + NodeTracerProviderMock, +} = vi.hoisted(() => { + const meterProviderShutdown = vi.fn().mockResolvedValue(undefined); + const loggerProviderShutdown = vi.fn().mockResolvedValue(undefined); + const tracerProviderShutdown = vi.fn().mockResolvedValue(undefined); + return { + meterProviderShutdown, + loggerProviderShutdown, + tracerProviderShutdown, + MeterProviderMock: vi.fn(() => ({ shutdown: meterProviderShutdown })), + LoggerProviderMock: vi.fn(() => ({ shutdown: loggerProviderShutdown })), + NodeTracerProviderMock: vi.fn(() => ({ + register: vi.fn(), + shutdown: tracerProviderShutdown, + })), + }; +}); + +vi.mock("@opentelemetry/sdk-metrics", () => ({ + MeterProvider: MeterProviderMock, + PeriodicExportingMetricReader: vi.fn(() => ({})), +})); +vi.mock("@opentelemetry/sdk-logs", () => ({ + LoggerProvider: LoggerProviderMock, + BatchLogRecordProcessor: vi.fn(() => ({})), +})); +vi.mock("@opentelemetry/sdk-trace-node", () => ({ + NodeTracerProvider: NodeTracerProviderMock, +})); +// Keep the real module (AppKitSampler needs SamplingDecision); only stub the +// span processor so no real exporter/timer is wired up. +vi.mock("@opentelemetry/sdk-trace-base", async () => { + const actual = await vi.importActual< + typeof import("@opentelemetry/sdk-trace-base") + >("@opentelemetry/sdk-trace-base"); + return { ...actual, BatchSpanProcessor: vi.fn(() => ({})) }; +}); +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.resourceFromAttributes({})), + }; +}); + +import { TelemetryManager } from "../telemetry-manager"; + +/** Reset the singleton and clear any globals a prior boot registered. */ +function resetTelemetry(): void { + TelemetryManager.reset(); + metrics.disable(); + logs.disable(); + trace.disable(); + context.disable(); +} + +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(); + resetTelemetry(); + }); + + afterEach(() => { + if (originalEndpoint === undefined) { + delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + } else { + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = originalEndpoint; + } + resetTelemetry(); + }); + + test("a reset() between boots rebuilds and re-flushes every provider", async () => { + // Boot 1: initialize() builds the meter + logger providers, start() the tracer. + TelemetryManager.initialize({}); + TelemetryManager.start(); + expect(MeterProviderMock).toHaveBeenCalledTimes(1); + expect(LoggerProviderMock).toHaveBeenCalledTimes(1); + expect(NodeTracerProviderMock).toHaveBeenCalledTimes(1); + + await TelemetryManager.getInstance().shutdown(); + expect(meterProviderShutdown).toHaveBeenCalledTimes(1); + expect(loggerProviderShutdown).toHaveBeenCalledTimes(1); + expect(tracerProviderShutdown).toHaveBeenCalledTimes(1); + + // Re-boot the way the harness does — reset() (dropCoreSingletons) then boot. + // A bare re-initialize() would be a no-op here (see the file header). + resetTelemetry(); + + TelemetryManager.initialize({}); + TelemetryManager.start(); + expect(MeterProviderMock).toHaveBeenCalledTimes(2); + expect(NodeTracerProviderMock).toHaveBeenCalledTimes(2); + + await TelemetryManager.getInstance().shutdown(); + expect(meterProviderShutdown).toHaveBeenCalledTimes(2); + expect(loggerProviderShutdown).toHaveBeenCalledTimes(2); + expect(tracerProviderShutdown).toHaveBeenCalledTimes(2); + + // A third cycle, to pin the general property rather than one transition. + resetTelemetry(); + TelemetryManager.initialize({}); + TelemetryManager.start(); + await TelemetryManager.getInstance().shutdown(); + expect(MeterProviderMock).toHaveBeenCalledTimes(3); + expect(meterProviderShutdown).toHaveBeenCalledTimes(3); + expect(tracerProviderShutdown).toHaveBeenCalledTimes(3); + }); + + test("concurrent shutdown() calls share one flush", async () => { + TelemetryManager.initialize({}); + TelemetryManager.start(); + const manager = TelemetryManager.getInstance(); + + await Promise.all([manager.shutdown(), manager.shutdown()]); + + // Clearing the provider refs synchronously is what makes this safe: the + // second caller finds no providers and awaits the first caller's memo. + expect(meterProviderShutdown).toHaveBeenCalledTimes(1); + expect(loggerProviderShutdown).toHaveBeenCalledTimes(1); + expect(tracerProviderShutdown).toHaveBeenCalledTimes(1); + }); + + test("shutdown() with no providers built resolves without flushing", async () => { + delete process.env.OTEL_EXPORTER_OTLP_ENDPOINT; + TelemetryManager.initialize({}); + TelemetryManager.start(); + const manager = TelemetryManager.getInstance(); + + await expect(manager.shutdown()).resolves.toBeUndefined(); + expect(meterProviderShutdown).not.toHaveBeenCalled(); + expect(loggerProviderShutdown).not.toHaveBeenCalled(); + expect(tracerProviderShutdown).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..159646980 --- /dev/null +++ b/packages/appkit/src/testing/create-test-app.ts @@ -0,0 +1,436 @@ +/** + * 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 { vi } from "vitest"; + +import { InMemoryStorage } from "../cache/storage/memory"; +import { ServiceContext } from "../context/service-context"; +import { AppKit, disposeApp } from "../core/appkit"; +import type { WorkspaceClient } from "../workspace-client"; +import type { OboOption } from "./fixtures"; +import { fakeUserContext, oboHeaders, setupDatabricksEnv } from "./fixtures"; +import type { CreateMockWorkspaceClientOptions } from "./mock-workspace-client"; +import { createMockWorkspaceClient } from "./mock-workspace-client"; +import { dropCoreSingletons } from "./reset-singletons"; + +// Loose shapes are intentional here; `noExplicitAny` is off repo-wide (see +// .oxlintrc.json), so a local alias keeps the intent readable. +type Any = any; + +/** + * The env snapshot to restore on close, and whether an app currently holds it. + * + * One live app at a time (see the guard in `createTestApp`), so a flag suffices. + */ +let envBaseline: NodeJS.ProcessEnv | undefined; +let harnessAppLive = false; + +/** Take the baseline on boot. */ +function acquireEnvBaseline(): void { + envBaseline = { ...process.env }; + harnessAppLive = true; +} + +/** Restore the baseline on close. */ +function releaseEnvBaseline(): void { + harnessAppLive = false; + if (!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. Refused when `client` is set. */ + responses?: CreateMockWorkspaceClientOptions["responses"]; + + /** + * Make the built-in mock throw when a path with no declared response is + * called, rather than resolving `undefined`. Refused when `client` is set — + * configure it on your own client instead. + */ + strict?: CreateMockWorkspaceClientOptions["strict"]; + + /** + * 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. + * + * Beyond refusing `development`, this decides error-response redaction: + * `errorHandlerMiddleware` returns the real message unless `NODE_ENV` is + * `production`, where a 5xx becomes `"Server error"`. Pass `"production"` to + * assert what a deployed app actually returns to a client. + */ + nodeEnv?: string; + + /** Defaults to in-memory, which is what keeps boot offline. */ + cache?: CacheConfig; +} + +/** 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; +} + +/** + * Point `ServiceContext.createUserContext` at the harness's mock so an `obo` + * request does not construct a real SDK client from `DATABRICKS_HOST`. + * + * Mirrors the `createUserContextSpy` in `fixtures.ts`; returns its restore. + */ +function stubUserContext(client: WorkspaceClient): () => void { + const spy = vi + .spyOn(ServiceContext, "createUserContext") + .mockImplementation((token, userId, userName, userEmail) => + fakeUserContext(client, ServiceContext.get())( + token, + userId, + userName, + userEmail, + ), + ); + return () => spy.mockRestore(); +} + +/** + * Wait for a server to finish binding and return the port it landed on. + * + * Needed with `port: 0`: `start()` returns once `listen()` is invoked, before + * the bind completes, so `address()` is null until the `listening` event fires. + * `createTestApp` does this for you — reach for it when hand-rolling a server. + */ +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, + strict, + client: suppliedClient, + env = {}, + server: serverOption, + nodeEnv = "test", + cache, + } = 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.", + ); + } + + // Refused rather than half-supported: AppKit's workspace client, cache and + // on-behalf-of fake are process-wide, so a second live app cannot own its + // own. Checked before any mutation, so a refused boot leaves the live app + // untouched. + if (harnessAppLive) { + throw new Error( + "createTestApp: a harness app is already open. AppKit's workspace " + + "client, cache, and on-behalf-of fake are process-wide, so a second " + + "app would not receive its own `client`/`responses`, and closing " + + "either would un-fake the other's on-behalf-of path. Close the first " + + "app before booting another — `await using`, or try/finally.", + ); + } + + // Wholesale rather than a whitelist: plugins read vars we cannot enumerate. + acquireEnvBaseline(); + + let app: Awaited> | undefined; + let restoreUserContext: (() => void) | undefined; + + // Runs the booted plugins' shutdown() hooks and closes the server it started; + // the harness owns dropping the singletons and restoring env. The `as Any` is + // the one escape hatch the symbol-keyed teardown forces on the harness. + const disposeBooted = (a: unknown): Promise => (a as Any)[disposeApp](); + + 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, + }); + + dropCoreSingletons(); + + // `responses` only seeds the built-in mock, so alongside a caller-supplied + // client it would silently do nothing. Refuse instead, matching the + // `server: false` conflict below. + if (suppliedClient && (responses !== undefined || strict !== undefined)) { + throw new Error( + "createTestApp: `responses` and `strict` configure the built-in mock " + + "client, so they do nothing when you also pass `client`. Drop them " + + "and configure your own client instead.", + ); + } + + // 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, strict }); + + // createApp({ client }) installs only the service-principal client. An `obo` + // request reaches ServiceContext.createUserContext, which builds a *real* + // client from process.env.DATABRICKS_HOST — so the user-scoped path is faked + // here too, or "no network" is false the moment a handler calls asUser. + restoreUserContext = stubUserContext(client); + + // 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 required to stay offline: without explicit storage the + // cache builds its own client and probes Lakebase, and without the opt-out + // TelemetryReporter fires an apiClient.request on boot. + app = await AppKit._createApp({ + plugins: bootPlugins as Any, + client, + cache: cache ?? { + storage: new InMemoryStorage({ enabled: true } as Any), + }, + disableInternalTelemetry: true, + // The harness boots repeatedly in one process and runs its own teardown, + // so it must not accumulate SIGTERM/SIGINT handlers across boots. + installSignalHandlers: false, + }); + + 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 disposeBooted(bootedApp); + } finally { + dropCoreSingletons(); + restoreUserContext?.(); + releaseEnvBaseline(); + } + })(); + 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. + // Lowercased first: `Headers` comma-joins case variants instead of + // replacing, so a mixed-case override would corrupt the value into + // "alice, bob" rather than win. + for (const [name, value] of Object.entries(reqOptions.headers ?? {})) { + headers[name.toLowerCase()] = value; + } + + 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. + if (app) { + try { + await disposeBooted(app); + } catch { + // The boot error is the interesting one; don't let teardown mask it. + } + } + // dispose() no longer drops the singletons, and a failure before boot may + // have left a partial init — drop unconditionally either way. + dropCoreSingletons(); + restoreUserContext?.(); + releaseEnvBaseline(); + 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/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index f3663b430..eaa8d8eed 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -1,3 +1,5 @@ +import { createHash } from "node:crypto"; + import type { Span, SpanOptions } from "@opentelemetry/api"; import type { IAppRouter } from "shared"; import { afterEach, beforeEach, vi } from "vitest"; @@ -5,10 +7,13 @@ import { afterEach, beforeEach, vi } from "vitest"; import { CacheManager } from "../cache"; import type { ServiceContextState } from "../context/service-context"; import { ServiceContext } from "../context/service-context"; +import { AuthenticationError } from "../errors"; import type { InstrumentConfig, ITelemetry } from "../telemetry/types"; +import { ApiError } from "../workspace-client"; +import { createMockWorkspaceClient } from "./mock-workspace-client"; // Test fixtures intentionally use loose shapes; `noExplicitAny` is disabled -// repo-wide (see biome.json), so a local alias keeps the intent readable. +// repo-wide (see .oxlintrc.json), so a local alias keeps the intent readable. type Any = any; /** @@ -121,8 +126,60 @@ export type OboOption = email?: string; }; -/** Build the forwarded identity headers an `obo` option implies. */ -function oboHeaders(obo: Exclude): Record { +/** + * The one fake of `ServiceContext.createUserContext` this kit uses, shared by + * `mockServiceContext` and `createTestApp`. + * + * Shared rather than duplicated because the two used to disagree, and neither + * matched production: a missing token went unrejected and `tokenFingerprint` + * was absent, which silently disables Lakebase pool rotation — `pool-manager` + * treats a missing fingerprint as "not stale", so the drain-and-recreate branch + * could never run under a fake. + * + * @internal + */ +export function fakeUserContext( + client: Any, + ids: { warehouseId?: Any; workspaceId: Any }, +) { + return ( + token: string, + userId: string, + userName?: string, + userEmail?: string, + ): Any => { + // Same rejection as production, so a path that forgets to forward the token + // fails here instead of only in a deployed app. + if (!token) throw AuthenticationError.missingToken("user token"); + return { + client, + userId, + userName, + userEmail, + // Derived from the token exactly as production does. Keyed on the user it + // would be constant across tokens, and rotation compares this value. + tokenFingerprint: createHash("sha256") + .update(token) + .digest("hex") + .slice(0, 16), + warehouseId: ids.warehouseId, + workspaceId: ids.workspaceId, + isUserContext: true, + }; + }; +} + +/** + * 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", @@ -284,6 +341,102 @@ export function setupDatabricksEnv(overrides: Record = {}) { Object.assign(process.env, overrides); } +/** + * Set environment variables and return a function that restores each key to its + * prior state — the prior value, or a delete when the key was previously unset. + * Shared capture/restore behind {@link withEnv} and the + * {@link createTestPluginContext} options path; not part of the public surface. + */ +export function applyEnv(vars: Record): () => void { + const prior = new Map(); + for (const key of Object.keys(vars)) { + prior.set(key, process.env[key]); + } + Object.assign(process.env, vars); + return () => { + for (const [key, value] of prior) { + if (value === undefined) { + delete process.env[key]; + } else { + process.env[key] = value; + } + } + }; +} + +/** + * Run a restore on an error path, suppressing any failure it throws so it + * cannot replace the caller's original error. + */ +function restoreQuietly(restore: () => void): void { + try { + restore(); + } catch (restoreError) { + // A failed env restore must not mask the caller's original error. + void restoreError; + } +} + +/** + * Sets environment variables for the duration of `fn`, then restores them to + * their prior state. Each key's prior value (or "was absent") is captured on + * entry; on exit, the prior value is restored, or the key is deleted only if + * it was previously unset. + * + * Supports both sync and async `fn`. If `fn` returns a thenable, `withEnv` + * returns that promise and restores in `.finally()`. Otherwise, it restores + * in a synchronous `finally` and returns the callback's return value. + * Restoration runs even if `fn` throws. Nested calls restore in LIFO order. + * + * @example + * ```ts + * // Sync: restores synchronously after fn + * withEnv({ MY_VAR: "test" }, () => { + * console.log(process.env.MY_VAR); // "test" + * }); + * console.log(process.env.MY_VAR); // prior value (or undefined) + * + * // Async: restores after promise settles + * await withEnv({ MY_VAR: "test" }, async () => { + * await fetch(...); + * }); + * ``` + */ +export function withEnv( + vars: Record, + fn: () => T | Promise, +): T | Promise { + const restore = applyEnv(vars); + + let result: T | Promise; + try { + result = fn(); + } catch (err) { + // Sync throw: restore, but never let a restore failure mask `err`. + restoreQuietly(restore); + throw err; + } + + // Async: restore after the promise settles. On rejection, guard the restore + // so it cannot replace the caller's error; on success, let a genuine restore + // failure surface. + if (result != null && typeof (result as Any).then === "function") { + return (result as Promise).then( + (value) => { + restore(); + return value; + }, + (err: unknown) => { + restoreQuietly(restore); + throw err; + }, + ); + } + + restore(); + return result; +} + /** * Clears AppKit's process-wide cache singleton so cached values don't leak * between tests in the same file. @@ -332,28 +485,6 @@ export interface TestContextOptions { workspaceId?: string; } -/** - * Creates a default mock WorkspaceClient for testing (SQL succeeds, warehouse - * RUNNING). - */ -export function createMockWorkspaceClient() { - return { - statementExecution: { - executeStatement: vi.fn().mockResolvedValue({ - status: { state: "SUCCEEDED" }, - result: { data: [] }, - }), - }, - // Analytics route now calls `warehouses.get` before issuing SQL to - // ensure the warehouse is RUNNING. Default to RUNNING so existing - // tests that only care about SQL behaviour aren't affected. - warehouses: { - get: vi.fn().mockResolvedValue({ state: "RUNNING" }), - start: vi.fn().mockResolvedValue(undefined), - }, - }; -} - /** * Builds a {@link ServiceContextState} value for testing without touching the * singleton. Internal building block for {@link mockServiceContext}, which @@ -396,17 +527,12 @@ export function mockServiceContext(options: TestContextOptions = {}) { 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, - }; - }); + .mockImplementation( + fakeUserContext( + options.userDatabricksClient || createMockWorkspaceClient(), + serviceContext, + ), + ); return { serviceContext, @@ -503,19 +629,19 @@ export async function runWithRequestContext( */ export function createSuccessfulSQLResponse( data: Any[][], - columns: Array<{ name: string; type_name?: string }>, + columns: Array<{ name: string; typeName?: string }>, ) { return { status: { state: "SUCCEEDED" }, - statement_id: `stmt-${Date.now()}`, + statementId: `stmt-${Date.now()}`, result: { - data_array: data, + dataArray: data, }, manifest: { schema: { columns: columns.map((col) => ({ name: col.name, - type_name: col.type_name ?? "STRING", + typeName: col.typeName ?? "STRING", })), }, }, @@ -531,42 +657,39 @@ export function createFailedSQLResponse(errorMessage: string) { message: errorMessage, }, }, - statement_id: `stmt-${Date.now()}`, + statementId: `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. + * Creates a genuine `ApiError` instance for testing error paths. Returns a real + * instance (where `error instanceof ApiError` holds), suitable for testing + * `instanceof` checks and `.statusCode` / `.errorCode` / `.message` accessors. + * + * @param options Error details: `statusCode`, `message`, and `errorCode`. All required. + * @returns A genuine `ApiError` instance. + * + * @example + * ```ts + * const error = createApiError({ + * statusCode: 404, + * message: "File not found", + * errorCode: "NOT_FOUND", + * }); + * expect(error).toBeInstanceOf(ApiError); + * expect(error.statusCode).toBe(404); + * ``` */ -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 createApiError(options: { + statusCode: number; + message: string; + errorCode: string; +}): ApiError { + return new ApiError( + options.message, + options.errorCode, + options.statusCode, + undefined, // response: sensible default for testing + [], // details: empty array + ); } diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts index 8565e4d5e..165f81ae7 100644 --- a/packages/appkit/src/testing/index.ts +++ b/packages/appkit/src/testing/index.ts @@ -9,9 +9,11 @@ * buffering, tool dispatch, timeout composition, user scoping — run under test * with no credentials. * - * Two entry points: + * 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. + * 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 @@ -41,6 +43,13 @@ // 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, @@ -51,13 +60,12 @@ export { type StreamSource, } from "./expect-stream"; export { - createConfigurableMockWorkspaceClient, + createApiError, createFailedSQLResponse, createMockRequest, createMockResponse, createMockRouter, createMockTelemetry, - createMockWorkspaceClient, createSuccessfulSQLResponse, mockServiceContext, type OboOption, @@ -67,7 +75,17 @@ export { setupDatabricksEnv, type TestContextOptions, useServiceContextMock, + withEnv, } from "./fixtures"; +export { + createMockWorkspaceClient, + type CreateMockWorkspaceClientOptions, + getMock, + type MockWorkspaceClient, +} from "./mock-workspace-client"; +export { createTestPlugin } from "./create-test-plugin"; +export { type TestAppHandle, useTestApp } from "./test-app"; +export { type TestCacheHandle, useTestCache } from "./test-cache"; export { createTestPluginContext, type FakeProvider, @@ -76,4 +94,5 @@ export { type RecordedRoute, type RecordedToolCall, type TestPluginContext, + type TestPluginContextOptions, } 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..59ec97635 --- /dev/null +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -0,0 +1,320 @@ +/** + * 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"; + +// Loose shapes are intentional here; `noExplicitAny` is off repo-wide (see +// .oxlintrc.json), so a local alias keeps the intent readable. +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; + /** + * Throw when a path with no declared response is *called*, instead of + * resolving `undefined`. + * + * Off by default: the never-crash floor is what lets a plugin touch services + * a test does not care about. Turn it on when a silently `undefined` return + * would let the test pass for the wrong reason. + */ + strict?: boolean; + + /** 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`. + * + * `statementExecution.executeStatement`, `warehouses.getWarehouse` and + * `warehouses.startWarehouse` must stay byte-identical to the old `fixtures.ts` + * values — suites reach them + * implicitly through `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.getWarehouse": { state: "RUNNING" }, + "warehouses.startWarehouse": undefined, + "currentUser.me": { + id: "test-service-user", + userName: "test-service-user", + }, +}; + +/** Generically proxied; `config`/`apiClient` are seeded below instead. */ +const FACADE_SERVICES = [ + "files", + "warehouses", + "genie", + "jobs", + "statementExecution", + "servingEndpoints", + "currentUser", +] as const; + +/** + * Answered with `undefined` rather than a minted mock. `then` must stay listed: + * 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", +]); + +/** `apiClient` members read synchronously; a seeded value must not be resolved. */ +const SYNC_API_CLIENT_MEMBERS: ReadonlySet = new Set(["userAgent"]); + +/** Distinguishes "no passthrough rule applied" from a rule answering `undefined`. */ +const NOT_PASSTHROUGH = Symbol("not-passthrough"); + +/** + * The passthrough rules every trap in this file shares: symbols delegate to the + * target, denied names answer `undefined`, and anything already on the target + * (seeded members, `Object.prototype`) wins over minting. + * + * Shared rather than copied so a key added to {@link PASSTHROUGH_DENY} cannot + * cover one trap and miss another. + */ +function passthroughFor(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 NOT_PASSTHROUGH; +} + +/** + * `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 => { + const passthrough = passthroughFor(target, prop); + if (passthrough !== NOT_PASSTHROUGH) return passthrough; + 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, + strict = false, + } = 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 getMock, 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 if (strict && !(path in merged)) { + // Thrown on call, never on mint: `getMock` mints to hand back a handle + // before the code under test runs, and that must not blow up. + fn.mockImplementation(() => { + throw new Error( + `createMockWorkspaceClient: "${path}" was called with no declared ` + + "response and `strict: true` is set. Add it to `responses`, or drop " + + "`strict` to have undeclared paths resolve undefined.", + ); + }); + } 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") { + // Match the member's own shape. `userAgent()` is read straight into a + // Headers value, so resolving a seed would put "[object Promise]" there — + // the failure the synchronous default exists to avoid. Seed a function to + // decide for yourself. + if (SYNC_API_CLIENT_MEMBERS.has(key)) fn.mockReturnValue(value); + else 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 => { + // Same passthrough rules as `neverCrashGet` — shared, not copied, so a + // key added to PASSTHROUGH_DENY cannot cover one trap and miss this + // one (miss `then` here and `await client` hangs again). + const passthrough = passthroughFor(target, prop); + if (passthrough !== NOT_PASSTHROUGH) return passthrough; + 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 getMock(client: MockWorkspaceClient, path: string): Mock { + const fns = clientFns.get(client); + if (!fns) { + throw new Error( + "getMock: 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( + `getMock: "${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-singletons.ts b/packages/appkit/src/testing/reset-singletons.ts new file mode 100644 index 000000000..834e712d4 --- /dev/null +++ b/packages/appkit/src/testing/reset-singletons.ts @@ -0,0 +1,35 @@ +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("testing"); + +/** + * Drop the process-wide singletons `AppKit._createApp` initializes — called by + * the harness on boot to clear a previous test's leakage, and on teardown. + * + * Kit-owned: the only caller is the test harness, so it lives here rather than + * in core. A pointer drop, not teardown — close the app first (the harness runs + * the shutdown phases before this) or the old app's storage and exporters leak. + * A caller that drops then reads `ServiceContext.get()` gets an + * `InitializationError`. + * @internal + */ +export function dropCoreSingletons(): 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/testing/test-app.ts b/packages/appkit/src/testing/test-app.ts new file mode 100644 index 000000000..d726cd30d --- /dev/null +++ b/packages/appkit/src/testing/test-app.ts @@ -0,0 +1,83 @@ +import type { PluginConstructor, PluginData } from "shared"; +import { afterEach, beforeEach } from "vitest"; + +import type { CreateTestAppOptions, TestApp } from "./create-test-app"; +import { createTestApp } from "./create-test-app"; + +/** Mirrors `create-test-app.ts`'s own constraint; not part of the public surface. */ +type Plugins = PluginData[]; + +/** + * The handle {@link useTestApp} returns: a live accessor for the harness app + * booted for the current test. + */ +export interface TestAppHandle { + /** + * The app booted for the current test. Read it inside a test body — each + * `beforeEach` boots a fresh app and each `afterEach` closes it. + */ + readonly current: TestApp; +} + +/** + * Boot a harness app before each test and close it after, so a suite that needs + * an app per test never hand-wires the hooks or risks a forgotten `close()`. + * + * Mirrors {@link useServiceContextMock} and {@link useTestCache}: call it at the + * top of a `describe` block (or module top-level), NOT inside a test — Vitest + * registers `beforeEach`/`afterEach` during collection. + * + * Reach for this when the app must outlive a single expression. `await using` + * covers one test more concisely, but it cannot carry an app from a `beforeEach` + * into the test body, and the harness allows only one open app at a time — so a + * `describe` that holds one in `beforeAll` cannot contain a test that boots its + * own. + * + * @example + * ```ts + * describe("my plugin over HTTP", () => { + * const app = useTestApp({ + * plugins: [myPlugin()], + * responses: { "jobs.getRun": { state: "TERMINATED" } }, + * }); + * + * test("answers a request", async () => { + * const res = await app.current.post("/api/my-plugin/run", { body: { id: 1 } }); + * expect(res.status).toBe(200); + * }); + * }); + * ``` + * + * @param options - Passed to {@link createTestApp} unchanged, for every boot. + * @returns `{ current }` — the app booted for the current test. + */ +export function useTestApp( + options: CreateTestAppOptions = {}, +): TestAppHandle { + let app: TestApp | undefined; + + beforeEach(async () => { + app = await createTestApp(options); + }); + + afterEach(async () => { + const booted = app; + // Cleared before the await so a close that throws cannot leave a stale + // handle readable by the next test. + app = undefined; + await booted?.close(); + }); + + return { + get current(): TestApp { + if (!app) { + throw new Error( + "useTestApp: no active app. Call useTestApp() at the top of a " + + "describe block (not inside a test), and read `.current` from " + + "within a test.", + ); + } + return app; + }, + }; +} diff --git a/packages/appkit/src/testing/test-cache.ts b/packages/appkit/src/testing/test-cache.ts new file mode 100644 index 000000000..6a1559010 --- /dev/null +++ b/packages/appkit/src/testing/test-cache.ts @@ -0,0 +1,90 @@ +import { afterEach, beforeEach } from "vitest"; + +import { CacheManager } from "../cache"; +import { InMemoryStorage } from "../cache/storage"; +import { resetTestCache } from "./fixtures"; + +/** + * The handle {@link useTestCache} returns: a live accessor for the real, + * in-memory {@link CacheManager} active in the current test. + */ +export interface TestCacheHandle { + /** + * The real (in-memory) cache for the current test. Each test's `beforeEach` + * seeds and clears the singleton, so reading this always sees a fresh cache — + * call `generateKey`, `get`, `has`, or `vi.spyOn(handle.current, "getOrExecute")` + * to assert real caching behaviour. + */ + readonly current: CacheManager; +} + +/** + * Stand up AppKit's real in-memory cache for a test file and clear it before + * each test, so a plugin's real caching path runs under test with no mock of + * the internal `cache` module. + * + * Boots the process-wide {@link CacheManager} singleton backed by + * {@link InMemoryStorage} (idempotent — an already-initialized singleton is + * reused and the storage argument ignored), then clears it via + * {@link resetTestCache} in `beforeEach` so each test starts empty. The + * singleton is left in place: this clears the cache's contents, never the + * pointer. + * + * Because the cache is booted before the test body runs, a plugin constructed + * in the test — whose constructor reads `CacheManager.getInstanceSync()` — + * binds `this.cache` to this real cache. So `getOrExecute` genuinely caches and + * the key is production's real `generateKey`, not a re-implemented fake. + * + * 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. + * + * @example + * ```ts + * describe("my plugin caches", () => { + * const testCache = useTestCache(); + * + * test("second identical request is a cache hit", async () => { + * const plugin = new MyPlugin(config); + * // ...drive the same request twice... + * expect(downstreamMock).toHaveBeenCalledTimes(1); + * }); + * + * test("metadata does not change the cache key", () => { + * const key = testCache.current.generateKey(["query", "SELECT 1"], "svc"); + * expect(key).toBe(testCache.current.generateKey(["query", "SELECT 1"], "svc")); + * }); + * }); + * ``` + * + * @returns `{ current }` — the active real {@link CacheManager} for the test. + */ +export function useTestCache(): TestCacheHandle { + let cache: CacheManager | undefined; + + beforeEach(async () => { + // Idempotent: reuses an existing singleton (ignoring the storage arg) or + // stands up a fresh in-memory one. Keeps the singleton either way. + cache = await CacheManager.getInstance({ + storage: new InMemoryStorage({}), + }); + // Fresh contents per test — clears storage without dropping the singleton. + await resetTestCache(); + }); + + afterEach(() => { + cache = undefined; + }); + + return { + get current(): CacheManager { + if (!cache) { + throw new Error( + "useTestCache: no active cache. Call useTestCache() at the top of a " + + "describe block (not inside a test), and read `.current` from " + + "within a test.", + ); + } + return cache; + }, + }; +} diff --git a/packages/appkit/src/testing/test-plugin-context.ts b/packages/appkit/src/testing/test-plugin-context.ts index 721e3cdbe..037235c2b 100644 --- a/packages/appkit/src/testing/test-plugin-context.ts +++ b/packages/appkit/src/testing/test-plugin-context.ts @@ -5,6 +5,7 @@ import type { IAppRequest, ToolProvider, } from "shared"; +import { afterEach, onTestFinished } from "vitest"; import { CacheManager } from "../cache"; import { InMemoryStorage } from "../cache/storage"; @@ -12,7 +13,8 @@ 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"; +import { applyEnv, createMockTelemetry, mockServiceContext } from "./fixtures"; +import { createMockWorkspaceClient } from "./mock-workspace-client"; /** * A concrete (non-function) fake tool response — returned as-is. Covers the @@ -52,6 +54,30 @@ export type FakeToolResponse = */ export type FakeProviders = Record>; +/** + * Options for {@link createTestPluginContext} when called with a second parameter. + * When provided, `createTestPluginContext` installs a service context seeded + * from a mock workspace client, plus optional environment variables. + */ +export interface TestPluginContextOptions { + /** + * Responses keyed by dotted path (`"jobs.getRun"`) for the mocked workspace + * client. Passed directly to {@link createMockWorkspaceClient}. + */ + responses?: Record; + /** + * Environment variables to set for the test. Captured on entry, restored + * (or deleted if they were unset) on exit via an `afterEach` hook and/or + * explicit {@link TestPluginContext.restore}. + */ + env?: Record; + /** + * If `true`, throw when a workspace client path with no declared response is + * called, instead of resolving `undefined`. Defaults to `false` (never crash). + */ + strict?: boolean; +} + /** A single dispatch observed by a fake provider. */ export interface RecordedToolCall { /** Registered plugin name (the key in {@link FakeProviders}). */ @@ -144,6 +170,14 @@ export interface TestPluginContext { * gate on `isReady`. Returns the same plugin for chaining. */ attach

(plugin: P): Promise

; + /** + * Restore the service context and environment variables to their pre-test state. + * Called automatically via `afterEach` when options were provided to + * `createTestPluginContext`. Can also be called explicitly for escape hatches + * (e.g., cleanup inside a test body). Idempotent — safe to call multiple times. + * Only present if the context was created with options. + */ + restore?: () => void; } /** @@ -164,18 +198,39 @@ export interface TestPluginContext { * Nothing about `PluginContext` is reimplemented. * * @param fakes - Canned tool responses keyed by plugin then tool name. + * @param options - When provided, installs a mock workspace client seeded from + * `responses`, mocks the service context, and sets `env`. Omit it for the + * original behavior. * * @example * ```ts + * // No options * 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 }); + * + * // With options — installs service context + seeded client + * const mock = createTestPluginContext( + * {}, + * { responses: { "jobs.getRun": { state: "DONE" } } }, + * ); * ``` */ export function createTestPluginContext( fakes: FakeProviders = {}, + options?: TestPluginContextOptions, ): TestPluginContext { + // No options: original behavior. + if (!options) { + return createTestPluginContextSync(fakes); + } + + // Options provided: install the seeded client, service context, and scoped env. + return createTestPluginContextWithOptions(fakes, options); +} + +function createTestPluginContextSync(fakes: FakeProviders): TestPluginContext { const telemetry = createMockTelemetry(); const ctx = new PluginContext({ telemetry }); @@ -351,6 +406,61 @@ export function createTestPluginContext( }; } +function createTestPluginContextWithOptions( + fakes: FakeProviders, + options: TestPluginContextOptions, +): TestPluginContext { + const { responses = {}, env: envVars = {}, strict = false } = options; + + // Build a mock workspace client seeded from responses + const client = createMockWorkspaceClient({ + responses, + strict, + }); + + // Install the mock service context with the seeded client + const serviceContextMock = mockServiceContext({ + serviceDatabricksClient: client, + }); + + // Set env (captured for restore) via the shared helper. + const restoreEnv = applyEnv(envVars); + + // Create the base context (without options this time, since we're handling everything) + const base = createTestPluginContextSync(fakes); + + // Restore function: restores env and service context (idempotent) + let hasRestored = false; + const restore = () => { + if (hasRestored) return; + hasRestored = true; + restoreEnv(); + serviceContextMock.restore(); + }; + + // Auto-restore after the current test. This helper is documented and used + // from inside a test body, where a runtime-registered `afterEach` does NOT run + // for that test (Vitest only collects `afterEach` before the test runs) — the + // reason the old `afterEach` here silently leaked. `onTestFinished` is the hook + // built for runtime registration and fires after the creating test. If called + // at collection scope instead, it throws, so fall back to `afterEach` there. + try { + onTestFinished(() => { + restore(); + }); + } catch { + afterEach(() => { + restore(); + }); + } + + // Return the context with the restore method + return { + ...base, + restore, + }; +} + function cacheReady(): boolean { try { CacheManager.getInstanceSync(); 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..994f950c9 --- /dev/null +++ b/packages/appkit/src/testing/tests/create-test-app.test.ts @@ -0,0 +1,849 @@ +import type { + IAppRequest, + IAppResponse, + IAppRouter, + PluginConstructor, + PluginData, + PluginManifest, +} from "shared"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; + +import { getWorkspaceClient } from "../../context"; +import { getUserContext } from "../../context/execution-context"; +import { ServiceContext } from "../../context/service-context"; +import { AuthenticationError } from "../../errors"; +import { Plugin, toPlugin } from "../../plugin"; +import type { WorkspaceClient } from "../../workspace-client"; +import type { CreateTestAppOptions, TestApp } from "../create-test-app"; +import { createTestApp } from "../create-test-app"; +import { expectStream } from "../expect-stream"; +import { createMockWorkspaceClient, getMock } 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; +} + +/** + * One probe for both halves of the harness: boot/data-plane concerns and the + * HTTP layer. Routes go through `this.route()`, the way real plugins register + * them — that is what wraps handlers in forwardAsyncErrors, so a rejection + * reaches errorHandlerMiddleware instead of hanging the request. + */ +class ProbePlugin extends Plugin { + static manifest = manifest("probe"); + + /** The client this plugin resolved at request time. */ + seenClient: WorkspaceClient | undefined; + + /** Incremented when harness teardown runs this plugin's shutdown() hook. */ + shutdownCalls = 0; + + async shutdown(): Promise { + this.shutdownCalls += 1; + } + + injectRoutes(router: IAppRouter): void { + const r = ( + method: "get" | "post" | "put" | "patch" | "delete", + path: string, + handler: (req: IAppRequest, res: IAppResponse) => Promise, + ) => + this.route(router, { name: `${method}${path}`, method, path, handler }); + + r("get", "/ping", async (_req, res) => { + res.json({ pong: true }); + }); + + // A non-default status, to prove the handler's status propagates. + r("get", "/created", async (_req, res) => { + res.status(201).json({ ok: true, method: "GET" }); + }); + + r("get", "/from-client", async (_req, res) => { + this.seenClient = getWorkspaceClient(); + res.json({ + run: await this.seenClient.jobs.getRun({ run_id: 1 } as never), + }); + }); + + r("post", "/echo", async (req, res) => { + res.json({ + body: req.body, + contentType: req.headers["content-type"] ?? null, + }); + }); + + r("get", "/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, + }); + }); + + // The real asUser path, so the forwarded identity has to be genuine. + r("get", "/as-user", async (req, res) => { + const ex = this.asUser(req).exports() as { + whoami: () => { userId?: string }; + }; + res.json(ex.whoami()); + }); + + // Calls the client *inside* asUser, unlike /as-user which only reads userId. + r("get", "/as-user-client", async (req, res) => { + const ex = this.asUser(req).exports() as { + probeClient: () => Promise; + }; + res.json({ run: (await ex.probeClient()) ?? null }); + }); + + r("get", "/boom", async () => { + throw new Error("handler exploded"); + }); + + r("post", "/stream", 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) { + r(method, "/verb", async (req, res) => { + res.json({ m: method.toUpperCase(), b: req.body }); + }); + } + r("delete", "/verb", async (_req, res) => { + res.json({ m: "DELETE" }); + }); + } + + exports() { + return { + seenClient: () => this.seenClient, + shutdownCalls: () => this.shutdownCalls, + whoami: () => ({ userId: getUserContext()?.userId }), + // Calls through the client so the harness's mock records it — a real + // OBO client would record nothing here. + probeClient: () => + getWorkspaceClient().jobs.getRun({ run_id: 42 } as never), + }; + } +} +const probe = toPlugin(ProbePlugin); + +/** Boot, run the body, always close. Collapses the try/finally every test needs. */ +async function withApp< + P extends PluginData[], +>( + options: CreateTestAppOptions

, + body: (app: TestApp

) => Promise, +): Promise { + const app = await createTestApp(options); + try { + await body(app); + } finally { + await app.close(); + } +} + +/** 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); + +/** + * Boots cleanly, then throws when the harness asks for the socket — the only way + * to reach the failure path *after* `createApp` has already returned an app. + * Named `server` so the harness uses it instead of adding the real one. + */ +class LateFailurePlugin extends Plugin { + static manifest = manifest("server"); + exports() { + return { + getServer: () => { + throw new Error("getServer exploded"); + }, + }; + } +} +const lateFailure = toPlugin(LateFailurePlugin); + +describe("createTestApp", () => { + test("boots with a single plugin and serves a real route", async () => { + await withApp({ plugins: [probe()] }, async (app) => { + expect(app.baseUrl).toMatch(/^http:\/\/127\.0\.0\.1:\d+$/); + expect(app.port).toBeGreaterThan(0); + + const res = await app.get("/api/probe/ping"); + expect(res.status).toBe(200); + await expect(res.json()).resolves.toEqual({ pong: true }); + }); + }); + + test("sequential apps in one file each bind their own ephemeral port", async () => { + // No EADDRINUSE and no hardcoded port, which is why fixed test ports are + // worth removing. Not asserting the two ports differ: the kernel may hand + // back the port just released, so that would flake. + for (const _ of [1, 2]) { + const app = await createTestApp({ plugins: [probe()] }); + try { + expect(app.port).toBeGreaterThan(0); + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await app.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 { + await withApp({ plugins: [probe()] }, async (app) => { + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + }); + } finally { + process.env = saved; + } + }); + + test("the default mock client reaches the plugin instead of crashing", async () => { + await withApp({ plugins: [probe()] }, async (app) => { + const res = await app.get("/api/probe/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({}); + }); + }); + + test("caller-supplied responses reach the plugin's client calls", async () => { + await withApp( + { + plugins: [probe()], + responses: { "jobs.getRun": { state: "TERMINATED" } }, + }, + async (app) => { + const res = await app.get("/api/probe/from-client"); + await expect(res.json()).resolves.toEqual({ + run: { state: "TERMINATED" }, + }); + expect(getMock(app.client, "jobs.getRun")).toHaveBeenCalledWith({ + run_id: 1, + }); + }, + ); + }); + + test("app.client is the same object a handler resolves", async () => { + await withApp({ plugins: [probe()] }, async (app) => { + await app.get("/api/probe/from-client"); + // Retires the "tribal seam knowledge" problem: no need to know that + // createApp({ client }) flows through ServiceContext to reach a handler. + expect(app.plugins.probe.seenClient()).toBe(app.client); + }); + }); + + test("apiClient.request has zero calls after boot", async () => { + await withApp({ plugins: [probe()] }, async (app) => { + // 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(getMock(app.client, "apiClient.request")).toHaveBeenCalledTimes(0); + }); + }); + + test("a caller-supplied server plugin is respected, and dedupes the injected one", async () => { + const { server: serverPlugin } = await import("../../plugins/server"); + await withApp( + { + plugins: [probe(), serverPlugin({ port: 0, host: "127.0.0.1" })], + }, + async (app) => { + expect(app.port).toBeGreaterThan(0); + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + }, + ); + }); + + test("server: false together with a server plugin is refused", async () => { + const { server: serverPlugin } = await import("../../plugins/server"); + await expect( + createTestApp({ + plugins: [probe(), 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 () => { + await withApp({ plugins: [probe()], server: false }, async (app) => { + expect(app.server).toBeUndefined(); + expect(() => app.baseUrl).toThrow(/no HTTP server/); + await expect(app.get("/api/probe/ping")).rejects.toThrow( + /no HTTP server/, + ); + }); + }); + + test("await using releases at scope exit", async () => { + let port: number | undefined; + { + await using app = await createTestApp({ plugins: [probe()] }); + port = app.port; + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + } + await expect(fetch(`http://127.0.0.1:${port}/health`)).rejects.toThrow(); + }); + + test("close runs a booted plugin's shutdown() hook and releases the socket", async () => { + // The path this PR rewired (close() -> disposeApp -> dispose()): a *real* + // registered plugin's shutdown() must fire through the boot->teardown wiring, + // and the server socket must actually be released — not just re-port-picked + // on the next boot. The lifecycle unit test covers this against a mocked + // context; this asserts the composed integration path end to end. + const app = await createTestApp({ plugins: [probe()] }); + const port = app.port; + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + + await app.close(); + + expect(app.plugins.probe.shutdownCalls()).toBe(1); + 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 () => { + await withApp( + { + plugins: [needsEnv(), probe()], + env: { MY_REQUIRED_SECRET: "s3cret" }, + }, + async (app) => { + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + }, + ); + // 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: [probe()], + 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: [probe()] }); + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + await app.close(); + }); + + test("a failure after createApp still tears the built app down", async () => { + // The other boot-failure test throws inside createApp, so `app` is never + // assigned and only the "nothing was built" branch runs. This one gets a + // live app first, exercising the branch that has to close it. + const before = { ...process.env }; + + await expect( + createTestApp({ plugins: [lateFailure()], env: { LEAKED: "no" } }), + ).rejects.toThrow(/getServer exploded/); + + expect(process.env.LEAKED).toBeUndefined(); + expect(Object.keys(process.env).sort()).toEqual( + Object.keys(before).sort(), + ); + + // The failed app's singletons have to be dropped too, or this boot + // reuses its half-closed ones. + const app = await createTestApp({ plugins: [probe()] }); + try { + await expect( + app.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await app.close(); + } + }); + + test("the obo token fingerprint tracks the token, matching production", async () => { + // Lakebase rotates its pool when this value changes (pool-manager compares + // it), so a fingerprint keyed on the user would be constant across tokens + // and rotation could never fire under the harness. + await using app = await createTestApp({ plugins: [probe()] }); + void app; + + // While a harness app is live, createUserContext *is* the stub. + const fingerprint = (token: string) => + ServiceContext.createUserContext(token, "same").tokenFingerprint; + + expect(fingerprint("tok-a")).toBe(fingerprint("tok-a")); + expect(fingerprint("tok-a")).not.toBe(fingerprint("tok-b")); + // A sha256 prefix, not a label derived from the user. + expect(fingerprint("tok-a")).toMatch(/^[0-9a-f]{16}$/); + expect(fingerprint("tok-a")).not.toContain("same"); + }); + + test("an obo request with no token is refused, as in production", async () => { + await using app = await createTestApp({ plugins: [probe()] }); + void app; + // The same error class production throws, not just a similar message — + // a handler that catches AuthenticationError must behave identically here. + expect(() => ServiceContext.createUserContext("", "nobody")).toThrow( + AuthenticationError, + ); + }); + + test("strict makes an undeclared data-plane call fail the request", async () => { + // Without strict the handler gets `undefined` and the route still 200s, + // so a forgotten response passes for the wrong reason. + await using app = await createTestApp({ + plugins: [probe()], + strict: true, + }); + const res = await app.get("/api/probe/from-client"); + expect(res.status).toBe(500); + await expect(res.text()).resolves.toMatch(/no declared response/); + }); + + test("passing both client and responses is refused, not silently ignored", async () => { + // `responses` only seeds the built-in mock, so with a supplied client it + // used to do nothing at all — the caller's seeded values never took effect + // and nothing said so. + await expect( + createTestApp({ + plugins: [probe()], + client: createMockWorkspaceClient(), + responses: { "jobs.getRun": { state: "IGNORED" } }, + }), + ).rejects.toThrow(/do nothing when you also pass `client`/); + }); + + test('nodeEnv: "development" is refused with an explanation', async () => { + // The get-port RangeError must never reach the user. + await expect( + createTestApp({ plugins: [probe()], 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: [probe()] }); + 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: [probe()] }); + const firstPort = first.port; + await first.close(); + + const second = await createTestApp({ plugins: [probe()] }); + try { + expect(second.port).not.toBe(firstPort); + await expect( + second.get("/api/probe/ping").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await second.close(); + } + }); + + test("repeated boot/close cycles leave no env residue", async () => { + const before = { ...process.env }; + + // Overlapping boots are refused, so each snapshot starts from an already + // restored env and the old "whichever closes last wins" hazard cannot + // arise. What still has to hold is that a cycle leaves nothing behind. + for (const [key, value] of [ + ["CYCLE_A", "a"], + ["CYCLE_B", "b"], + ] as const) { + const app = await createTestApp({ + plugins: [probe()], + env: { [key]: value }, + }); + expect(process.env[key]).toBe(value); + await app.close(); + expect(process.env[key]).toBeUndefined(); + } + + const leaked = Object.keys(process.env).filter((k) => !(k in before)); + expect(leaked).toEqual([]); + expect(Object.keys(process.env).sort()).toEqual( + Object.keys(before).sort(), + ); + }); + + test("close() is idempotent", async () => { + const app = await createTestApp({ plugins: [probe()] }); + await app.close(); + await expect(app.close()).resolves.toBeUndefined(); + }); + }); +}); + +describe("createTestApp HTTP layer", () => { + let app: TestApp<[ReturnType]>; + + beforeAll(async () => { + app = await createTestApp({ plugins: [probe()] }); + }); + + afterAll(async () => { + await app?.close(); + }); + + test("GET returns the plugin's JSON body and status", async () => { + const res = await app.get("/api/probe/created"); + 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/probe/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/probe/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", + }); + }); + + // All three hit /headers and differ only in what `obo`/`headers` should produce. + test.each([ + [ + "obo: true sets the forwarded identity", + { obo: true as const }, + { user: "test-user", token: "test-user-token" }, + ], + [ + "obo object overrides the identity", + { obo: { userId: "alice", email: "alice@example.com" } }, + { user: "alice", email: "alice@example.com" }, + ], + [ + "explicit headers win over what obo generated", + { + obo: true as const, + headers: { "x-custom": "hello", "x-forwarded-user": "override" }, + }, + { custom: "hello", user: "override", token: "test-user-token" }, + ], + [ + "a mixed-case override wins too, rather than comma-joining", + { obo: true as const, headers: { "X-Forwarded-User": "override" } }, + { user: "override", token: "test-user-token" }, + ], + ])("%s", async (_name, options, expected) => { + const res = await app.get("/api/probe/headers", options); + await expect(res.json()).resolves.toMatchObject(expected); + }); + + test("a handler using asUser resolves the forwarded test user", async () => { + const res = await app.get("/api/probe/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/probe/stream"); + await expectStream(res).toEmit("status", "result"); + }); + + test("a throwing handler produces the real error-middleware response", async () => { + const res = await app.get("/api/probe/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/probe/nope"); + expect(res.status).toBe(404); + }); + + test("put, patch, and delete reach their handlers", async () => { + await expect( + app.put("/api/probe/verb", { body: { a: 1 } }).then((r) => r.json()), + ).resolves.toEqual({ m: "PUT", b: { a: 1 } }); + await expect( + app.patch("/api/probe/verb", { body: { a: 2 } }).then((r) => r.json()), + ).resolves.toEqual({ m: "PATCH", b: { a: 2 } }); + await expect( + app.delete("/api/probe/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/probe/json", { signal: controller.signal }), + ).rejects.toThrow(); + }); +}); + +describe("createTestApp — one app at a time", () => { + // Top level on purpose: these boot their own apps, so they cannot live + // inside a describe that holds one open in beforeAll. + test("an obo request gets the mock client, not a real one", async () => { + await withApp( + { plugins: [probe()], responses: { "jobs.getRun": { via: "mock" } } }, + async (app) => { + const res = await app.get("/api/probe/as-user-client", { + obo: { userId: "carol" }, + }); + expect(res.status).toBe(200); + + // Asserting the host would not discriminate — a real client built from + // DATABRICKS_HOST carries the same string. What only the mock can do is + // record the call and return the declared response. + await expect(res.json()).resolves.toEqual({ run: { via: "mock" } }); + expect(getMock(app.client, "jobs.getRun")).toHaveBeenCalledWith({ + run_id: 42, + }); + }, + ); + }); + + test("refuses a second app while one is open", async () => { + const a = await createTestApp({ plugins: [probe()] }); + try { + await expect(createTestApp({ plugins: [probe()] })).rejects.toThrow( + /a harness app is already open/, + ); + } finally { + await a.close(); + } + }); + + test("a refused boot leaves the open app fully working", async () => { + // The guard runs before any mutation, so the refusal must not disturb the + // live app's env baseline, singletons, or on-behalf-of fake. Without that + // ordering the first app would be collateral damage of someone else's bug. + const a = await createTestApp({ + plugins: [probe()], + responses: { "jobs.getRun": { via: "mock" } }, + }); + try { + await expect(createTestApp({ plugins: [probe()] })).rejects.toThrow(); + + await expect( + a.get("/api/probe/from-client").then((r) => r.status), + ).resolves.toBe(200); + + // /as-user-client calls the client *inside* asUser, so it is the one + // route that notices a disturbed on-behalf-of fake. Asserting the + // declared response discriminates: a real client cannot invent it. + const oboRes = await a.get("/api/probe/as-user-client", { + obo: { userId: "u@example.com" }, + }); + expect(oboRes.status).toBe(200); + await expect(oboRes.json()).resolves.toEqual({ run: { via: "mock" } }); + expect(getMock(a.client, "jobs.getRun")).toHaveBeenCalledWith({ + run_id: 42, + }); + } finally { + await a.close(); + } + }); + + test("a refused boot does not strand the env baseline", async () => { + const before = { ...process.env }; + const a = await createTestApp({ plugins: [probe()] }); + await expect(createTestApp({ plugins: [probe()] })).rejects.toThrow(); + await a.close(); + + // The refused boot must leave harnessAppLive and the saved baseline alone, + // or a's close would not restore the env. + expect(process.env.NODE_ENV).toBe(before.NODE_ENV); + expect(process.env.DATABRICKS_WORKSPACE_ID).toBe( + before.DATABRICKS_WORKSPACE_ID, + ); + }); + + test("a new app boots cleanly once the previous one closed", async () => { + const a = await createTestApp({ plugins: [probe()] }); + await expect( + a.get("/api/probe/from-client").then((r) => r.status), + ).resolves.toBe(200); + await a.close(); + + // The singletons A dropped on close must be rebuilt for B, not inherited. + const b = await createTestApp({ plugins: [probe()] }); + try { + await expect( + b.get("/api/probe/from-client").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await b.close(); + } + }); + + test("a stale handle's second close cannot reset a newer app", async () => { + const first = await createTestApp({ plugins: [probe()] }); + await first.close(); + + const second = await createTestApp({ plugins: [probe()] }); + try { + // close() is memoized, so this stale call is a no-op — not a drop that + // would reset second's singletons. + await first.close(); + await expect( + second.get("/api/probe/from-client").then((r) => r.status), + ).resolves.toBe(200); + } finally { + await second.close(); + } + }); +}); 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/fixtures.test.ts b/packages/appkit/src/testing/tests/fixtures.test.ts index 39fc2fef1..9f7a906f8 100644 --- a/packages/appkit/src/testing/tests/fixtures.test.ts +++ b/packages/appkit/src/testing/tests/fixtures.test.ts @@ -3,10 +3,15 @@ import { afterEach, describe, expect, test, vi } from "vitest"; import { CacheManager } from "../../cache"; import { InMemoryStorage } from "../../cache/storage"; import { ServiceContext } from "../../context"; +import { AuthenticationError } from "../../errors"; +import { ApiError } from "../../workspace-client"; import { + createApiError, createMockRequest, + mockServiceContext, resetTestCache, useServiceContextMock, + withEnv, } from "../fixtures"; describe("createMockRequest — obo option", () => { @@ -102,6 +107,50 @@ describe("resetTestCache", () => { }); }); +describe("mockServiceContext — user context matches production", () => { + test("the fingerprint is derived from the token, not the user", () => { + // Lakebase rotates its pool by comparing this. A user-keyed value would be + // constant across tokens, so `pool-manager` would never see a change and the + // drain-and-recreate branch could never run under this fake. + const mock = mockServiceContext(); + try { + const fp = (token: string) => + ServiceContext.createUserContext(token, "same").tokenFingerprint; + expect(fp("tok-a")).toBe(fp("tok-a")); + expect(fp("tok-a")).not.toBe(fp("tok-b")); + expect(fp("tok-a")).toMatch(/^[0-9a-f]{16}$/); + } finally { + mock.restore(); + } + }); + + test("a missing token is refused, with production's error class", () => { + const mock = mockServiceContext(); + try { + expect(() => ServiceContext.createUserContext("", "nobody")).toThrow( + AuthenticationError, + ); + } finally { + mock.restore(); + } + }); + + test("userEmail is carried through", () => { + const mock = mockServiceContext(); + try { + const ctx = ServiceContext.createUserContext( + "tok", + "u-1", + "Alice", + "alice@example.com", + ); + expect(ctx.userEmail).toBe("alice@example.com"); + } finally { + mock.restore(); + } + }); +}); + describe("useServiceContextMock", () => { const ctx = useServiceContextMock({ warehouseId: "wh-1" }); @@ -138,3 +187,184 @@ describe("useServiceContextMock — restores after the block", () => { expect(vi.isMockFunction(ServiceContext.isInitialized)).toBe(false); }); }); + +describe("withEnv — environment variable restoration", () => { + test("sets a var inside fn, restores it to its prior value afterward", () => { + const original = process.env.TEST_VAR; + process.env.TEST_VAR = "original"; + + const result = withEnv({ TEST_VAR: "modified" }, () => { + expect(process.env.TEST_VAR).toBe("modified"); + return "done"; + }); + + expect(result).toBe("done"); + expect(process.env.TEST_VAR).toBe("original"); + + // Cleanup + if (original === undefined) { + delete process.env.TEST_VAR; + } else { + process.env.TEST_VAR = original; + } + }); + + test("a key that was UNSET before is deleted (not left set) after fn returns", () => { + if (process.env.NEVER_SET_VAR !== undefined) { + delete process.env.NEVER_SET_VAR; + } + + withEnv({ NEVER_SET_VAR: "temp" }, () => { + expect(process.env.NEVER_SET_VAR).toBe("temp"); + }); + + expect(process.env.NEVER_SET_VAR).toBeUndefined(); + }); + + test("a key that PRE-EXISTED is restored to its original value, not deleted", () => { + process.env.PRE_EXISTING = "before"; + + withEnv({ PRE_EXISTING: "changed" }, () => { + expect(process.env.PRE_EXISTING).toBe("changed"); + }); + + expect(process.env.PRE_EXISTING).toBe("before"); + + // Cleanup + delete process.env.PRE_EXISTING; + }); + + test("restores even when fn throws", () => { + process.env.THROW_TEST = "before"; + + expect(() => { + withEnv({ THROW_TEST: "during" }, () => { + expect(process.env.THROW_TEST).toBe("during"); + throw new Error("test error"); + }); + }).toThrow("test error"); + + expect(process.env.THROW_TEST).toBe("before"); + + // Cleanup + delete process.env.THROW_TEST; + }); + + test("async form: await withEnv({...}, async () => …) restores after the promise settles", async () => { + process.env.ASYNC_TEST = "before"; + + await withEnv({ ASYNC_TEST: "during" }, async () => { + expect(process.env.ASYNC_TEST).toBe("during"); + await new Promise((resolve) => setTimeout(resolve, 10)); + }); + + expect(process.env.ASYNC_TEST).toBe("before"); + + // Cleanup + delete process.env.ASYNC_TEST; + }); + + test("async form restores even when the promise rejects", async () => { + process.env.ASYNC_REJECT_TEST = "before"; + + await expect( + withEnv({ ASYNC_REJECT_TEST: "during" }, async () => { + expect(process.env.ASYNC_REJECT_TEST).toBe("during"); + throw new Error("async error"); + }), + ).rejects.toThrow("async error"); + + expect(process.env.ASYNC_REJECT_TEST).toBe("before"); + + // Cleanup + delete process.env.ASYNC_REJECT_TEST; + }); + + test("nested withEnv calls restore in reverse order (LIFO)", () => { + process.env.NESTED_VAR = "original"; + const log: string[] = []; + + withEnv({ NESTED_VAR: "level1" }, () => { + log.push(`L1-during: ${process.env.NESTED_VAR}`); + + withEnv({ NESTED_VAR: "level2" }, () => { + log.push(`L2-during: ${process.env.NESTED_VAR}`); + }); + + log.push(`L1-after: ${process.env.NESTED_VAR}`); + }); + + log.push(`outside: ${process.env.NESTED_VAR}`); + + expect(log).toEqual([ + "L1-during: level1", + "L2-during: level2", + "L1-after: level1", + "outside: original", + ]); + + // Cleanup + delete process.env.NESTED_VAR; + }); +}); + +describe("createApiError — genuine ApiError factory", () => { + test("returns a genuine ApiError instance", () => { + const error = createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "NOT_FOUND", + }); + expect(error).toBeInstanceOf(ApiError); + }); + + test("preserves statusCode", () => { + const error = createApiError({ + statusCode: 500, + message: "Server error", + errorCode: "INTERNAL_ERROR", + }); + expect(error.statusCode).toBe(500); + }); + + test("preserves message", () => { + const error = createApiError({ + statusCode: 400, + message: "Bad request input", + errorCode: "INVALID_ARGUMENT", + }); + expect(error.message).toBe("Bad request input"); + }); + + test("preserves errorCode", () => { + const error = createApiError({ + statusCode: 403, + message: "Access denied", + errorCode: "PERMISSION_DENIED", + }); + expect(error.errorCode).toBe("PERMISSION_DENIED"); + }); + + test("works in production-shaped instanceof checks", () => { + const error = createApiError({ + statusCode: 404, + message: "Not found", + errorCode: "NOT_FOUND", + }); + // This is the production pattern from files/client.ts + const isNotFoundError = + error instanceof ApiError && error.statusCode === 404; + expect(isNotFoundError).toBe(true); + }); + + test("has sensible defaults for optional fields", () => { + const error = createApiError({ + statusCode: 400, + message: "Bad request", + errorCode: "BAD_REQUEST", + }); + // Should have response (undefined or null) and details (array) + expect(error).toHaveProperty("response"); + expect(error).toHaveProperty("errorInfoType"); + }); +}); 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..5ec90a0d9 --- /dev/null +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -0,0 +1,309 @@ +import { inspect } from "node:util"; + +import { describe, expect, test, vi } from "vitest"; + +import { ServiceContext } from "../../context/service-context"; +import { mockServiceContext } from "../fixtures"; +import { createMockWorkspaceClient, getMock } from "../mock-workspace-client"; + +const mk = createMockWorkspaceClient; +/** Service methods are SDK-typed, so calling an arbitrary one needs a cast. */ +const svc = (client: unknown, name: string) => + (client as Record unknown>>)[ + name + ]; + +const SUCCEEDED = { status: { state: "SUCCEEDED" }, result: { data: [] } }; +const TEST_USER = { id: "test-service-user", userName: "test-service-user" }; + +describe("createMockWorkspaceClient", () => { + describe("the never-crash floor", () => { + // Never-crash is the headline claim, so all nine are asserted, not sampled. + test.each([ + ["files", "listDirectory", undefined], + ["genie", "getMessage", undefined], + ["jobs", "getRun", undefined], + ["servingEndpoints", "get", undefined], + ["warehouses", "getWarehouse", { state: "RUNNING" }], + ["warehouses", "startWarehouse", undefined], + ["statementExecution", "executeStatement", SUCCEEDED], + ["currentUser", "me", TEST_USER], + ])("%s.%s resolves its default", async (service, method, expected) => { + const client = mk(); + expect(client[service as "jobs"]).toBeDefined(); + await expect(svc(client, service)[method]({})).resolves.toEqual(expected); + }); + + test("config and apiClient are reachable, and not mocks where it matters", () => { + const client = mk(); + // Both read directly by production code — a Promise or mock here breaks it. + expect(typeof client.config.host).toBe("string"); + expect(client.config.host).toBeTruthy(); + expect(typeof client.apiClient.userAgent()).toBe("string"); + }); + + test("a seeded userAgent stays synchronous, not a Promise", async () => { + // Resolving it would put "[object Promise]" into a Headers value. + const client = mk({ responses: { "apiClient.userAgent": "custom/9" } }); + expect(client.apiClient.userAgent()).toBe("custom/9"); + const headers = new Headers(); + headers.set("user-agent", client.apiClient.userAgent()); + expect(headers.get("user-agent")).toBe("custom/9"); + // `request` is genuinely async, so its seeds still resolve. + await expect( + mk({ responses: { "apiClient.request": { ok: 1 } } }).apiClient.request( + {} as never, + ), + ).resolves.toEqual({ ok: 1 }); + }); + + test("apiClient.request is depth-2 and destructurable", async () => { + await expect(mk().apiClient.request({} as never)).resolves.toEqual({}); + const client = mk({ + responses: { "apiClient.request": { results: [] } }, + }); + await expect(client.apiClient.request({} as never)).resolves.toEqual({ + results: [], + }); + }); + }); + + describe("responses", () => { + test("a declared value resolves, and overrides a default", async () => { + const client = mk({ + responses: { + "jobs.getRun": { state: "TERMINATED" }, + "statementExecution.executeStatement": { status: { state: "MINE" } }, + }, + }); + await expect(client.jobs.getRun({} as never)).resolves.toEqual({ + state: "TERMINATED", + }); + await expect( + client.statementExecution.executeStatement({} as never), + ).resolves.toEqual({ status: { state: "MINE" } }); + }); + + test("a function receives the arguments, and its rejection propagates", async () => { + const fn = vi.fn().mockResolvedValue({ ok: true }); + await mk({ responses: { "jobs.getRun": fn } }).jobs.getRun({ + run_id: 456, + } as never); + expect(fn).toHaveBeenCalledWith({ run_id: 456 }); + + const err = new Error("boom"); + const rejecting = mk({ + responses: { "jobs.getRun": () => Promise.reject(err) }, + }); + await expect(rejecting.jobs.getRun({} as never)).rejects.toBe(err); + }); + + test("{ defaults: false } leaves the canned paths unresolved", async () => { + const client = mk({ defaults: false }); + await expect( + client.statementExecution.executeStatement({} as never), + ).resolves.toBeUndefined(); + }); + + test("the config option overrides defaults and adds members", () => { + const authenticate = vi.fn(); + const client = mk({ + config: { host: "https://custom.example.com", authenticate }, + }); + expect(client.config.host).toBe("https://custom.example.com"); + expect(client.config.authenticate).toBe(authenticate); + }); + + test('a "config.host" response stays a raw string, not a mock', () => { + expect( + mk({ responses: { "config.host": "https://a.b" } }).config.host, + ).toBe("https://a.b"); + }); + + test("config.authenticate stamps a header; ensureResolved resolves", async () => { + const client = mk(); + const headers = new Headers(); + // Asserting only "was called" would pass against a mock that does nothing. + await client.config.authenticate(headers); + expect(headers.get("Authorization")).toBe("Bearer test-token"); + await expect(client.config.ensureResolved()).resolves.toBeUndefined(); + }); + + test("an unknown member of a seeded namespace still hits the floor", () => { + expect( + typeof (mk().config as never as Record).nope, + ).toBe("function"); + }); + }); + + describe("memoization (call assertions depend on it)", () => { + test("methods, namespaces, and the legacy view share one identity", () => { + const client = mk(); + expect(client.jobs.getRun).toBe(client.jobs.getRun); + expect(client.jobs).toBe(client.jobs); + expect(client.toLegacyWorkspaceClient().jobs.getRun).toBe( + client.jobs.getRun, + ); + }); + + test("un-faceted legacy services also work", async () => { + const legacy = mk().toLegacyWorkspaceClient(); + await expect(svc(legacy, "clusters").list({})).resolves.toBeUndefined(); + }); + }); + + describe("footguns", () => { + test("a service is not thenable, so await does not hang", async () => { + const client = mk(); + expect((client.jobs as never as { then?: unknown }).then).toBeUndefined(); + await expect(Promise.resolve(client.jobs)).resolves.toBe(client.jobs); + }); + + test("formatting and structural equality neither throw nor recurse", () => { + const client = mk(); + // ownKeys stays default, so a service inspects as {} instead of minting a + // mock per probed property. + expect(inspect(client.jobs)).toBe("{}"); + expect(inspect(client.toLegacyWorkspaceClient())).toBe("{}"); + expect(inspect(client)).toContain("https://test.databricks.com"); + expect(() => JSON.stringify(client.config)).not.toThrow(); + expect(() => expect(client.jobs).toEqual({})).not.toThrow(); + + const log = vi.spyOn(console, "log").mockImplementation(() => {}); + try { + expect(() => console.log("%O", client)).not.toThrow(); + } finally { + log.mockRestore(); + } + }); + }); + + describe("strict", () => { + test("an undeclared call throws, naming the path", async () => { + const client = mk({ strict: true, responses: { "jobs.getRun": {} } }); + await expect(client.jobs.getRun({} as never)).resolves.toEqual({}); + expect(() => client.jobs.cancelRun({} as never)).toThrow( + /"jobs.cancelRun" was called with no declared response/, + ); + }); + + test("it throws on call, not on mint, so getMock still hands back a handle", () => { + // Grabbing the handle before the code under test runs is the normal + // pattern; minting must not be what explodes. + const client = mk({ strict: true }); + const fn = getMock(client, "jobs.cancelRun"); + expect(fn).toHaveBeenCalledTimes(0); + expect(() => fn({} as never)).toThrow(/no declared response/); + }); + + test("the canned defaults still count as declared", async () => { + // Otherwise `strict` would break every harness boot, which reads + // currentUser.me through this client. + const client = mk({ strict: true }); + await expect(client.currentUser.me({} as never)).resolves.toMatchObject({ + id: "test-service-user", + }); + }); + + test("with defaults: false even the canned paths are undeclared", async () => { + const client = mk({ strict: true, defaults: false }); + expect(() => client.currentUser.me({} as never)).toThrow( + /no declared response/, + ); + }); + + test("off by default — an undeclared call still resolves undefined", async () => { + await expect(mk().jobs.cancelRun({} as never)).resolves.toBeUndefined(); + }); + }); + + describe("getMock", () => { + test("mints before first use and stays stable after", async () => { + const client = mk(); + const getRun = getMock(client, "jobs.getRun"); + expect(getRun).toHaveBeenCalledTimes(0); + + await client.jobs.getRun({ run_id: 7 } as never); + expect(getRun).toBe(getMock(client, "jobs.getRun")); + expect(getRun).toHaveBeenCalledWith({ run_id: 7 }); + }); + + test("resolves seeded members and rejects non-function paths", () => { + const client = mk(); + expect(getMock(client, "apiClient.request")).toBe( + client.apiClient.request, + ); + expect(() => getMock(client, "config.host")).toThrow( + /not a mocked function/, + ); + expect(() => getMock({} as never, "jobs.getRun")).toThrow( + /not a createMockWorkspaceClient/, + ); + }); + }); + + describe("convergence with mockServiceContext (D4)", () => { + test("the historical canned defaults are byte-identical", async () => { + const client = mk(); + await expect( + client.statementExecution.executeStatement({} as never), + ).resolves.toEqual(SUCCEEDED); + await expect( + client.warehouses.getWarehouse({} as never), + ).resolves.toEqual({ + state: "RUNNING", + }); + await expect( + client.warehouses.startWarehouse({} as never), + ).resolves.toBeUndefined(); + }); + + test("the service and user clients are both faked, and neither crashes", async () => { + const mock = mockServiceContext(); + try { + // Before convergence this threw "Cannot read properties of undefined". + await expect( + mock.serviceContext.client.jobs.getRun({} as never), + ).resolves.toBeUndefined(); + await expect( + mock.serviceContext.client.statementExecution.executeStatement( + {} as never, + ), + ).resolves.toMatchObject({ status: { state: "SUCCEEDED" } }); + + const user = ServiceContext.createUserContext("tok", "u-1", "alice"); + await expect( + user.client.jobs.getRun({} as never), + ).resolves.toBeUndefined(); + } finally { + mock.restore(); + } + }); + }); + + /** + * Enforced by `tsc --noEmit`, not at runtime: a `@ts-expect-error` that stops + * being an error fails the typecheck. + */ + describe("compile-time contract", () => { + test("unknown members and misspelled methods are compile errors", () => { + const client = mk(); + + // @ts-expect-error - `jbos` is not a facade member + expect(client.jbos).toBeUndefined(); + // @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; + + // `host` is `string | undefined` in the SDK, so the honest claim is that it + // narrows to a string — not that it is non-optional. + const host = client.config.host; + expect(typeof host).toBe("string"); + + const getRun = getMock(client, "jobs.getRun"); + getRun.mockResolvedValue({ state: "TERMINATED" }); + expect(getRun.mock.calls).toEqual([]); + }); + }); +}); 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..52f5f1925 --- /dev/null +++ b/packages/appkit/src/testing/tests/published-surface.integration.test.ts @@ -0,0 +1,157 @@ +import * as testing from "@databricks/appkit/testing"; +import { + createMockRequest, + createTestApp, + createTestPluginContext, + expectStream, + getMock, +} 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("every documented export is reachable from the entry", () => { + // Importing a few symbols proves the entry resolves, not that the surface is + // intact — everything else could be dropped from the barrel and this file + // would still pass. `tsc` would catch it via other suites, but this test + // claims to guard the surface, so it should. + const expected = [ + "createTestApp", + "createTestPlugin", + "createTestPluginContext", + "createMockWorkspaceClient", + "getMock", + "getListeningPort", + "expectStream", + "mockServiceContext", + "createMockRequest", + "createMockResponse", + "createMockRouter", + "createMockTelemetry", + "createSuccessfulSQLResponse", + "createFailedSQLResponse", + "createApiError", + "parseSSEResponse", + "resetTestCache", + "runWithRequestContext", + "setupDatabricksEnv", + "useServiceContextMock", + "useTestApp", + "useTestCache", + "withEnv", + ]; + const missing = expected.filter( + (name) => + typeof (testing as Record)[name] !== "function", + ); + expect(missing).toEqual([]); + }); + + test("createTestPluginContext runs its real dispatch through the entry", async () => { + // The name check above only proves the barrel exports *something*. This + // drives the context's real tool registry and on-behalf-of path, so a + // hollowed-out export fails here instead of shipping. + const mock = createTestPluginContext({ + widget: { lookup: (args) => ({ echoed: args }) }, + }); + + const req = createMockRequest({ obo: { userId: "analyst@example.com" } }); + const result = await mock.ctx.executeTool( + req as never, + "widget", + "lookup", + { + id: 7, + }, + ); + + expect(result).toEqual({ echoed: { id: 7 } }); + expect(mock.toolCalls[0]).toMatchObject({ + plugin: "widget", + tool: "lookup", + asUser: true, + userId: "analyst@example.com", + }); + }); + + 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(getMock(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-app.test.ts b/packages/appkit/src/testing/tests/test-app.test.ts new file mode 100644 index 000000000..0bd75d0e1 --- /dev/null +++ b/packages/appkit/src/testing/tests/test-app.test.ts @@ -0,0 +1,87 @@ +import type { PluginManifest } from "shared"; +import { describe, expect, test } from "vitest"; + +import { Plugin, toPlugin } from "../../plugin"; +import { createTestApp } from "../create-test-app"; +import { useTestApp } from "../test-app"; + +/** + * The behaviour that matters is the hook wiring: a fresh app per test, closed + * after each, with no `close()` for the caller to forget. The harness allows one + * open app at a time, so "the previous test's app was really closed" is + * observable — a leak makes the next boot throw. + */ + +class ProbePlugin extends Plugin { + static manifest = { + name: "probe", + displayName: "Probe", + version: "0.0.0", + description: "useTestApp probe", + resources: { required: [], optional: [] }, + } as unknown as PluginManifest; + + injectRoutes(router: never): void { + this.route(router, { + name: "ping", + method: "get", + path: "/ping", + handler: async (_req, res) => { + res.json({ pong: true }); + }, + }); + } +} +const probe = toPlugin(ProbePlugin); + +describe("useTestApp", () => { + const app = useTestApp({ plugins: [probe()] }); + + test("boots an app and serves a route inside a test", async () => { + const res = await app.current.get("/api/probe/ping"); + expect(res.status).toBe(200); + await expect(res.json()).resolves.toEqual({ pong: true }); + }); + + test("the previous test's app was closed, so this boot succeeded", async () => { + // If afterEach had not closed it, the one-app-at-a-time guard would have + // thrown during this test's beforeEach and never reached the body. + expect(app.current.port).toBeGreaterThan(0); + }); + + test("hands out a different app than the previous test", async () => { + const res = await app.current.get("/api/probe/ping"); + expect(res.status).toBe(200); + }); +}); + +describe("useTestApp passes options through", () => { + const app = useTestApp({ + plugins: [probe()], + env: { USE_TEST_APP_PROBE: "set" }, + }); + + test("env reaches the boot", () => { + expect(process.env.USE_TEST_APP_PROBE).toBe("set"); + }); +}); + +describe("useTestApp cleans up after the file's suites", () => { + test("env from the previous suite was restored on close", () => { + expect(process.env.USE_TEST_APP_PROBE).toBeUndefined(); + }); + + test("no app is held open, so a manual boot is allowed", async () => { + // The guard makes this the discriminating check: it only passes if every + // app useTestApp booted above was actually closed. + await using manual = await createTestApp({ plugins: [probe()] }); + expect(manual.port).toBeGreaterThan(0); + }); +}); + +describe("useTestApp misuse", () => { + test("reading .current outside a registered test explains itself", () => { + const stray = useTestApp({ plugins: [probe()] }); + expect(() => stray.current).toThrow(/no active app/); + }); +}); diff --git a/packages/appkit/src/testing/tests/test-cache.test.ts b/packages/appkit/src/testing/tests/test-cache.test.ts new file mode 100644 index 000000000..d0c3732fd --- /dev/null +++ b/packages/appkit/src/testing/tests/test-cache.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, test, vi } from "vitest"; + +import { CacheManager } from "../../cache"; +import { useTestCache } from "../test-cache"; + +describe("useTestCache", () => { + const testCache = useTestCache(); + + test("boots the cache and exposes it inside a test", () => { + expect(testCache.current).toBeInstanceOf(CacheManager); + }); + + test("generateKey is production's key fn: stable for equal parts, differs by userKey", () => { + const a = testCache.current.generateKey(["op", 1], "user-1"); + const b = testCache.current.generateKey(["op", 1], "user-1"); + const c = testCache.current.generateKey(["op", 1], "user-2"); + expect(a).toBe(b); + expect(a).not.toBe(c); + }); + + test("getOrExecute caches: same key runs fn once, different keys run it twice", async () => { + const fn = vi.fn(async () => "value"); + await testCache.current.getOrExecute(["op", 1], fn, "user-1"); + await testCache.current.getOrExecute(["op", 1], fn, "user-1"); + expect(fn).toHaveBeenCalledTimes(1); + + const fn2 = vi.fn(async () => "value2"); + await testCache.current.getOrExecute(["op", 2], fn2, "user-1"); + await testCache.current.getOrExecute(["op", 3], fn2, "user-1"); + expect(fn2).toHaveBeenCalledTimes(2); + }); + + test("is spy-able: getOrExecute records the key parts a caller passes", async () => { + const spy = vi.spyOn(testCache.current, "getOrExecute"); + await testCache.current.getOrExecute( + ["listing", "/a"], + async () => 1, + "svc", + ); + expect(spy).toHaveBeenCalledWith( + ["listing", "/a"], + expect.any(Function), + "svc", + ); + }); +}); + +// Proves the per-test clear: these two tests run in source order (Vitest's +// default within a file), and the second must not see the first's write. +describe("useTestCache clears between tests", () => { + const testCache = useTestCache(); + const parts = ["shared"]; + const user = "u"; + + test("writes a value", async () => { + const key = testCache.current.generateKey(parts, user); + await testCache.current.set(key, "first"); + expect(await testCache.current.get(key)).toBe("first"); + }); + + test("does not see the previous test's value", async () => { + const key = testCache.current.generateKey(parts, user); + expect(await testCache.current.get(key)).toBeNull(); + }); +}); + +describe("useTestCache keeps the singleton", () => { + useTestCache(); + + test("getInstanceSync still returns an instance during a test", () => { + expect(() => CacheManager.getInstanceSync()).not.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 index 60b30b917..fe276a90b 100644 --- a/packages/appkit/src/testing/tests/test-plugin-context.test.ts +++ b/packages/appkit/src/testing/tests/test-plugin-context.test.ts @@ -1,9 +1,10 @@ import type express from "express"; -import { describe, expect, test } from "vitest"; +import { describe, expect, test, vi } from "vitest"; import { PluginContext } from "../../core/plugin-context"; import { Plugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; +import { createMockRequest } from "../fixtures"; import { createTestPluginContext } from "../test-plugin-context"; // A minimal real plugin for exercising attach() end-to-end. @@ -40,11 +41,7 @@ function mockReq( "x-forwarded-user": "alice", }, ): express.Request { - return { - body: {}, - headers, - header: (name: string) => headers[name.toLowerCase()], - } as unknown as express.Request; + return createMockRequest({ headers }) as unknown as express.Request; } describe("createTestPluginContext — construction", () => { @@ -327,3 +324,213 @@ describe("createTestPluginContext — attach()", () => { expect(result).toBe("fake"); }); }); + +describe("createTestPluginContext — optional second parameter (options overload)", () => { + test("returns synchronously in both forms (not a promise)", () => { + const noOptions = createTestPluginContext(); + expect(noOptions).not.toBeInstanceOf(Promise); + expect(noOptions.ctx).toBeInstanceOf(PluginContext); + + const withOptions = createTestPluginContext( + {}, + { responses: { "jobs.getRun": { state: "DONE" } } }, + ); + expect(withOptions).not.toBeInstanceOf(Promise); + expect(withOptions.ctx).toBeInstanceOf(PluginContext); + withOptions.restore?.(); + }); + + test("no-options call behaves exactly as before (non-breaking)", async () => { + const mock = createTestPluginContext({ + analytics: { query: [{ id: 1 }] }, + }); + + const result = await mock.ctx.executeTool( + mockReq(), + "analytics", + "query", + {}, + ); + + expect(result).toEqual([{ id: 1 }]); + expect(mock.toolCalls).toHaveLength(1); + }); + + test("with options and responses, installs a mock workspace client seeded from responses", async () => { + createTestPluginContext( + {}, + { + responses: { + "jobs.getRun": { job_id: 42, state: "RUNNING" }, + }, + }, + ); + + // The service context is installed, so getWorkspaceClient() returns the mocked client. + const { getWorkspaceClient } = await import("../../context"); + const run = await getWorkspaceClient().jobs.getRun({ run_id: 42 } as never); + expect(run).toEqual({ job_id: 42, state: "RUNNING" }); + }); + + test("with env in options, sets env vars during the test", async () => { + const prior = process.env.TEST_VAR_ABC; + delete process.env.TEST_VAR_ABC; + + createTestPluginContext( + {}, + { + env: { TEST_VAR_ABC: "test-value" }, + }, + ); + + expect(process.env.TEST_VAR_ABC).toBe("test-value"); + + // Cleanup + delete process.env.TEST_VAR_ABC; + if (prior !== undefined) { + process.env.TEST_VAR_ABC = prior; + } + }); + + test("env is restored after the test", async () => { + const prior = process.env.TEST_VAR_XYZ; + delete process.env.TEST_VAR_XYZ; + + const mock = createTestPluginContext( + {}, + { + env: { TEST_VAR_XYZ: "value1" }, + }, + ); + + expect(process.env.TEST_VAR_XYZ).toBe("value1"); + + // Call restore explicitly + if ("restore" in mock && typeof mock.restore === "function") { + mock.restore(); + } + + // After explicit restore, env should be gone (it was unset before) + expect(process.env.TEST_VAR_XYZ).toBeUndefined(); + + // Cleanup + if (prior !== undefined) { + process.env.TEST_VAR_XYZ = prior; + } + }); + + test("restore() is idempotent", async () => { + const prior = process.env.TEST_VAR_IDEMPOTENT; + process.env.TEST_VAR_IDEMPOTENT = "prior"; + + const mock = createTestPluginContext( + {}, + { + env: { TEST_VAR_IDEMPOTENT: "changed" }, + }, + ); + + expect(process.env.TEST_VAR_IDEMPOTENT).toBe("changed"); + + if ("restore" in mock && typeof mock.restore === "function") { + mock.restore(); + } + + expect(process.env.TEST_VAR_IDEMPOTENT).toBe("prior"); + + // Calling restore again should not error + if ("restore" in mock && typeof mock.restore === "function") { + mock.restore(); + } + + expect(process.env.TEST_VAR_IDEMPOTENT).toBe("prior"); + + // Cleanup + if (prior !== undefined) { + process.env.TEST_VAR_IDEMPOTENT = prior; + } else { + delete process.env.TEST_VAR_IDEMPOTENT; + } + }); + + test("strict: true passes through to the mock client", async () => { + createTestPluginContext( + {}, + { + responses: { "jobs.getRun": { state: "DONE" } }, + strict: true, + }, + ); + + const { getWorkspaceClient } = await import("../../context"); + const client = getWorkspaceClient(); + + // Declared response should work + const run = await client.jobs.getRun({ run_id: 42 } as never); + expect(run).toEqual({ state: "DONE" }); + + // Undeclared path should throw when called with strict: true + // Access the method through the facade (which returns a service proxy) + const undeclaredFn = (client as any).warehouses.undeclaredMethod; + try { + await undeclaredFn({ foo: "bar" }); + expect.fail("Should have thrown"); + } catch (err) { + expect((err as Error).message).toContain("no declared response"); + } + }); + + // Proves the options overload auto-restores the service-context spies after + // the creating test WITHOUT a manual restore(). Split across two ordered tests + // because the cleanup fires between them: the second test would fail if the + // hook did not run for the first (the runtime-`afterEach` bug this replaced). + describe("service context auto-restore (no manual restore)", () => { + test("the mock is active inside the test that created it", async () => { + const { ServiceContext } = await import("../../context/service-context"); + + // Intentionally NO manual restore() — auto-cleanup must handle it. + createTestPluginContext( + {}, + { responses: { "jobs.getRun": { state: "DONE" } } }, + ); + + expect(ServiceContext.isInitialized()).toBe(true); + expect(vi.isMockFunction(ServiceContext.isInitialized)).toBe(true); + }); + + test("the previous test's service-context spies were auto-restored", async () => { + const { ServiceContext } = await import("../../context/service-context"); + + // If auto-restore fired after the test above, the spy is gone and the real + // static method is back. + expect(vi.isMockFunction(ServiceContext.isInitialized)).toBe(false); + }); + }); + + test("combines fakes and responses in a single call", async () => { + const mock = createTestPluginContext( + { + analytics: { query: [{ result: "fake" }] }, + }, + { + responses: { + "jobs.getRun": { state: "DONE" }, + }, + }, + ); + + // Fakes work + const fakeResult = await mock.ctx.executeTool( + mockReq(), + "analytics", + "query", + {}, + ); + expect(fakeResult).toEqual([{ result: "fake" }]); + + // Responses work (via seeded mock client in service context) + const { getWorkspaceClient } = await import("../../context"); + const run = await getWorkspaceClient().jobs.getRun({ run_id: 42 } as never); + expect(run).toEqual({ state: "DONE" }); + }); +}); diff --git a/packages/appkit/src/type-generator/mv-registry/describe.ts b/packages/appkit/src/type-generator/mv-registry/describe.ts index db2302761..a01f26190 100644 --- a/packages/appkit/src/type-generator/mv-registry/describe.ts +++ b/packages/appkit/src/type-generator/mv-registry/describe.ts @@ -25,7 +25,7 @@ export function parseDescribeTableExtendedJson( throw new Error(`DESCRIBE TABLE EXTENDED failed: ${msg}`); } - const rows = response.result?.data_array ?? []; + const rows = response.result?.dataArray ?? []; if (rows.length === 0) { throw new Error( "DESCRIBE TABLE EXTENDED returned no rows. Verify the FQN points to a metric view.", diff --git a/packages/appkit/src/type-generator/query-registry.ts b/packages/appkit/src/type-generator/query-registry.ts index db9cf57b9..b822a9986 100644 --- a/packages/appkit/src/type-generator/query-registry.ts +++ b/packages/appkit/src/type-generator/query-registry.ts @@ -158,7 +158,7 @@ function formatParametersType(sql: string): string { /** * Decode a base64 Arrow IPC attachment from a DESCRIBE QUERY response and * extract column metadata. Returns the same shape as rows parsed from the - * legacy data_array path. + * legacy dataArray path. * * IMPORTANT: a DESCRIBE QUERY response is itself a result *table* with rows * shaped like `(col_name, data_type, comment)` describing the user query's @@ -196,7 +196,7 @@ export function convertToQueryType( sql: string, queryName: string, ): { type: string; hasResults: boolean } { - const dataRows = result.result?.data_array || []; + const dataRows = result.result?.dataArray || []; let columns = dataRows.map((row) => ({ name: row[0] || "", type_name: row[1]?.toUpperCase() || "STRING", @@ -204,10 +204,10 @@ export function convertToQueryType( })); // Fallback: serverless warehouses return ARROW_STREAM format with an inline - // base64 attachment instead of data_array. Decode the Arrow IPC rows (the + // base64 attachment instead of dataArray. Decode the Arrow IPC rows (the // DESCRIBE QUERY result table) to extract column names and types. if (columns.length === 0 && result.result?.attachment) { - logger.debug("data_array empty, decoding Arrow IPC attachment for schema"); + logger.debug("dataArray empty, decoding Arrow IPC attachment for schema"); try { columns = columnsFromArrowAttachment(result.result.attachment); } catch (err) { @@ -849,7 +849,7 @@ export async function generateQueriesFromDescribe( "DESCRIBE result for %s: state=%s, rows=%d, hasAttachment=%s", queryName, result.status.state, - result.result?.data_array?.length ?? 0, + result.result?.dataArray?.length ?? 0, !!result.result?.attachment, ); diff --git a/packages/appkit/src/type-generator/statement-result.ts b/packages/appkit/src/type-generator/statement-result.ts index 7ae091aaf..9f5f791fe 100644 --- a/packages/appkit/src/type-generator/statement-result.ts +++ b/packages/appkit/src/type-generator/statement-result.ts @@ -7,18 +7,18 @@ const logger = createLogger("type-generator:statement-result"); /** * Normalize a Statement Execution response so downstream parsers can always - * read rows from `result.data_array`, regardless of the wire format the + * read rows from `result.dataArray`, regardless of the wire format the * warehouse chose. * * `@databricks/sdk-experimental`'s `executeStatement` defaults to an * `ARROW_STREAM` disposition. With an `INLINE` disposition the single * DESCRIBE row is returned as a base64-encoded Arrow IPC stream in - * `result.attachment` and `result.data_array` is left undefined. The metric - * and query type generators only ever read `result.data_array`, so without + * `result.attachment` and `result.dataArray` is left undefined. The metric + * and query type generators only ever read `result.dataArray`, so without * this normalization an Arrow response reads as "returned no rows" — the * registry ships empty and the runtime fail-closed gate 503s every affected * metric/query. (A warehouse configured to return `JSON_ARRAY` populates - * `data_array` directly and needs no decoding — that path, and every mocked + * `dataArray` directly and needs no decoding — that path, and every mocked * test, flows through here unchanged.) */ export async function normalizeResultRows( @@ -30,18 +30,18 @@ export async function normalizeResultRows( // types. A deliberate throw — unlike the best-effort decode below — that both // callers catch per-entry as a loud per-key/per-query failure. if ( - response.result?.next_chunk_index != null || - response.result?.next_chunk_internal_link != null + response.result?.nextChunkIndex != null || + response.result?.nextChunkInternalLink != null ) { throw new Error( - "DESCRIBE result is multi-chunk (truncated); refusing to emit partial types — see next_chunk_index", + "DESCRIBE result is multi-chunk (truncated); refusing to emit partial types — see nextChunkIndex", ); } // Passthrough: rows already materialized (JSON_ARRAY warehouses + every - // mocked test). `data_array` being an empty array still counts as present — + // mocked test). `dataArray` being an empty array still counts as present — // that is a genuine "no rows" answer we must not overwrite with a decode. - if (response.result?.data_array !== undefined) { + if (response.result?.dataArray !== undefined) { return response; } @@ -78,7 +78,7 @@ export async function normalizeResultRows( ...response, result: { ...response.result, - data_array: dataArray, + dataArray: dataArray, }, }; } catch (err) { @@ -149,7 +149,7 @@ function isFormatRejection( /** * Run a DESCRIBE and return a response whose rows are readable via - * `result.data_array`, adapting to the warehouse's result-format capability. + * `result.dataArray`, adapting to the warehouse's result-format capability. * * No single format is portable: standard DBSQL (PRO/CLASSIC) serves * `INLINE`+`JSON_ARRAY` and rejects `INLINE`+`ARROW_STREAM`; the Reyden engine @@ -175,12 +175,16 @@ export async function describeAdaptive( let lastError: unknown; for (const format of formats) { try { + // Narrow the modular SDK's camelCase StatementResponse straight onto our + // subset. The only gap is `dataArray` cells (the SDK types them as + // `JsonValue[][]`); for a DESCRIBE they are always string/null, so the + // assertion is safe. `attachment` survives via the pinned pnpm patch. const response = (await client.statementExecution.executeStatement({ statement, - warehouse_id: warehouseId, + warehouseId, // Synchronous wait: without it the call can return PENDING/RUNNING with // no rows, which downstream misreads as a no-result degrade. - wait_timeout: "30s", + waitTimeout: "30s", format, disposition: "INLINE", })) as DatabricksStatementExecutionResponse; @@ -189,7 +193,7 @@ export async function describeAdaptive( normalized.status?.state === "FAILED" && isFormatRejection( normalized.status.error?.message, - normalized.status.error?.error_code, + normalized.status.error?.errorCode, ) ) { lastResponse = normalized; diff --git a/packages/appkit/src/type-generator/tests/generate-queries.test.ts b/packages/appkit/src/type-generator/tests/generate-queries.test.ts index 9117cbcb9..d29b12cd7 100644 --- a/packages/appkit/src/type-generator/tests/generate-queries.test.ts +++ b/packages/appkit/src/type-generator/tests/generate-queries.test.ts @@ -30,7 +30,10 @@ vi.mock("../../workspace-client", async (importOriginal) => { ...actual, createWorkspaceClient: () => ({ statementExecution: { executeStatement: mocks.executeStatement }, - warehouses: { get: mocks.getWarehouse, start: mocks.startWarehouse }, + warehouses: { + getWarehouse: mocks.getWarehouse, + startWarehouse: mocks.startWarehouse, + }, }), }; }); @@ -82,15 +85,15 @@ const lastSavedQueries = () => function succeededResult(columns: [string, string, string | null][]) { return { - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "SUCCEEDED" }, - result: { data_array: columns }, + result: { dataArray: columns }, }; } /** * Build a SUCCEEDED DESCRIBE QUERY response whose rows arrive only as a base64 - * Arrow IPC `attachment` (no `data_array`) — the ARROW_STREAM/INLINE wire shape + * Arrow IPC `attachment` (no `dataArray`) — the ARROW_STREAM/INLINE wire shape * the fetcher now requests. The describeOne path pipes this through * normalizeResultRows, which decodes the attachment so convertToQueryType can * read the columns. Each [name, type, comment] triple becomes one DESCRIBE row. @@ -108,10 +111,10 @@ async function succeededArrowAttachmentResult( "base64", ); return { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, - // No data_array — rows live in the attachment, like a real INLINE Arrow + // No dataArray — rows live in the attachment, like a real INLINE Arrow // response. This is the condition the silent-degrade bug left unread. result: { attachment }, }; @@ -157,7 +160,7 @@ describe("generateQueriesFromDescribe", () => { test("ARROW attachment path — decodes Arrow rows into a real query schema", async () => { // The warehouse answers ARROW_STREAM/INLINE: columns arrive only as a - // base64 Arrow IPC attachment with data_array undefined. describeOne pipes + // base64 Arrow IPC attachment with dataArray undefined. describeOne pipes // this through normalizeResultRows before convertToQueryType, so the schema // resolves to real columns instead of the degraded `result: unknown`. mocks.readdir.mockResolvedValue(["users.sql"]); @@ -203,8 +206,8 @@ describe("generateQueriesFromDescribe", () => { expect(mocks.executeStatement).toHaveBeenCalledTimes(1); expect(mocks.executeStatement.mock.calls[0][0]).toMatchObject({ - warehouse_id: "wh-123", - wait_timeout: "30s", + warehouseId: "wh-123", + waitTimeout: "30s", format: "JSON_ARRAY", disposition: "INLINE", }); @@ -214,7 +217,7 @@ describe("generateQueriesFromDescribe", () => { mocks.readdir.mockResolvedValue(["bad_table.sql"]); mocks.readFile.mockResolvedValue("SELECT * FROM bad_table"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-2", + statementId: "stmt-2", status: { state: "FAILED", error: { message: "Table or view not found: bad_table" }, @@ -234,7 +237,7 @@ describe("generateQueriesFromDescribe", () => { mocks.readdir.mockResolvedValue(["query.sql"]); mocks.readFile.mockResolvedValue("SELECT 1"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-3", + statementId: "stmt-3", status: { state: "FAILED" }, }); @@ -256,7 +259,7 @@ describe("generateQueriesFromDescribe", () => { mocks.executeStatement .mockResolvedValueOnce(succeededResult([["id", "INT", null]])) .mockResolvedValueOnce({ - statement_id: "stmt-fail", + statementId: "stmt-fail", status: { state: "FAILED", error: { message: "Table not found" }, @@ -288,7 +291,7 @@ describe("generateQueriesFromDescribe", () => { mocks.executeStatement .mockRejectedValueOnce(new Error("Connection refused")) .mockResolvedValueOnce({ - statement_id: "stmt-fail-2", + statementId: "stmt-fail-2", status: { state: "FAILED", error: { message: "Table not found" } }, }); @@ -468,7 +471,7 @@ describe("generateQueriesFromDescribe", () => { .mockResolvedValueOnce("SELECT * FROM whatever"); mocks.executeStatement .mockResolvedValueOnce({ - statement_id: "stmt-syntax", + statementId: "stmt-syntax", status: { state: "FAILED", error: { message: "Table not found" }, @@ -625,7 +628,7 @@ describe("generateQueriesFromDescribe", () => { // state with no result rows. Must degrade like a transient outage, not be // misreported as EMPTY (which would discard a good cached type). mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-1", + statementId: "stmt-1", status: { state: "PENDING" }, }); @@ -660,7 +663,7 @@ describe("generateQueriesFromDescribe", () => { mocks.executeStatement .mockResolvedValueOnce(succeededResult([["id", "INT", null]])) .mockResolvedValueOnce({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "RUNNING" }, }); @@ -687,7 +690,7 @@ describe("generateQueriesFromDescribe", () => { mocks.readdir.mockResolvedValue(["broken.sql"]); mocks.readFile.mockResolvedValue("SELECT * FROM missing"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { message: "Table or view not found: missing" }, diff --git a/packages/appkit/src/type-generator/tests/index.test.ts b/packages/appkit/src/type-generator/tests/index.test.ts index 21b5a7e0c..088b2bcf5 100644 --- a/packages/appkit/src/type-generator/tests/index.test.ts +++ b/packages/appkit/src/type-generator/tests/index.test.ts @@ -298,10 +298,10 @@ describe("generateFromEntryPoint — metric-view emission", () => { const metricFile = path.join(metricsDir, "generated", "metric-views.d.ts"); const describeResponse: DatabricksStatementExecutionResponse = { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ [ JSON.stringify({ columns: [ @@ -460,7 +460,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { // the statement still PENDING — no rows yet. Previously this fell // into the "returned no rows" failure with per-key warns. metricFetcher: async () => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }), }), @@ -568,7 +568,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { expect(mocks.executeStatement).toHaveBeenCalledWith( expect.objectContaining({ statement: "DESCRIBE TABLE EXTENDED `demo`.`sales`.`revenue` AS JSON", - warehouse_id: "wh-1", + warehouseId: "wh-1", }), ); const declarations = fs.readFileSync(metricFile, "utf-8"); @@ -708,7 +708,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { warehouseId: "wh-1", mode: "blocking", metricFetcher: async () => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }), }), @@ -893,7 +893,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { // The fall-through DESCRIBE hits a still-cold warehouse: non-terminal // response, which classifies as degraded (never an error). mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1157,10 +1157,10 @@ describe("generateFromEntryPoint — metric cache section", () => { const describeResponseFor = ( measure: string, ): DatabricksStatementExecutionResponse => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ [ JSON.stringify({ columns: [ @@ -1494,26 +1494,26 @@ describe("generateFromEntryPoint — metric cache section", () => { } if (statement.includes("failed_stmt")) { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }; } if (statement.includes("no_rows")) { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }; } if (statement.includes("no_columns")) { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [[JSON.stringify({ unrelated: true })]] }, + result: { dataArray: [[JSON.stringify({ unrelated: true })]] }, }; } if (statement.includes("pending")) { - return { statement_id: "stmt-mock", status: { state: "PENDING" } }; + return { statementId: "stmt-mock", status: { state: "PENDING" } }; } return describeResponseFor("total_revenue"); }, @@ -1573,7 +1573,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const firstWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1610,7 +1610,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1627,7 +1627,7 @@ describe("generateFromEntryPoint — metric cache section", () => { vi.clearAllMocks(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const error = await run({ mode: "blocking" }).then( @@ -1665,7 +1665,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -1707,7 +1707,7 @@ describe("generateFromEntryPoint — metric cache section", () => { writeConfig({ revenue: { source: "demo.sales.revenue" } }); mocks.getWarehouseState.mockResolvedValue("RUNNING"); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -2013,7 +2013,7 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { warehouseId: "wh-1", mode: "blocking", metricFetcher: async () => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }), }), @@ -2060,10 +2060,10 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { ); const describeResponse: DatabricksStatementExecutionResponse = { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ [ JSON.stringify({ columns: [ diff --git a/packages/appkit/src/type-generator/tests/mv-registry.test.ts b/packages/appkit/src/type-generator/tests/mv-registry.test.ts index 90b07185c..50c69fe1c 100644 --- a/packages/appkit/src/type-generator/tests/mv-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/mv-registry.test.ts @@ -45,10 +45,10 @@ function mockDescribeResponse( payload: unknown, ): DatabricksStatementExecutionResponse { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, result: { - data_array: [[JSON.stringify(payload)]], + dataArray: [[JSON.stringify(payload)]], }, }; } @@ -359,7 +359,7 @@ describe("resolveMetricConfig — FQN naming (UC-accurate)", () => { // crashing the pass — exactly the pre-existing degrade behavior. const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); const { schemas, failures } = await syncMetrics(resolution, fetcher); @@ -556,7 +556,7 @@ describe("parseDescribeTableExtendedJson", () => { test("throws on a FAILED status", () => { expect(() => parseDescribeTableExtendedJson({ - statement_id: "x", + statementId: "x", status: { state: "FAILED", error: { message: "no such table" } }, }), ).toThrowError(/no such table/); @@ -565,9 +565,9 @@ describe("parseDescribeTableExtendedJson", () => { test("throws when the response is empty", () => { expect(() => parseDescribeTableExtendedJson({ - statement_id: "x", + statementId: "x", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }), ).toThrowError(/no rows/); }); @@ -575,9 +575,9 @@ describe("parseDescribeTableExtendedJson", () => { test("throws when the cell is not a JSON string", () => { expect(() => parseDescribeTableExtendedJson({ - statement_id: "x", + statementId: "x", status: { state: "SUCCEEDED" }, - result: { data_array: [[null]] }, + result: { dataArray: [[null]] }, }), ).toThrowError(/JSON string/); }); @@ -680,8 +680,8 @@ describe("createWorkspaceDescribeFetcher", () => { expect(statements).toHaveLength(1); expect(statements[0]).toMatchObject({ statement: "DESCRIBE TABLE EXTENDED `demo`.`sales`.`revenue` AS JSON", - warehouse_id: "wh-1", - wait_timeout: "30s", + warehouseId: "wh-1", + waitTimeout: "30s", // describeAdaptive tries JSON_ARRAY first (standard DBSQL); it falls back // to ARROW_STREAM only if the warehouse rejects that format. format: "JSON_ARRAY", @@ -691,7 +691,7 @@ describe("createWorkspaceDescribeFetcher", () => { test("decodes an Arrow attachment-only response into parseable columns (fetcher → normalizer → parser)", async () => { // The warehouse answers ARROW_STREAM/INLINE: rows arrive as a base64 Arrow - // IPC attachment with `data_array` undefined. Before the normalizer was + // IPC attachment with `dataArray` undefined. Before the normalizer was // wired in, parseDescribeTableExtendedJson read this as "no rows" and the // metric shipped degraded. Now the fetcher pipes the response through // normalizeResultRows, so the real describe doc is recovered end-to-end. @@ -701,12 +701,12 @@ describe("createWorkspaceDescribeFetcher", () => { executeStatement: async (req: Record) => { statements.push(req); return { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, - // Only an attachment — no data_array (the bug's trigger condition). + // Only an attachment — no dataArray (the bug's trigger condition). result: { attachment: ARROW_ATTACHMENT_B64 }, - } as DatabricksStatementExecutionResponse; + }; }, }, } as unknown as Parameters[0]; @@ -715,7 +715,7 @@ describe("createWorkspaceDescribeFetcher", () => { const response = await fetcher("appkit_demo.public.revenue_metrics"); // The fetcher decoded the attachment: rows are now readable. - expect(response.result?.data_array).toBeDefined(); + expect(response.result?.dataArray).toBeDefined(); const parsed = parseDescribeTableExtendedJson(response); const cols = extractMetricColumns(parsed); // The real revenue_metrics describe doc carries measures and dimensions. @@ -1005,7 +1005,7 @@ describe("syncMetrics", () => { test("a multi-chunk (truncated) DESCRIBE surfaces as a loud failure, not a crash (fetcher → normalizer → syncMetrics)", async () => { // End-to-end loudness check for the truncation guard. The warehouse paginates - // the DESCRIBE result (sets next_chunk_index on the first chunk); the fetcher + // the DESCRIBE result (sets nextChunkIndex on the first chunk); the fetcher // pipes the response through normalizeResultRows, which THROWS rather than // emit partial types. That throw must be caught inside describeOne and // recorded as a MetricSyncFailure — never an uncaught crash that aborts the @@ -1015,16 +1015,15 @@ describe("syncMetrics", () => { }); const client = { statementExecution: { - executeStatement: async () => - ({ - statement_id: "stmt-chunked", - status: { state: "SUCCEEDED" }, - manifest: { format: "ARROW_STREAM" }, - result: { - attachment: ARROW_ATTACHMENT_B64, - next_chunk_index: 1, - }, - }) as DatabricksStatementExecutionResponse, + executeStatement: async () => ({ + statementId: "stmt-chunked", + status: { state: "SUCCEEDED" }, + manifest: { format: "ARROW_STREAM" }, + result: { + attachment: ARROW_ATTACHMENT_B64, + nextChunkIndex: 1, + }, + }), }, } as unknown as Parameters[0]; const fetcher = createWorkspaceDescribeFetcher(client, "wh-1"); @@ -1134,24 +1133,24 @@ describe("syncMetrics — failure transience (D′)", () => { [ "a FAILED statement", { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }, ], [ "a SUCCEEDED statement with zero rows", { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }, ], [ "an unparseable payload", { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [["{not json"]] }, + result: { dataArray: [["{not json"]] }, }, ], ["zero extracted columns", mockDescribeResponse({ unrelated: true })], @@ -1198,7 +1197,7 @@ describe("syncMetrics — DESCRIBE state classification", () => { test(`a non-terminal ${state} response degrades the schema without recording a failure`, async () => { const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state }, }); @@ -1222,7 +1221,7 @@ describe("syncMetrics — DESCRIBE state classification", () => { test("a FAILED response stays a genuine failure (and its schema is degraded)", async () => { const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "FAILED", error: { message: "no such table" } }, }); @@ -1241,9 +1240,9 @@ describe("syncMetrics — DESCRIBE state classification", () => { // wrong FQN, not warehouse readiness. const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }); const { schemas, failures } = await syncMetrics( @@ -1430,7 +1429,7 @@ describe("syncMetrics — bounded-concurrency scheduling", () => { throw new Error(`boom ${key}`); } if (key === nonTerminal) { - return { statement_id: "stmt-mock", status: { state: "PENDING" } }; + return { statementId: "stmt-mock", status: { state: "PENDING" } }; } return mockDescribeResponse({ columns: [ @@ -1602,7 +1601,7 @@ describe("generateMetricTypeDeclarations — snapshot", () => { ): Promise => fqn.endsWith("cold_metric") ? // Stopped/cold warehouse: wait_timeout elapsed → non-terminal, no rows. - { statement_id: "stmt-mock", status: { state: "PENDING" } } + { statementId: "stmt-mock", status: { state: "PENDING" } } : // Genuinely measure-less view: SUCCEEDED with dimension columns only. mockDescribeResponse({ columns: [{ name: "region", type: "STRING", is_measure: false }], @@ -1767,7 +1766,7 @@ describe("metric metadata bundle", () => { // Non-terminal DESCRIBE → degraded schema (empty column arrays). const fetcher = async (): Promise => ({ - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "PENDING" }, }); const { schemas } = await syncMetrics(resolution, fetcher); diff --git a/packages/appkit/src/type-generator/tests/query-registry.test.ts b/packages/appkit/src/type-generator/tests/query-registry.test.ts index 1f8156a2f..00d78e5fb 100644 --- a/packages/appkit/src/type-generator/tests/query-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/query-registry.test.ts @@ -307,10 +307,10 @@ describe("defaultForType", () => { describe("convertToQueryType", () => { // DESCRIBE QUERY returns rows as [col_name, data_type, comment] const mockResponse: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, result: { - data_array: [ + dataArray: [ ["id", "STRING", null], ["name", "STRING", null], ["count", "INT", null], @@ -373,10 +373,10 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("uses column comment when available", () => { const responseWithComment: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, result: { - data_array: [["total", "DECIMAL", "Total amount in USD"]], + dataArray: [["total", "DECIMAL", "Total amount in USD"]], }, }; @@ -391,10 +391,10 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("quotes invalid column identifiers", () => { const responseWithInvalidName: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, result: { - data_array: [["(1 = 1)", "BOOLEAN", null]], + dataArray: [["(1 = 1)", "BOOLEAN", null]], }, }; @@ -414,9 +414,9 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("returns hasResults: false when no columns exist", () => { const emptyResponse: DatabricksStatementExecutionResponse = { - statement_id: "test-123", + statementId: "test-123", status: { state: "SUCCEEDED" }, - result: { data_array: [] }, + result: { dataArray: [] }, }; const { hasResults } = convertToQueryType( emptyResponse, @@ -439,7 +439,7 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` { col_name: "active", data_type: "BOOLEAN", comment: null }, ]); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-arrow", + statementId: "test-arrow", status: { state: "SUCCEEDED" }, result: { attachment }, }; @@ -467,7 +467,7 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` { col_name: "id", data_type: "int", comment: null }, ]); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-arrow", + statementId: "test-arrow", status: { state: "SUCCEEDED" }, result: { attachment }, }; @@ -477,15 +477,15 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` expect(type).toContain("id: number"); }); - test("prefers data_array over attachment when both are present", () => { + test("prefers dataArray over attachment when both are present", () => { const attachment = describeQueryAttachment([ { col_name: "from_arrow", data_type: "STRING", comment: null }, ]); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-both", + statementId: "test-both", status: { state: "SUCCEEDED" }, result: { - data_array: [["from_data_array", "INT", null]], + dataArray: [["from_data_array", "INT", null]], attachment, }, }; @@ -498,7 +498,7 @@ SELECT * FROM users WHERE date = :startDate AND count = :count AND name = :name` test("logs a warning and yields the unknown-result fallback on malformed attachment", () => { mockLoggerWarn.mockClear(); const response: DatabricksStatementExecutionResponse = { - statement_id: "test-bad", + statementId: "test-bad", status: { state: "SUCCEEDED" }, result: { attachment: "not-valid-arrow-ipc" }, }; diff --git a/packages/appkit/src/type-generator/tests/statement-result.test.ts b/packages/appkit/src/type-generator/tests/statement-result.test.ts index 4221cd705..1d1f43a6e 100644 --- a/packages/appkit/src/type-generator/tests/statement-result.test.ts +++ b/packages/appkit/src/type-generator/tests/statement-result.test.ts @@ -43,9 +43,9 @@ const ARROW_REORDERED_FIELDS_B64 = fs.readFileSync( ); describe("normalizeResultRows", () => { - test("decodes an Arrow attachment into data_array (real fixture)", async () => { + test("decodes an Arrow attachment into dataArray (real fixture)", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_ATTACHMENT_B64 }, @@ -54,10 +54,10 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); // One row, one cell — the JSON-string DESCRIBE payload. - expect(normalized.result?.data_array).toHaveLength(1); - expect(normalized.result?.data_array?.[0]).toHaveLength(1); + expect(normalized.result?.dataArray).toHaveLength(1); + expect(normalized.result?.dataArray?.[0]).toHaveLength(1); - const cell = normalized.result?.data_array?.[0]?.[0]; + const cell = normalized.result?.dataArray?.[0]?.[0]; expect(typeof cell).toBe("string"); // The real describe doc parses to an object with a non-empty `columns` array. @@ -66,9 +66,9 @@ describe("normalizeResultRows", () => { expect(parsed.columns.length).toBeGreaterThan(0); }); - test("preserves status, statement_id, and manifest when decoding", async () => { + test("preserves status, statementId, and manifest when decoding", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-arrow", + statementId: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_ATTACHMENT_B64 }, @@ -76,46 +76,46 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); - expect(normalized.statement_id).toBe("stmt-arrow"); + expect(normalized.statementId).toBe("stmt-arrow"); expect(normalized.status.state).toBe("SUCCEEDED"); expect(normalized.manifest?.format).toBe("ARROW_STREAM"); - // The attachment is left in place; only data_array is added. + // The attachment is left in place; only dataArray is added. expect(normalized.result?.attachment).toBe(ARROW_ATTACHMENT_B64); }); - test("passes through unchanged when data_array is already present", async () => { + test("passes through unchanged when dataArray is already present", async () => { // JSON_ARRAY warehouses (and every mocked test) take this path: no decode. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-json", + statementId: "stmt-json", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY" }, - result: { data_array: [['{"columns":[]}']] }, + result: { dataArray: [['{"columns":[]}']] }, }; const normalized = await normalizeResultRows(response); expect(normalized).toBe(response); - expect(normalized.result?.data_array).toEqual([['{"columns":[]}']]); + expect(normalized.result?.dataArray).toEqual([['{"columns":[]}']]); }); - test("treats an empty data_array as present (genuine no-rows, no decode)", async () => { + test("treats an empty dataArray as present (genuine no-rows, no decode)", async () => { // An empty array is a real "no rows" answer — it must not be overwritten by // an attachment decode even if an attachment is somehow also present. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-empty", + statementId: "stmt-empty", status: { state: "SUCCEEDED" }, - result: { data_array: [], attachment: ARROW_ATTACHMENT_B64 }, + result: { dataArray: [], attachment: ARROW_ATTACHMENT_B64 }, }; const normalized = await normalizeResultRows(response); expect(normalized).toBe(response); - expect(normalized.result?.data_array).toEqual([]); + expect(normalized.result?.dataArray).toEqual([]); }); - test("returns response unchanged when neither data_array nor attachment is present", async () => { + test("returns response unchanged when neither dataArray nor attachment is present", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-bare", + statementId: "stmt-bare", status: { state: "SUCCEEDED" }, result: {}, }; @@ -123,12 +123,12 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); expect(normalized).toBe(response); - expect(normalized.result?.data_array).toBeUndefined(); + expect(normalized.result?.dataArray).toBeUndefined(); }); test("returns response unchanged when result is entirely absent", async () => { const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-noresult", + statementId: "stmt-noresult", status: { state: "RUNNING" }, }; @@ -141,13 +141,13 @@ describe("normalizeResultRows", () => { test("does not throw when Arrow decoding rejects; degrades to no usable rows", async () => { // Bytes that look like an Arrow IPC header but aren't make `tableFromIPC` // throw. The decoder must swallow that so the generation pass does not - // crash — it leaves data_array absent and the downstream "returned no + // crash — it leaves dataArray absent and the downstream "returned no // rows" degrade fires instead. const notArrow = Buffer.from( "hello world this is plainly not an arrow ipc stream", ).toString("base64"); const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-corrupt", + statementId: "stmt-corrupt", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: notArrow }, @@ -160,19 +160,19 @@ describe("normalizeResultRows", () => { })(), ).resolves.toBeUndefined(); - // No fabricated rows: decode rejected, so data_array stays absent. - expect(normalized.result?.data_array).toBeUndefined(); + // No fabricated rows: decode rejected, so dataArray stays absent. + expect(normalized.result?.dataArray).toBeUndefined(); // The (bad) attachment is preserved; nothing was invented. expect(normalized.result?.attachment).toBe(notArrow); }); - test("decodes garbage that yields an empty Arrow table to an empty data_array", async () => { + test("decodes garbage that yields an empty Arrow table to an empty dataArray", async () => { // Some malformed payloads decode without throwing into a zero-row table - // (e.g. truncated/garbage bytes). That surfaces as an empty data_array — + // (e.g. truncated/garbage bytes). That surfaces as an empty dataArray — // which is itself a valid "no rows" answer and degrades correctly // downstream, never a fabricated row. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-garbage", + statementId: "stmt-garbage", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: "not-valid-base64-arrow-ipc!!!" }, @@ -182,58 +182,57 @@ describe("normalizeResultRows", () => { // Either absent or empty — both mean "no usable rows". Crucially: no // non-empty fabricated row. - expect(normalized.result?.data_array ?? []).toHaveLength(0); + expect(normalized.result?.dataArray ?? []).toHaveLength(0); }); - test("throws on a multi-chunk result flagged by next_chunk_index", async () => { + test("throws on a multi-chunk result flagged by nextChunkIndex", async () => { // A DESCRIBE result that exceeds INLINE's size limit is paginated. The - // first chunk carries `next_chunk_index`; decoding it alone would silently + // first chunk carries `nextChunkIndex`; decoding it alone would silently // cache partial types. The normalizer must throw (loud) rather than degrade // — distinct from the malformed-attachment path which degrades silently. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-chunked-json", + statementId: "stmt-chunked-json", status: { state: "SUCCEEDED" }, manifest: { format: "JSON_ARRAY" }, - // data_array present (first chunk) — but the guard runs ABOVE the + // dataArray present (first chunk) — but the guard runs ABOVE the // passthrough, so truncation still throws instead of returning rows. result: { - data_array: [["col_a", "STRING", null]], - next_chunk_index: 1, + dataArray: [["col_a", "STRING", null]], + nextChunkIndex: 1, }, }; await expect(normalizeResultRows(response)).rejects.toThrow(/multi-chunk/i); await expect(normalizeResultRows(response)).rejects.toThrow( - /next_chunk_index/, + /nextChunkIndex/, ); }); - test("throws on a multi-chunk result flagged by next_chunk_internal_link", async () => { + test("throws on a multi-chunk result flagged by nextChunkInternalLink", async () => { // The attachment transport can paginate too: first chunk arrives as an - // Arrow attachment with `next_chunk_internal_link` set. The guard runs + // Arrow attachment with `nextChunkInternalLink` set. The guard runs // before the decode, so this throws rather than emitting first-chunk types. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-chunked-arrow", + statementId: "stmt-chunked-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_ATTACHMENT_B64, - next_chunk_internal_link: - "/api/2.0/sql/statements/stmt/result/chunks/1", + nextChunkInternalLink: "/api/2.0/sql/statements/stmt/result/chunks/1", }, }; await expect(normalizeResultRows(response)).rejects.toThrow(/multi-chunk/i); }); - test("throws on a multi-chunk result with neither data_array nor attachment", async () => { + test("throws on a multi-chunk result with neither dataArray nor attachment", async () => { // Even when the first chunk somehow carries no inline rows, the chunk // markers alone mean the answer is truncated — refuse, do not fall through // to the "no rows" degrade. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-chunked-bare", + statementId: "stmt-chunked-bare", status: { state: "SUCCEEDED" }, - result: { next_chunk_index: 2 }, + result: { nextChunkIndex: 2 }, }; await expect(normalizeResultRows(response)).rejects.toThrow( @@ -249,7 +248,7 @@ describe("normalizeResultRows", () => { // scrambling the [col_name, data_type, comment] triple. The `[...row]` // iterator preserves field order. const response: DatabricksStatementExecutionResponse = { - statement_id: "stmt-reordered", + statementId: "stmt-reordered", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, result: { attachment: ARROW_REORDERED_FIELDS_B64 }, @@ -257,7 +256,7 @@ describe("normalizeResultRows", () => { const normalized = await normalizeResultRows(response); - expect(normalized.result?.data_array).toEqual([ + expect(normalized.result?.dataArray).toEqual([ ["revenue", "DOUBLE", "total revenue"], ]); }); @@ -271,14 +270,16 @@ describe("describeAdaptive", () => { | Promise; // Minimal WorkspaceClient stub: records the formats requested and delegates - // each executeStatement to behavior(format), which may resolve or throw. + // each executeStatement to behavior(format), which may resolve or throw. The + // fixtures are the camelCase domain type — the same shape executeStatement + // returns — so describeAdaptive consumes them directly (no adapter needed). function stubClient(behavior: StubBehavior) { const formats: string[] = []; const client = { statementExecution: { executeStatement: async (req: { format: string }) => { formats.push(req.format); - return behavior(req.format); + return await behavior(req.format); }, }, } as unknown as WorkspaceClient; @@ -288,9 +289,9 @@ describe("describeAdaptive", () => { const rows = ( data: (string | null)[][], ): DatabricksStatementExecutionResponse => ({ - statement_id: "stmt", + statementId: "stmt", status: { state: "SUCCEEDED" }, - result: { data_array: data }, + result: { dataArray: data }, }); test("standard DBSQL: JSON_ARRAY succeeds, memoized, no fallback", async () => { @@ -307,7 +308,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["schema"]]); + expect(result.result?.dataArray).toEqual([["schema"]]); expect(memo.format).toBe("JSON_ARRAY"); expect(formats).toEqual(["JSON_ARRAY"]); }); @@ -328,7 +329,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["arrow-decoded"]]); + expect(result.result?.dataArray).toEqual([["arrow-decoded"]]); expect(memo.format).toBe("ARROW_STREAM"); expect(formats).toEqual(["JSON_ARRAY", "ARROW_STREAM"]); }); @@ -338,7 +339,7 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { message: "merge_json_arrays" } }, result: {}, } as DatabricksStatementExecutionResponse; @@ -353,7 +354,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["arrow-decoded"]]); + expect(result.result?.dataArray).toEqual([["arrow-decoded"]]); expect(memo.format).toBe("ARROW_STREAM"); expect(formats).toEqual(["JSON_ARRAY", "ARROW_STREAM"]); }); @@ -375,7 +376,7 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { message: "[TABLE_OR_VIEW_NOT_FOUND]" }, @@ -410,11 +411,11 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { - error_code: "TABLE_OR_VIEW_NOT_FOUND", + errorCode: "TABLE_OR_VIEW_NOT_FOUND", message: "table x has no disposition column; format unknown", }, }, @@ -433,7 +434,7 @@ describe("describeAdaptive", () => { // The real diagnostic survives unmasked, and no second format was probed. expect(result.status.state).toBe("FAILED"); - expect(result.status.error?.error_code).toBe("TABLE_OR_VIEW_NOT_FOUND"); + expect(result.status.error?.errorCode).toBe("TABLE_OR_VIEW_NOT_FOUND"); expect(memo.format).toBeUndefined(); expect(formats).toEqual(["JSON_ARRAY"]); }); @@ -446,11 +447,11 @@ describe("describeAdaptive", () => { const { client, formats } = stubClient((format) => { if (format === "JSON_ARRAY") { return { - statement_id: "stmt", + statementId: "stmt", status: { state: "FAILED", error: { - error_code: "INVALID_PARAMETER_VALUE", + errorCode: "INVALID_PARAMETER_VALUE", message: "disposition must be one of INLINE, EXTERNAL_LINKS; format must be JSON_ARRAY, ARROW_STREAM", }, @@ -468,7 +469,7 @@ describe("describeAdaptive", () => { memo, ); - expect(result.result?.data_array).toEqual([["arrow-decoded"]]); + expect(result.result?.dataArray).toEqual([["arrow-decoded"]]); expect(memo.format).toBe("ARROW_STREAM"); expect(formats).toEqual(["JSON_ARRAY", "ARROW_STREAM"]); }); diff --git a/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts b/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts index a0a4e9995..27a2580b9 100644 --- a/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts +++ b/packages/appkit/src/type-generator/tests/sync-metric-views-types.test.ts @@ -63,9 +63,9 @@ function mockDescribeResponse( payload: unknown, ): DatabricksStatementExecutionResponse { return { - statement_id: "stmt-mock", + statementId: "stmt-mock", status: { state: "SUCCEEDED" }, - result: { data_array: [[JSON.stringify(payload)]] }, + result: { dataArray: [[JSON.stringify(payload)]] }, }; } @@ -232,7 +232,7 @@ describe("syncMetricViewsTypes", () => { mode: "blocking", suppressDegradedWrite: true, metricFetcher: async () => ({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "PENDING" }, }), }); diff --git a/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts b/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts index 6134f8348..0ffc0ed71 100644 --- a/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts +++ b/packages/appkit/src/type-generator/tests/unreachable-warehouse-gate.test.ts @@ -21,7 +21,7 @@ vi.mock("../../workspace-client", async (importOriginal) => { ...actual, createWorkspaceClient: () => ({ statementExecution: { executeStatement: mocks.executeStatement }, - warehouses: { get: mocks.getWarehouse, start: vi.fn() }, + warehouses: { getWarehouse: mocks.getWarehouse, startWarehouse: vi.fn() }, }), }; }); @@ -140,7 +140,7 @@ describe("--wait gate: environmental query failures (real query path)", () => { test("non-terminal DESCRIBE + no committed types → crashes instead of silently exiting 0", async () => { mocks.getWarehouse.mockResolvedValue({ state: "RUNNING" }); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "PENDING" }, }); @@ -162,7 +162,7 @@ describe("--wait gate: environmental query failures (real query path)", () => { test("non-terminal DESCRIBE + committed types → warns unavailable and keeps them", async () => { mocks.getWarehouse.mockResolvedValue({ state: "RUNNING" }); mocks.executeStatement.mockResolvedValue({ - statement_id: "stmt-pending", + statementId: "stmt-pending", status: { state: "RUNNING" }, }); fs.mkdirSync(path.dirname(outFile), { recursive: true }); diff --git a/packages/appkit/src/type-generator/tests/warehouse-status.test.ts b/packages/appkit/src/type-generator/tests/warehouse-status.test.ts index 882188e34..4a1b0f7b3 100644 --- a/packages/appkit/src/type-generator/tests/warehouse-status.test.ts +++ b/packages/appkit/src/type-generator/tests/warehouse-status.test.ts @@ -8,15 +8,15 @@ import { } from "../warehouse-status"; /** - * Build a minimal WorkspaceClient stub exposing only `warehouses.get`, the one - * method these helpers touch. Cast through `unknown` to the SDK type so callers - * type-check without us constructing a real client. + * Build a minimal WorkspaceClient stub exposing only `warehouses.getWarehouse`, + * the one method these helpers touch. Cast through `unknown` to the SDK type so + * callers type-check without us constructing a real client. */ function makeClient(get: ReturnType): WorkspaceClient { - return { warehouses: { get } } as unknown as WorkspaceClient; + return { warehouses: { getWarehouse: get } } as unknown as WorkspaceClient; } -/** A warehouses.get resolution carrying a given lifecycle state. */ +/** A warehouses.getWarehouse resolution carrying a given lifecycle state. */ const stateResponse = (state: WarehouseState) => ({ state }); describe("getWarehouseState", () => { diff --git a/packages/appkit/src/type-generator/types.ts b/packages/appkit/src/type-generator/types.ts index 954bde706..124a990c5 100644 --- a/packages/appkit/src/type-generator/types.ts +++ b/packages/appkit/src/type-generator/types.ts @@ -2,34 +2,39 @@ * Databricks statement execution response interface for DESCRIBE QUERY / * DESCRIBE TABLE EXTENDED. * + * A hand-written camelCase subset of the modular SDK's `StatementResponse` — + * only the fields the type generators read. `describeAdaptive` narrows the SDK + * response into this type directly (the sole difference is `dataArray`, which + * the SDK types as `JsonValue[][]`; DESCRIBE cells are always string/null). + * * Two result shapes matter here: - * - `result.data_array` — rows already materialized as JSON arrays. Present + * - `result.dataArray` — rows already materialized as JSON arrays. Present * when the warehouse returns `JSON_ARRAY` (and what every mocked test * builds). * - `result.attachment` — a base64-encoded Arrow IPC stream. Present when the * statement runs with `format: "ARROW_STREAM"` + `disposition: "INLINE"`, * which is the SDK's default disposition. The single row lands here and - * `data_array` is left undefined. {@link normalizeResultRows} decodes this - * back into `data_array` so downstream parsers stay shape-agnostic. + * `dataArray` is left undefined. {@link normalizeResultRows} decodes this + * back into `dataArray` so downstream parsers stay shape-agnostic. * - * @property statement_id - the id of the statement + * @property statementId - the id of the statement * @property status - the status of the statement * @property manifest - result metadata; `manifest.format` echoes the wire * format (`ARROW_STREAM`, `JSON_ARRAY`, ...) the warehouse chose. - * @property result - the result; either `data_array` (rows as + * @property result - the result; either `dataArray` (rows as * `[col_name, data_type, comment]` arrays) or `attachment` (base64 Arrow IPC) */ export interface DatabricksStatementExecutionResponse { - statement_id: string; + statementId: string; status: { state: string; - error?: { error_code?: string; message?: string }; + error?: { errorCode?: string; message?: string }; }; manifest?: { format?: string; }; result?: { - data_array?: (string | null)[][]; + dataArray?: (string | null)[][]; /** Base64-encoded Arrow IPC stream (ARROW_STREAM + INLINE disposition). */ attachment?: string; /** @@ -37,9 +42,9 @@ export interface DatabricksStatementExecutionResponse { * limit). Its presence means this response holds only the FIRST chunk; * {@link normalizeResultRows} throws rather than emit truncated types. */ - next_chunk_index?: number; - /** Companion to {@link next_chunk_index}: link to fetch the next chunk. */ - next_chunk_internal_link?: string; + nextChunkIndex?: number; + /** Companion to {@link nextChunkIndex}: link to fetch the next chunk. */ + nextChunkInternalLink?: string; }; } diff --git a/packages/appkit/src/type-generator/warehouse-status.ts b/packages/appkit/src/type-generator/warehouse-status.ts index 27a0afeb5..8aae70e5e 100644 --- a/packages/appkit/src/type-generator/warehouse-status.ts +++ b/packages/appkit/src/type-generator/warehouse-status.ts @@ -71,14 +71,14 @@ export async function getWarehouseState( client: WorkspaceClient, warehouseId: string, ): Promise { - const response = await client.warehouses.get({ id: warehouseId }); + const response = await client.warehouses.getWarehouse({ id: warehouseId }); return response.state as WarehouseState; } /** * Initiate a start of a stopped/stopping SQL warehouse. * - * Only KICKS OFF the start: the SDK's `start()` returns a Waiter, but we + * Only KICKS OFF the start: the SDK's `startWarehouse()` returns a Waiter, but we * deliberately do not `.wait()` on it. Blocking on the full cold-start isn't our * job here — {@link waitUntilRunning} is the poller that watches the warehouse * the rest of the way to RUNNING. We just nudge it out of the stopped state. @@ -90,7 +90,7 @@ export async function startWarehouse( client: WorkspaceClient, warehouseId: string, ): Promise { - await client.warehouses.start({ id: warehouseId }); + await client.warehouses.startWarehouse({ id: warehouseId }); } /** diff --git a/packages/appkit/src/workspace-client/index.ts b/packages/appkit/src/workspace-client/index.ts index 581cb79a8..a7d7e1c34 100644 --- a/packages/appkit/src/workspace-client/index.ts +++ b/packages/appkit/src/workspace-client/index.ts @@ -13,15 +13,8 @@ export { Time, TimeUnits, } from "shared"; -export type { - CancellationToken, - ClientOptions, - files, - GenieMessage, - jobs, - serving, - sql, - Waiter, - WorkspaceClient, - WorkspaceClientOptions, -} from "shared/workspace-client"; +// Forwards every wrapper type — legacy service namespaces (files/jobs/serving), +// the client option/waiter types, and the modular SDK client + model types +// (warehouses, statementExecution). `sql` is gone: its statement + warehouse +// types now come from the modular SDK. +export type * from "shared/workspace-client"; 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/shared/package.json b/packages/shared/package.json index a25379b08..4e127007e 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -47,7 +47,12 @@ "dependencies": { "@ast-grep/napi": "0.37.0", "@clack/prompts": "1.0.1", + "@databricks/sdk-auth": "0.46.0", + "@databricks/sdk-core": "0.46.0", "@databricks/sdk-experimental": "0.17.0", + "@databricks/sdk-options": "0.46.0", + "@databricks/sdk-statementexecution": "0.46.0", + "@databricks/sdk-warehouses": "0.46.0", "@standard-schema/spec": "1.1.0", "commander": "12.1.0", "dotenv": "16.6.1", diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index 18ef76a17..adf30c8db 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -12,11 +12,19 @@ import { type LegacyWorkspaceClient, type WorkspaceClientOptions, } from "./legacy"; +import { + buildStatementExecutionClient, + buildWarehousesClient, + type StatementExecutionClient, + type WarehousesClient, +} from "./modular"; import type { WorkspaceClient } from "./types"; export class AppKitWorkspaceClient implements WorkspaceClient { readonly #opts: WorkspaceClientOptions; #legacy?: LegacyWorkspaceClient; + #warehouses?: WarehousesClient; + #statementExecution?: StatementExecutionClient; constructor(opts: WorkspaceClientOptions) { this.#opts = opts; @@ -26,8 +34,12 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#getLegacy().files; } - get warehouses() { - return this.#getLegacy().warehouses; + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get warehouses(): WarehousesClient { + if (!this.#warehouses) { + this.#warehouses = buildWarehousesClient(this.#opts); + } + return this.#warehouses; } get genie() { @@ -38,8 +50,12 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#getLegacy().jobs; } - get statementExecution() { - return this.#getLegacy().statementExecution; + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get statementExecution(): StatementExecutionClient { + if (!this.#statementExecution) { + this.#statementExecution = buildStatementExecutionClient(this.#opts); + } + return this.#statementExecution; } get servingEndpoints() { diff --git a/packages/shared/src/workspace-client/index.ts b/packages/shared/src/workspace-client/index.ts index 91921efeb..2b981ebec 100644 --- a/packages/shared/src/workspace-client/index.ts +++ b/packages/shared/src/workspace-client/index.ts @@ -23,3 +23,5 @@ export { TimeUnits, } from "./legacy"; export type { files, jobs, serving, sql, WorkspaceClient } from "./types"; +// Modular SDK client + model types (warehouses). +export type * from "./modular"; diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts new file mode 100644 index 000000000..44a31fbec --- /dev/null +++ b/packages/shared/src/workspace-client/modular.ts @@ -0,0 +1,188 @@ +/** + * The single module allowed to import the modular `@databricks/sdk-*` SDK + * directly — the new-SDK sibling of {@link ./legacy.ts}. Every other AppKit + * module reaches these clients through the {@link WorkspaceClient} facade and + * the type re-exports below, so the modular SDK stays isolated exactly like the + * legacy one (the oxlint `no-restricted-imports` boundary walls `@databricks/sdk-*` + * off everywhere outside `packages/shared/src/workspace-client/`). + * + * Migrated services are built here as per-service clients; the facade delegates + * their accessors to these instead of the legacy monolithic client. Currently + * `warehouses` and `statementExecution` are migrated; every other service still + * routes through `legacy.ts`. + * + * NOTE: statementExecution relies on a pinned pnpm patch + * (`patches/@databricks__sdk-statementexecution@0.46.0.patch`) that restores the + * undocumented Reyden `attachment` response field, which the SDK's generated + * unmarshal transform would otherwise strip. + */ +import { + newM2mCredentials, + newPatCredentials, +} from "@databricks/sdk-auth/credentials"; +import { type HttpClient, newFetchHttpClient } from "@databricks/sdk-core/http"; +import type { ClientOptions } from "@databricks/sdk-options/client"; +import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; +import { WarehousesClient } from "@databricks/sdk-warehouses/v1"; + +import type { WorkspaceClientOptions } from "./legacy"; + +/** + * Prepend `https://` to a scheme-less host. The legacy SDK normalized the host + * this way; the modular SDK does NOT — it passes the host straight into `fetch`, + * so a bare `DATABRICKS_HOST=my-workspace.cloud.databricks.com` (the common form, + * and what the Databricks Apps runtime sets) yields `TypeError: Invalid URL`. + */ +function normalizeHost(host: string | undefined): string | undefined { + const trimmed = host?.trim(); + if (!trimmed) return undefined; + return /^https?:\/\//i.test(trimmed) ? trimmed : `https://${trimmed}`; +} + +/** + * Map wrapper options onto the modular SDK's `ClientOptions`, reproducing the + * legacy SDK's auth resolution: explicit token → PAT (the OBO path); profile → + * profile file; otherwise the service principal from the environment. It carries + * the privilege-escalation guard — check `token !== undefined` (NOT truthiness) + * so an explicitly-passed token, even an empty string, pins the PAT path and + * fails loudly at request time rather than silently falling through to the + * service-principal env credentials (which would be an OBO privilege escalation). + */ +function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { + const clientOptions: ClientOptions = {}; + // Resolve + scheme-normalize the host the way the legacy SDK did. Explicit + // `opts.host` wins; otherwise fall back to `DATABRICKS_HOST` (env is where the + // Apps runtime and dev set it). When a profile is selected without an explicit + // host, defer to the SDK's profile-file resolution instead of the env. + const host = normalizeHost( + opts.host ?? (opts.profile ? undefined : process.env.DATABRICKS_HOST), + ); + if (host) { + clientOptions.host = host; + } + if (opts.token !== undefined) { + // Explicit token (this is the OBO path: `asUser` passes the user's token). + clientOptions.credentials = newPatCredentials(opts.token); + } else if (opts.profile) { + clientOptions.profileOptions = { profile: opts.profile }; + } else { + // No token, no profile: authenticate as the service principal from the + // environment. The SDK's own default chain DOES read the DATABRICKS_* env + // vars (host, client id/secret, token) — but its M2M strategy feeds the RAW + // `DATABRICKS_HOST` straight into OAuth token-endpoint discovery, and the + // Databricks Apps runtime sets that host scheme-less (e.g. + // `x.cloud.databricks.com`), so discovery fails with `Invalid URL` and every + // request dies as "Warehouse readiness check failed". Resolve the SP here + // with the scheme-normalized `host` instead: M2M from client id + secret + // (what Apps injects), else PAT from `DATABRICKS_TOKEN`, else fall through to + // the SDK default chain (local dev, where a `~/.databrickscfg` host already + // carries a scheme). + const clientId = process.env.DATABRICKS_CLIENT_ID; + const clientSecret = process.env.DATABRICKS_CLIENT_SECRET; + const envToken = process.env.DATABRICKS_TOKEN; + if (host && clientId && clientSecret) { + clientOptions.credentials = newM2mCredentials({ + host, + clientId, + clientSecret, + }); + } else if (envToken) { + clientOptions.credentials = newPatCredentials(envToken); + } + // Otherwise leave credentials unset and let the SDK walk its profile-based + // default chain (local dev with `~/.databrickscfg`). + } + const httpClient = buildHttpClient(opts); + if (httpClient) { + clientOptions.httpClient = httpClient; + } + return clientOptions; +} + +/** + * Wrap the SDK's default fetch transport to prepend AppKit's product segment to + * the outgoing `User-Agent`, preserving the exact legacy string (e.g. + * `@databricks/appkit/0.75.1`) that Databricks-side dashboards match on. + * + * Why the transport and not `setProduct`: the modular SDK's client-info API + * validates the product as a simple token and rejects `@databricks/appkit` (the + * `@`/`/`), and it is process-global. Setting the header on the `httpClient` + * instead keeps the literal product name, is per-client, and — unlike a pnpm + * patch — ships inside appkit's bundled `dist`, so it also reaches deployed apps + * (npm-installed from the tarball, where pnpm patches do not apply). The SDK's + * own client-info (`sdk-js-core/…`, runtime) is already on `request.headers`, so + * prepending keeps it intact after AppKit's segment. + * + * Returns `undefined` when no product is configured (build-time callers), leaving + * the SDK's default User-Agent untouched — matching the legacy behavior where + * build-time clients carried no AppKit UA. + */ +function buildHttpClient(opts: WorkspaceClientOptions): HttpClient | undefined { + const co = opts.clientOptions; + if (!co?.product || !co?.productVersion) { + return undefined; + } + const segments = [`${co.product}/${co.productVersion}`]; + if (co.userAgentExtra) { + for (const [key, value] of Object.entries(co.userAgentExtra)) { + segments.push(`${key}/${String(value)}`); + } + } + const appkitUserAgent = segments.join(" "); + const base = newFetchHttpClient(); + return { + send(request) { + const existing = request.headers.get("User-Agent"); + request.headers.set( + "User-Agent", + existing ? `${appkitUserAgent} ${existing}` : appkitUserAgent, + ); + return base.send(request); + }, + }; +} + +/** Build a modular Warehouses client from wrapper options. */ +export function buildWarehousesClient( + opts: WorkspaceClientOptions, +): WarehousesClient { + return new WarehousesClient(mapToClientOptions(opts)); +} + +/** Build a modular Statement Execution client from wrapper options. */ +export function buildStatementExecutionClient( + opts: WorkspaceClientOptions, +): StatementExecutionClient { + return new StatementExecutionClient(mapToClientOptions(opts)); +} + +// ── Client type re-exports (for the facade accessor types) ─────────────── +export type { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; +export type { WarehousesClient } from "@databricks/sdk-warehouses/v1"; + +// ── Model type re-exports ──────────────────────────────────────────────── +// AppKit modules import request/response/enum types from the wrapper rather +// than the SDK, so the import boundary holds. Type-only: the connector compares +// state against string literals, which satisfy the SDK's `Enum | (string & {})` +// field unions — no runtime enum values needed. +export type { + ColumnInfo, + Disposition, + ExecuteStatementRequest, + ExternalLink, + Format, + ResultData, + ResultManifest, + Schema, + ServiceError, + StatementParameter, + StatementResponse, + StatementStatus, + StatementStatus_State, +} from "@databricks/sdk-statementexecution/v1"; +export type { + EndpointHealth, + EndpointInfo, + EndpointState, + GetWarehouseResponse, +} from "@databricks/sdk-warehouses/v1"; diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts new file mode 100644 index 000000000..51814a293 --- /dev/null +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -0,0 +1,226 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +// The wrapper's own tests are the one place allowed to mock the SDK directly. +// Capture the `ClientOptions` the modular `WarehousesClient` constructor receives +// so we can assert how wrapper options map onto the modular SDK's config. +const { ctorOpts, patTokens, m2mOpts } = vi.hoisted(() => ({ + ctorOpts: [] as Array>, + patTokens: [] as string[], + m2mOpts: [] as Array>, +})); + +vi.mock("@databricks/sdk-warehouses/v1", () => ({ + WarehousesClient: vi.fn().mockImplementation((opts) => { + ctorOpts.push(opts); + return { opts }; + }), +})); +vi.mock("@databricks/sdk-statementexecution/v1", () => ({ + StatementExecutionClient: vi.fn().mockImplementation((opts) => ({ opts })), +})); +vi.mock("@databricks/sdk-auth/credentials", () => ({ + newPatCredentials: vi.fn((token: string) => { + patTokens.push(token); + return { kind: "pat", token }; + }), + newM2mCredentials: vi.fn((opts: Record) => { + m2mOpts.push(opts); + return { kind: "m2m", ...opts }; + }), +})); +// The default transport: its `send` echoes the final request headers so tests +// can assert the User-Agent the wrapper set before delegating. +vi.mock("@databricks/sdk-core/http", () => ({ + newFetchHttpClient: vi.fn(() => ({ + send: vi.fn((request: { headers: Headers }) => + Promise.resolve({ + statusCode: 200, + headers: request.headers, + body: null, + }), + ), + })), +})); + +import { buildWarehousesClient } from "../modular"; + +/** Drive the wrapped httpClient with one request and return the UA it set. */ +async function sentUserAgent( + httpClient: unknown, + seedUserAgent?: string, +): Promise { + const headers = new Headers( + seedUserAgent ? { "User-Agent": seedUserAgent } : undefined, + ); + await ( + httpClient as { + send: (r: { + url: string; + method: string; + headers: Headers; + }) => Promise; + } + ).send({ url: "https://x", method: "GET", headers }); + return headers.get("User-Agent"); +} + +describe("modular mapToClientOptions (via buildWarehousesClient)", () => { + // Auth resolution reads these env vars; snapshot + clear them so the dev + // machine's own DATABRICKS_* values never leak into a case. + const AUTH_ENV = [ + "DATABRICKS_HOST", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + "DATABRICKS_TOKEN", + ] as const; + const originalEnv: Record = {}; + + beforeEach(() => { + ctorOpts.length = 0; + patTokens.length = 0; + m2mOpts.length = 0; + for (const key of AUTH_ENV) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } + }); + + afterEach(() => { + for (const key of AUTH_ENV) { + if (originalEnv[key] === undefined) delete process.env[key]; + else process.env[key] = originalEnv[key]; + } + }); + + test("prepends https:// to a scheme-less explicit host", () => { + buildWarehousesClient({ host: "ws.cloud.databricks.com" }); + expect(ctorOpts[0].host).toBe("https://ws.cloud.databricks.com"); + }); + + test("leaves an explicit host that already has a scheme unchanged", () => { + buildWarehousesClient({ host: "https://ws.cloud.databricks.com" }); + expect(ctorOpts[0].host).toBe("https://ws.cloud.databricks.com"); + }); + + test("falls back to DATABRICKS_HOST (scheme-normalized) when no host is passed", () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + buildWarehousesClient({}); + expect(ctorOpts[0].host).toBe("https://envhost.cloud.databricks.com"); + }); + + test("a token takes the PAT path and pins the resolved host", () => { + buildWarehousesClient({ token: "abc", host: "https://x" }); + expect(patTokens).toEqual(["abc"]); + expect(ctorOpts[0].host).toBe("https://x"); + expect(ctorOpts[0].credentials).toEqual({ kind: "pat", token: "abc" }); + }); + + test("an empty-string token still uses PAT (no silent fall-through to default auth)", () => { + buildWarehousesClient({ token: "", host: "https://x" }); + expect(patTokens).toEqual([""]); + expect(ctorOpts[0].credentials).toEqual({ kind: "pat", token: "" }); + }); + + test("a profile sets profileOptions and defers host to the SDK (ignores env)", () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + buildWarehousesClient({ profile: "myprofile" }); + expect(ctorOpts[0].profileOptions).toEqual({ profile: "myprofile" }); + expect(ctorOpts[0].host).toBeUndefined(); + expect(patTokens).toEqual([]); + }); + + test("no host, no token, no profile, no env → empty options (SDK default chain)", () => { + buildWarehousesClient({}); + expect(ctorOpts[0].host).toBeUndefined(); + expect(ctorOpts[0].credentials).toBeUndefined(); + expect(ctorOpts[0].profileOptions).toBeUndefined(); + }); + + test("service-principal by default: DATABRICKS_CLIENT_ID/SECRET + host env → M2M creds", () => { + // The Databricks Apps runtime injects the app's SP credentials via env. The + // SDK's default chain would read them too, but its M2M strategy uses the raw + // scheme-less DATABRICKS_HOST for OAuth discovery (→ Invalid URL); we resolve + // M2M here with the scheme-normalized host so discovery succeeds. + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({}); + expect(m2mOpts).toEqual([ + { + host: "https://envhost.cloud.databricks.com", + clientId: "sp-client-id", + clientSecret: "sp-secret", + }, + ]); + expect(ctorOpts[0].credentials).toEqual({ + kind: "m2m", + host: "https://envhost.cloud.databricks.com", + clientId: "sp-client-id", + clientSecret: "sp-secret", + }); + expect(patTokens).toEqual([]); + }); + + test("falls back to DATABRICKS_TOKEN (PAT) when no client id/secret is set", () => { + process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; + process.env.DATABRICKS_TOKEN = "env-pat"; + buildWarehousesClient({}); + expect(patTokens).toEqual(["env-pat"]); + expect(ctorOpts[0].credentials).toEqual({ kind: "pat", token: "env-pat" }); + expect(m2mOpts).toEqual([]); + }); + + test("an explicit (OBO) token wins over env SP credentials — no escalation", () => { + // asUser passes the user's token; it must NOT be shadowed by the SP env + // creds the deployed runtime also sets. + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({ token: "user-token", host: "https://x" }); + expect(patTokens).toEqual(["user-token"]); + expect(ctorOpts[0].credentials).toEqual({ + kind: "pat", + token: "user-token", + }); + expect(m2mOpts).toEqual([]); + }); + + test("M2M needs a host: client id/secret with no resolvable host falls through to the default chain", () => { + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + buildWarehousesClient({}); + expect(m2mOpts).toEqual([]); + expect(ctorOpts[0].credentials).toBeUndefined(); + }); + + test("User-Agent: prepends the exact @databricks/appkit product segment (dashboards match on it)", async () => { + // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` + // (the `@`/`/`). We set the UA on the httpClient transport instead, keeping + // the literal legacy product string that Databricks-side dashboards match. + buildWarehousesClient({ + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + userAgentExtra: { mode: "dev" }, + }, + } as never); + // The SDK's own client-info UA is already on the request; ours prepends. + const ua = await sentUserAgent(ctorOpts[0].httpClient, "sdk-js-core/1.0.0"); + expect(ua).toBe("@databricks/appkit/0.64.0 mode/dev sdk-js-core/1.0.0"); + }); + + test("User-Agent: sets the product segment even when the request has no prior UA", async () => { + buildWarehousesClient({ + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + }, + } as never); + const ua = await sentUserAgent(ctorOpts[0].httpClient); + expect(ua).toBe("@databricks/appkit/0.64.0"); + }); + + test("no product configured (build-time) → no httpClient override (SDK default UA)", () => { + buildWarehousesClient({ host: "https://x" }); + expect(ctorOpts[0].httpClient).toBeUndefined(); + }); +}); diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index 398a4afe0..6d9865ace 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -14,10 +14,17 @@ * as each service migrates. */ import type { LegacyWorkspaceClient } from "./legacy"; +import type { StatementExecutionClient, WarehousesClient } from "./modular"; -// SDK type namespaces, re-exported so AppKit modules import them from the -// wrapper rather than the SDK directly. +// Legacy SDK type namespaces for un-migrated services, re-exported so AppKit +// modules import them from the wrapper rather than the SDK directly. `sql` +// stays only for the dev-mode warehouse listing in service-context, which reads +// the raw (snake_case) `/api/2.0/sql/warehouses` body via the still-legacy +// `apiClient` and types it as `sql.EndpointInfo[]`. Statement + warehouse +// service types now come from `./modular`. export type { files, jobs, serving, sql } from "@databricks/sdk-experimental"; +// Modular SDK client + model types (warehouses, statementExecution). +export type * from "./modular"; /** * AppKit's workspace client facade. Mirrors the multi-client shape of the @@ -31,8 +38,8 @@ export interface WorkspaceClient { /** UC Volumes / Files API. */ readonly files: LegacyWorkspaceClient["files"]; - /** SQL Warehouses. */ - readonly warehouses: LegacyWorkspaceClient["warehouses"]; + /** SQL Warehouses (modular SDK). */ + readonly warehouses: WarehousesClient; /** Genie / dashboards. */ readonly genie: LegacyWorkspaceClient["genie"]; @@ -40,8 +47,8 @@ export interface WorkspaceClient { /** Jobs. */ readonly jobs: LegacyWorkspaceClient["jobs"]; - /** Statement Execution. */ - readonly statementExecution: LegacyWorkspaceClient["statementExecution"]; + /** Statement Execution (modular SDK). */ + readonly statementExecution: StatementExecutionClient; /** Serving Endpoints. */ readonly servingEndpoints: LegacyWorkspaceClient["servingEndpoints"]; diff --git a/patches/@databricks__sdk-statementexecution@0.46.0.patch b/patches/@databricks__sdk-statementexecution@0.46.0.patch new file mode 100644 index 000000000..c206b65c4 --- /dev/null +++ b/patches/@databricks__sdk-statementexecution@0.46.0.patch @@ -0,0 +1,34 @@ +diff --git a/dist/v1/model.d.ts b/dist/v1/model.d.ts +index e8d95659ea348b384a3d32b6a3d4f754287b38b6..705b9bed203981a2f3cde5417ed8019ff3a7065c 100644 +--- a/dist/v1/model.d.ts ++++ b/dist/v1/model.d.ts +@@ -385,6 +385,8 @@ interface QueryTag { + * link is returned.) + */ + interface ResultData { ++ /** PATCH(appkit): Reyden's non-standard INLINE ARROW_STREAM payload (base64 Arrow IPC). */ ++ attachment?: string | undefined; + externalLinks?: ExternalLink[] | undefined; + /** + * The `JSON_ARRAY` format is an array of arrays of values, where each non-null value is +diff --git a/dist/v1/model.js b/dist/v1/model.js +index fc35e28bbad5e7e873c14f6492696f9086ec9280..3bd78f3ad2dea39cdc883fa0daddb940e99e83b9 100644 +--- a/dist/v1/model.js ++++ b/dist/v1/model.js +@@ -177,10 +177,15 @@ const unmarshalResultDataSchema = z.object({ + z.string() + ]).transform((v) => BigInt(v)).optional(), + next_chunk_index: z.number().optional(), +- next_chunk_internal_link: z.string().optional() ++ next_chunk_internal_link: z.string().optional(), ++ // PATCH(appkit): preserve Reyden's non-standard INLINE ARROW_STREAM `attachment` ++ // (base64 Arrow IPC). The generated schema + rebuild-transform would otherwise ++ // strip it, breaking the inline-arrow delivery path. See patches/ for rationale. ++ attachment: z.string().optional() + }).transform((d) => ({ + externalLinks: d.external_links, + dataArray: d.data_array, ++ attachment: d.attachment, + chunkIndex: d.chunk_index, + rowOffset: d.row_offset, + rowCount: d.row_count, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 8c6660f4d..f14b66053 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -11,6 +11,11 @@ overrides: qs@<6.15.2: 6.15.2 size-sensor: 1.0.3 +patchedDependencies: + '@databricks/sdk-statementexecution@0.46.0': + hash: a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d + path: patches/@databricks__sdk-statementexecution@0.46.0.patch + importers: .: @@ -258,9 +263,24 @@ importers: '@databricks/lakebase': specifier: workspace:* version: link:../lakebase + '@databricks/sdk-auth': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-core': + specifier: 0.46.0 + version: 0.46.0 '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-options': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-statementexecution': + specifier: 0.46.0 + version: 0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d) + '@databricks/sdk-warehouses': + specifier: 0.46.0 + version: 0.46.0 '@mlflow/core': specifier: 0.4.0 version: 0.4.0(bufferutil@4.0.9) @@ -573,9 +593,24 @@ importers: '@clack/prompts': specifier: 1.0.1 version: 1.0.1 + '@databricks/sdk-auth': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-core': + specifier: 0.46.0 + version: 0.46.0 '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-options': + specifier: 0.46.0 + version: 0.46.0 + '@databricks/sdk-statementexecution': + specifier: 0.46.0 + version: 0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d) + '@databricks/sdk-warehouses': + specifier: 0.46.0 + version: 0.46.0 '@standard-schema/spec': specifier: 1.1.0 version: 1.1.0 @@ -1935,6 +1970,14 @@ packages: engines: {node: ^20 || ^22 || ^24 || ^25, pnpm: '>=10'} hasBin: true + '@databricks/sdk-auth@0.46.0': + resolution: {integrity: sha512-cMrwxsFtpiEKFxta5dKHchKdrgmHkQ6upJ2C4OacmlHrOHJ+ChzQBVTpAiVtjX62xj+YjNgp/29IpSdKKYUVDA==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-core@0.46.0': + resolution: {integrity: sha512-Q2LAGWYIi+jyeKR9OIqvkgyde2GdzqfSG8lewxA9Xu/C9RJBBFbSfg5Nh8ZC66TKElGIosVOecoEJdbxnNMuxw==} + engines: {node: '>=22.0.0'} + '@databricks/sdk-experimental@0.15.0': resolution: {integrity: sha512-HkoMiF7dNDt6WRW0xhi7oPlBJQfxJ9suJhEZRFt08VwLMaWcw2PiF8monfHlkD4lkufEYV6CTxi5njQkciqiHA==} engines: {node: '>=22.0', npm: '>=10.0.0'} @@ -1943,6 +1986,18 @@ packages: resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} + '@databricks/sdk-options@0.46.0': + resolution: {integrity: sha512-UtADlR+41rYEoOCycZvJh1g96uDN6GVWgqQk+72cHBzcxi+koxKSJXHcYRI997Sc4Fcc1d2oyAC2I2ddVhurjA==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-statementexecution@0.46.0': + resolution: {integrity: sha512-VJA3e7UHmxRxN42/mV5VtKeINME0vCz3Na3hrwmta3tZqJWZbBU7XTfUdD1yOQ5Z1JU5UIM65OlWq8gc4IzHFg==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-warehouses@0.46.0': + resolution: {integrity: sha512-9r/gbdTb6ASiWCiibCwAOF8QizqNacidIw78uwzJYKK9dbdqmWIfNK0pF/jt3BG1sUiXz2b1I6URPdX7Qi0oLg==} + engines: {node: '>=22.0.0'} + '@date-fns/tz@1.4.1': resolution: {integrity: sha512-P5LUNhtbj6YfI3iJjw5EL9eUAG6OitD0W3fWQcpQjDRc/QIsL0tRNuO1PcDvPccWL1fSTXXdE1ds+l95DV/OFA==} @@ -2609,6 +2664,10 @@ packages: '@js-sdsl/ordered-map@4.4.2': resolution: {integrity: sha512-iUKgm52T8HOE/makSxjqoWhe95ZJA1/G1sYsGev2JDKUSS14KAgg1LHb+Ba+IPow0xflbnSkOsZcO08C7w1gYw==} + '@js-temporal/polyfill@0.5.1': + resolution: {integrity: sha512-hloP58zRVCRSpgDxmqCWJNlizAlUgJFqG2ypq79DCvyv9tHjRYMDOcPFjzfl/A1/YxDvRCZz8wvZvmapQnKwFQ==} + engines: {node: '>=12'} + '@jsep-plugin/assignment@1.3.0': resolution: {integrity: sha512-VVgV+CXrhbMI3aSusQyclHkenWSAm95WaiKrMxRFam3JSUiIaQjoMIw2sEs/OX4XifnqeQUN4DYbJjlA8EfktQ==} engines: {node: '>= 10.16.0'} @@ -8668,6 +8727,9 @@ packages: resolution: {integrity: sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==} hasBin: true + jsbi@4.3.2: + resolution: {integrity: sha512-9fqMSQbhJykSeii05nxKl4m6Eqn2P6rOlYiS+C5Dr/HPIU/7yZxu5qzbs40tgaFORiw2Amd0mirjxatXYMkIew==} + jsdom@27.0.0: resolution: {integrity: sha512-lIHeR1qlIRrIN5VMccd8tI2Sgw6ieYXSVktcSHaNe3Z5nE/tcPQYQWOq00wxMvYOsz+73eAkNenVvmPC6bba9A==} engines: {node: '>=20'} @@ -14215,6 +14277,16 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-auth@0.46.0': + dependencies: + '@databricks/sdk-core': 0.46.0 + zod: 4.3.6 + + '@databricks/sdk-core@0.46.0': + dependencies: + json-bigint: 1.0.0 + zod: 4.3.6 + '@databricks/sdk-experimental@0.15.0': dependencies: google-auth-library: 10.5.0 @@ -14233,6 +14305,29 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-options@0.46.0': + dependencies: + '@databricks/sdk-auth': 0.46.0 + '@databricks/sdk-core': 0.46.0 + + '@databricks/sdk-statementexecution@0.46.0(patch_hash=a0fde44d73faf28cc107a930fea00967d00dd4e74796002c77686bb5bf56569d)': + dependencies: + '@databricks/sdk-auth': 0.46.0 + '@databricks/sdk-core': 0.46.0 + '@databricks/sdk-options': 0.46.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + + '@databricks/sdk-warehouses@0.46.0': + dependencies: + '@databricks/sdk-auth': 0.46.0 + '@databricks/sdk-core': 0.46.0 + '@databricks/sdk-options': 0.46.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + '@date-fns/tz@1.4.1': {} '@discoveryjs/json-ext@0.5.7': {} @@ -15431,6 +15526,10 @@ snapshots: '@js-sdsl/ordered-map@4.4.2': {} + '@js-temporal/polyfill@0.5.1': + dependencies: + jsbi: 4.3.2 + '@jsep-plugin/assignment@1.3.0(jsep@1.4.0)': dependencies: jsep: 1.4.0 @@ -22044,6 +22143,8 @@ snapshots: dependencies: argparse: 2.0.1 + jsbi@4.3.2: {} + jsdom@27.0.0(bufferutil@4.0.9)(postcss@8.5.6): dependencies: '@asamuzakjp/dom-selector': 6.6.2 diff --git a/template/server/example.test.ts b/template/server/example.test.ts index 9140c2e24..535f800fd 100644 --- a/template/server/example.test.ts +++ b/template/server/example.test.ts @@ -1,5 +1,5 @@ -import { Plugin, type PluginManifest } from '@databricks/appkit'; -import { expectStream, createTestPluginContext } from '@databricks/appkit/testing'; +import { Plugin, type PluginManifest, toPlugin } from '@databricks/appkit'; +import { createTestApp, createTestPluginContext, expectStream } from '@databricks/appkit/testing'; import { describe, expect, test } from 'vitest'; /** @@ -9,10 +9,13 @@ import { describe, expect, test } from 'vitest'; * network — so these tests run anywhere, including CI. Delete this file, or use * it as a starting point for testing your own plugins. * - * Two headline helpers are shown below: + * 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. + * run under test. No boot, no socket — the fastest option for unit tests. * - `expectStream(...).toEmit(...)` — assert the ordered event types a * streaming handler emits. * @@ -43,8 +46,25 @@ class GreeterPlugin extends Plugin { 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(); @@ -61,4 +81,20 @@ describe('testing kit example', () => { 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/dist-appkit.ts b/tools/dist-appkit.ts index fee75f96a..83bc975a0 100644 --- a/tools/dist-appkit.ts +++ b/tools/dist-appkit.ts @@ -23,6 +23,24 @@ const pkg = JSON.parse(fs.readFileSync("package.json", "utf-8")); // "shared" is intentionally excluded: it is bundled directly into appkit/appkit-ui via noExternal. const WORKSPACE_PACKAGE_REPLACEMENTS = ["@databricks/lakebase"]; +// Modular SDK packages we carry a pnpm patch for (`patches/`). A pnpm patch is +// applied only at *this* monorepo's install and does NOT travel through a normal +// `npm install`, so a deployed consumer would otherwise resolve the unpatched +// registry copy. We ship the patched copy *inside* appkit's own tarball via +// npm `bundledDependencies`: it lands in `/node_modules/@databricks/ +// appkit/node_modules/`, which Node's nested resolution prefers for appkit's +// own imports — on any package manager, including the Databricks Apps runtime's +// npm. Extensible: when a future patch is added, list its package here. +// (Their transitive deps — sdk-core/sdk-auth/zod — stay external and are already +// declared in appkit's dependencies, so they resolve from the consumer.) +const BUNDLED_PATCHED_PACKAGES = ["@databricks/sdk-statementexecution"]; + +// This script builds BOTH the appkit and appkit-ui tarballs, so only bundle a +// patched package into the tarball whose package actually declares it as a +// dependency (captured before the CLI-dependency merge below). appkit-ui does +// not depend on the modular SDK packages, so it bundles nothing. +const ownDependencyNames = new Set(Object.keys(pkg.dependencies ?? {})); + if (prerelease) { pkg.version = `${pkg.version}-pr.${prerelease}`; } @@ -64,6 +82,34 @@ if (fs.existsSync(sharedPostinstall)) { pkg.dependencies = pkg.dependencies || {}; Object.assign(pkg.dependencies, CLI_DEPENDENCIES); +// Ship pnpm-patched SDK packages inside the tarball via `bundledDependencies` +// (see BUNDLED_PATCHED_PACKAGES above). Copy the patched copy from this +// monorepo's node_modules (pnpm applies the patch there) into the tarball's +// node_modules and list it as bundled so npm packs it and the consumer resolves +// appkit's imports to it. `dereference: true` resolves pnpm's symlink to the +// real (patched) files. +const bundled: string[] = []; +for (const depName of BUNDLED_PATCHED_PACKAGES) { + // Skip packages this tarball's package doesn't depend on (e.g. appkit-ui). + if (!ownDependencyNames.has(depName)) { + continue; + } + const src = path.resolve("node_modules", depName); + if (!fs.existsSync(src)) { + throw new Error( + `bundledDependencies: patched package not found at ${src} — is it installed + patched?`, + ); + } + const dest = path.join("tmp/node_modules", depName); + fs.rmSync(dest, { recursive: true, force: true }); + fs.mkdirSync(path.dirname(dest), { recursive: true }); + fs.cpSync(src, dest, { recursive: true, dereference: true }); + bundled.push(depName); +} +if (bundled.length > 0) { + pkg.bundledDependencies = bundled; +} + fs.writeFileSync("tmp/package.json", JSON.stringify(pkg, null, 2)); fs.cpSync("dist", "tmp/dist", { recursive: true }); diff --git a/tools/test-helpers.ts b/tools/test-helpers.ts index 5a88312da..b09c58524 100644 --- a/tools/test-helpers.ts +++ b/tools/test-helpers.ts @@ -5,13 +5,16 @@ * `@tools/test-helpers` importers keep working; new code (inside or outside * this repo) should import from `@databricks/appkit/testing` instead. * + * 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,