Skip to main content

Testing

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. Plugin tests stay fast and run in CI, where no workspace is available.

Goal

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 three entry points plus a set of fixture helpers:

  • createTestApp({ plugins }) — boot a real app and call it over real HTTP. Start here.
  • createTestPluginContext() — build a real PluginContext with faked edges and attach it to a plugin, with no boot and no socket.
  • expectStream(...).toEmit(...) — assert the ordered event types a stream emits.
  • FixturescreateMockRequest, createMockResponse, createMockWorkspaceClient, mockServiceContext, and SQL response builders.

The kit uses Vitest's vi for its mocks, so vitest is an optional peer dependency: you already have it (you're writing Vitest tests), and the kit resolves to your copy rather than bundling a second one. Because it's optional, it is not installed into apps that never import @databricks/appkit/testing — production installs stay free of the test framework. Any Vitest v3 or v4 works.

Testing your plugin

createTestApp({ plugins }) boots a real AppKit app, with the real Express wiring, routes, and resource validation, then hands you methods to call it like a client would:

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?

createTestAppcreateTestPluginContext
Boots the appYesNo
Binds a socketYes (ephemeral port)No
Express middleware, error handlerRealNot involved
Resource / env validationReal, and strictNot involved
Workspace clientFaked and injectedFake it yourself with mockServiceContext
Needs close()YesNo
SpeedFast, but pays for a socketFastest

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 — "<service>.<method>" on AppKit's workspace-client facade:

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 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:

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.
  • obotrue 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:

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:

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:

// 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().

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:

EdgeHow it's faked
TelemetryA no-op mock provider — no OpenTelemetry pipeline needed.
Tool providersFakes registered through the real registerToolProvider, keyed by plugin then tool name.
RoutesThe real addRoute/addMiddleware are wrapped to record what a plugin registers.

Because the context is real, executeTool still resolves the user scope via asUser(req) and still composes the abort signal from your timeout — so those paths are genuinely under test.

Registering fake tool responses

Pass canned responses keyed by plugin name, then tool name. A response is either a static value or a function of the call arguments and the composed abort signal:

import { createTestPluginContext } from "@databricks/appkit/testing";

const mock = createTestPluginContext({
  analytics: {
    // static response
    top_users: [{ user: "alice", events: 42 }],
    // function response — assert on args, or simulate slow/aborting work
    query: (args, signal) => runFakeQuery(args, signal),
  },
});

Attaching to a plugin

attach() wires the context to a plugin the production way: it seeds an in-memory cache (if AppKit hasn't already initialized one), then calls the plugin's attachContext, which rebuilds telemetry and flips isReady to true. Await it before exercising any handler that reads this.context, this.cache, or gates on isReady:

const plugin = new MyAgentPlugin({});
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 workspace client and the on-behalf-of stub are process-wide too, not per app: ServiceContext holds one client, and the createUserContext fake is a single spy. Because of that, createTestApp allows one open app at a time and throws if you boot a second before closing the first — with two open, the second one's client and responses would not reach the handlers, and closing either would remove the shared OBO fake from the other. Vitest isolates test files in separate workers, so this only constrains apps within a single file. One consequence worth knowing: a describe that holds an app open in beforeAll cannot contain a test that boots its own.

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():

import { resetTestCache } from "@databricks/appkit/testing";

beforeEach(async () => {
  await resetTestCache(); // no-op if the cache isn't initialized yet
});

It also helps within a single test — clear the cache to force a miss, then assert the following call is a hit.

Inspecting what happened

The returned object exposes live views you read after the action under test runs:

await someHandler(req, res);

// Every cross-plugin tool dispatch, in order.
expect(mock.toolCalls[0]).toMatchObject({
  plugin: "analytics",
  tool: "query",
  asUser: true, // proves the on-behalf-of path ran
});

// Every route the plugin registered (raw handlers, before wrapping).
expect(mock.routes).toContainEqual(
  expect.objectContaining({ method: "post", path: "/invocations" }),
);

// The injected telemetry provider records the context's own spans — i.e. the
// span PluginContext.executeTool opens around each cross-plugin tool call.
expect(mock.telemetry.getTracer().startActiveSpan).toHaveBeenCalled();

mock.telemetry is injected into the PluginContext, so it captures the spans the context opens (notably executeTool). It is not the plugin's own telemetry: attachContext rebuilds this.telemetry from the real TelemetryManager, so spans a plugin opens internally do not land on mock.telemetry.

RecordedToolCall.asUser is the field to assert for cross-plugin calls: because the fake asUser enforces the same token precondition as the real Plugin.asUser, a dispatch that records asUser: true (with userId set) genuinely resolved the caller's user scope, and a request missing x-forwarded-access-token rejects instead — the OBO distinction that silent { executeTool } stubs cannot verify. Assert both directions: a well-formed request records the expected userId, and a token-less one throws.

The fake replicates asUser's token precondition, not its internal dev-mode telemetry marker: in NODE_ENV=development the real Plugin.asUser skips impersonation and sets an OTel isDevOboFallback() flag, which the fake does not reproduce. Assert OBO through the recorded asUser/userId fields rather than isDevOboFallback().

expectStream(...)

AppKit plugins stream Server-Sent Events. expectStream consumes a stream and asserts the ordered event types it emits. It accepts an async iterable (an agent adapter's run()), a plain array of events, an SSE Response (or a promise of one) whose body it parses, or a createMockResponse() whose captured writes it replays.

import { expectStream } from "@databricks/appkit/testing";

// In-order subsequence match — interleaved events (heartbeats, deltas) are ignored.
await expectStream(agent.adapter.run(input)).toEmit("tool_call", "message_delta");

// Exact match — the stream's full shape, in order, with nothing else.
await expectStream(events).toEmitExactly("warehouse_status", "result");

// Or collect without asserting.
const types = await expectStream(res).collectTypes();

Asserting a plugin's streaming route

Most plugins stream SSE from a route handler (res.write(...)), not a bare generator. createMockResponse() captures those writes, and expectStream reads them straight back — drive the real handler, then assert:

import { createMockRequest, createMockResponse, expectStream } from "@databricks/appkit/testing";

const res = createMockResponse();
await plugin._handleStream(createMockRequest({ obo: true }), res);

// The mock captured the SSE the handler wrote; expectStream parses it.
await expectStream(res).toEmit("status", "result");

expectStream(res) and expectStream(res.sseResponse()) are equivalent; the latter hands you the raw Response if you want it. Do not pass the SSE body as a string: a string is an iterable of characters, so expectStream rejects it with a pointer to sseResponse() rather than emitting one "event" per character.

toEmit checks that the expected types appear in order but tolerates other events before, between, or after them — which is what you want for streams that interleave bookkeeping events like heartbeats or metadata. Use toEmitExactly when the stream's shape is fully determined.

expectStream buffers the whole source before asserting, so a stream that never terminates would otherwise hang until the test runner's own timeout. Pass { timeout } to fail fast with a clear error instead:

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 now covers both. createTestApp fakes the data plane for you by injecting a mock workspace client at the real seam; below that, mockServiceContext spies the singleton directly, and createMockWorkspaceClient builds the client either of them installs.

The kit re-exports the request/response/context fixtures AppKit uses internally:

  • createMockRequest(overrides?) / createMockResponse() — Express request/response doubles, including the streaming flags (headersSent, writableEnded). Pass obo: true (or obo: { userId, token, email }) to set the forwarded identity headers asUser requires, instead of hand-adding them. createMockResponse() also captures everything a handler writes; pass it to expectStream (or call sseResponse()) to assert a streaming route's SSE. (Plugins resolve the workspace client through getWorkspaceClient(), not the request — use mockServiceContext to control it.)

  • mockServiceContext(options?) — spy the ServiceContext singleton so code that resolves the service principal or a user context gets test doubles. Call in beforeEach, and call the returned restore() in afterEach.

  • useServiceContextMock(options?) — the same, in one line: it registers the beforeEach install and afterEach restore for you. Call it at the top of a describe block (not inside a test), and read the live .current handle from within a test:

    describe("my plugin", () => {
      const ctx = useServiceContextMock();
      test("...", async () => {
        await handler(createMockRequest({ obo: true }), res);
        expect(ctx.current.createUserContextSpy).toHaveBeenCalled();
      });
    });
  • createSuccessfulSQLResponse(rows, columns) / createFailedSQLResponse(message) — build SQL Warehouse statement responses.

  • setupDatabricksEnv(overrides?) — set DATABRICKS_HOST / DATABRICKS_WAREHOUSE_ID to test values.

  • resetTestCache() — clear the shared cache singleton between (or within) tests; no-ops if the cache isn't initialized yet. 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.

  • 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:

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.
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.

const app = await createTestApp({ plugins: [myPlugin()], strict: true });
// a handler calling an undeclared path now fails the request

TypeScript covers more of this than you might expect: because each accessor is typed against the SDK's own service class, both a misspelled service (client.jbos) and a misspelled method (client.jobs.getRunz) are compile errors. The gap is a real method with no declared response — and any call that bypasses the types with a cast.

One more divergence: a service's methods are minted on access, so they are callable but not enumerable. typeof client.jobs.getRun is "function", but 'getRun' in client.jobs is false and Object.keys(client.jobs) is []. Plugin code that feature-detects with in or reflects over a service will therefore take a different branch than it does in production. This is deliberate: reporting those keys would make util.inspect probe each one, minting a mock per probe, which is the runaway recursion the default traps avoid.

Separately, createLakebasePool({ workspaceClient }) will build a pool whose password callback resolves to a mock: the pool exists but cannot connect. A Lakebase test needs a real database or a purpose-built fake pool, not this.

Full example

For a plugin you wrote, instantiate the class directly with new. The analytics() / agents() factory functions you pass to createApp return a descriptor for the app to construct, not an instance.

When you want an instance from one of those factories, use createTestPlugin rather than reaching through the descriptor:

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.

import { Plugin, type PluginManifest } from "@databricks/appkit";
import { expectStream, createMockRequest, createTestPluginContext } from "@databricks/appkit/testing";
import { describe, expect, test } from "vitest";

// A small plugin that registers a route and streams two events.
class GreeterPlugin extends Plugin {
  static manifest = {
    name: "greeter",
    displayName: "Greeter",
    description: "Example plugin",
    resources: { required: [], optional: [] },
  } as PluginManifest<"greeter">;

  async setup() {
    this.context?.addRoute("get", "/hello", (_req, res) => res.end());
  }

  async *greet(name: string) {
    yield { type: "greeting_start", name };
    yield { type: "greeting_end", message: `Hello, ${name}!` };
  }
}

describe("greeter plugin", () => {
  test("registers its route through the context", async () => {
    const mock = createTestPluginContext();
    const plugin = new GreeterPlugin({});

    await mock.attach(plugin);
    await plugin.setup();

    expect(mock.routes).toContainEqual(
      expect.objectContaining({ method: "get", path: "/hello" }),
    );
  });

  test("streams events in order", async () => {
    const plugin = new GreeterPlugin({});
    await expectStream(plugin.greet("world")).toEmit(
      "greeting_start",
      "greeting_end",
    );
  });
});

To test a plugin that dispatches cross-plugin tool calls, register fake providers and assert on mock.toolCalls — including asUser, which confirms the on-behalf-of path ran:

const mock = createTestPluginContext({ analytics: { query: [{ n: 1 }] } });
const plugin = new MyAgentPlugin({});
await mock.attach(plugin);

// `obo` sets the forwarded identity headers `asUser` needs — without them the
// dispatch would (correctly) reject with "Missing user token".
const req = createMockRequest({ obo: true });
await plugin.runSomethingThatCallsAnalytics(req);

expect(mock.toolCalls[0]).toMatchObject({
  plugin: "analytics",
  tool: "query",
  asUser: true,
});

See also

Databricks Developer Hub

Ready to ship your next agentic app in minutes?

Read docs