From ac23567c3c336496b897fb131c65fa5f4fc9c9e2 Mon Sep 17 00:00:00 2001 From: IamGalymzhan <62868459+IamGalymzhan@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:28:09 +0500 Subject: [PATCH 01/12] feat(appkit): add createTestApp, a never-crash mock client (#540) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(appkit): make PluginContext telemetry injectable The testing kit needs to construct a real PluginContext without a live OpenTelemetry pipeline. Add an optional constructor dependency for the telemetry provider, defaulting to the shared "plugin-context" provider so the production path is unchanged. This is the single production edit required to wrap the real class in tests rather than reimplementing it. Signed-off-by: Galymzhan * feat(appkit): ship @databricks/appkit/testing and migrate first stub Wire the testing kit as a published subpath and prove it against the first of the two hand-rolled context stubs (the design gate): - Add ./testing to both exports maps (dev + publishConfig) following the ./type-generator shape, add src/testing/index.ts to the tsdown entry, and declare vitest as an optional peerDependency. Build passes attw + publint; dist/testing/{index,mock-plugin-context,expect-stream,fixtures}.{js,d.ts} are emitted and vitest stays external to the main entry. - Migrate dispatch-tool-call.test.ts: replace (plugin as any).context = { executeTool } with mockPluginContext. executeTool is now the REAL method, so the forwarded toolCallTimeoutMs is asserted through actual signal composition, the on-behalf-of (asUser) path is verified, and a new test proves the forwarded timeout actually aborts a slow toolkit tool end-to-end. This is the primary win from the plan: executeTool's OBO and timeout paths gain real assertions instead of a stub that proved nothing. Signed-off-by: Galymzhan * test(appkit): migrate route-handler-errors context stub to mockPluginContext Replace the second and final hand-rolled stub — (plugin as any).context = { addRoute } — with the real PluginContext from mockPluginContext. The kit's route recorder captures raw handlers, so the alias assertion (both /invocations and /responses mount the same handler reference) holds against the real class, where forwardAsyncErrors wrapping would otherwise break reference identity. Both context stubs the plan identified are now migrated. Signed-off-by: Galymzhan * docs(appkit): document the testing kit and ship a template example test - Add docs/docs/development/testing.md covering mockPluginContext(), expectStream(), and the fixture helpers, with a full end-to-end example. Cross-links to local-development, custom-plugins, and execution-context. - Add template/server/example.test.ts: a self-contained, plugin-agnostic example that scaffolded apps ship with — it defines a tiny custom plugin and exercises both mockPluginContext (route recording) and expectStream (ordered event assertions), running with no workspace or network. Ships the kit to users, satisfying the plan's acceptance criteria that a docs page exists and the template carries at least one example test. Signed-off-by: Galymzhan * docs(appkit): fix testing-kit examples to instantiate the plugin class Validation by scaffolding a real app with `databricks apps init` surfaced that the examples called the `analytics()`/`toPlugin()` factory and then treated the result as a plugin instance — but a factory returns a { plugin, config, name } descriptor for createApp to construct, so `.attachContext`/handler methods are absent. Rewrite both the template example test and the docs "Full example" to instantiate the plugin class directly (`new GreeterPlugin({})`), matching how the migrated agents suites use the kit. The scaffolded app's `npm test` and `tsc` both pass against the published `@databricks/appkit/testing` subpath with no workspace or network. Signed-off-by: Galymzhan * refactor(appkit): tighten FakeToolResponse so a missing value is a type error Drop `undefined` from the static FakeToolValue union. `resolve()` treats an undefined map entry as "unregistered tool" and throws, so allowing undefined as a declared response made `{ query: undefined }` a confusing runtime error instead of a compile error. A function returning undefined still works for the rare "returns nothing" case. Add a test pinning that a null response is returned as a value, not misread as a missing tool. Signed-off-by: Galymzhan * refactor(appkit): make tools/test-helpers a shim over the shipped testing kit The plan's step 5 was to MOVE the fixtures into the package, not copy them. The shipped kit (src/testing/fixtures.ts) duplicated all 15 exports of tools/test-helpers.ts, which would drift over time. Collapse the original into a thin re-export of @databricks/appkit/testing so src/testing is the single source of truth while the 18 existing @tools/test-helpers importers keep working unchanged. The re-exported mockServiceContext is now synchronous; every call site either awaits it (no-op on a non-promise) or reads it through Awaited>, so all suites pass unchanged (full appkit suite: 3117 passed, 1 pre-existing skip). Signed-off-by: Galymzhan * fix(appkit): normalize CRLF in expectStream SSE parsing; sharpen testing docs Code review follow-ups: - expectStream's parseSSEBody split frames on \n\n, so a spec-compliant SSE stream delimited by \r\n\r\n (from a real server) collapsed into one event. AppKit's own writer uses \n\n so existing tests were unaffected, but expectStream is public API that accepts any Response. Normalize CRLF to LF before splitting; add a CRLF regression test. - Docs: instantiate the plugin CLASS in the attach() snippet (the factory returns a descriptor, not an instance), and note that the cache attach() seeds is a per-process singleton shared by tests within a file. Signed-off-by: Galymzhan * fix(appkit): resolve repo-wide Biome error blocking CI CI's "Lint & Type Check" job runs `pnpm run check` over the whole repo, so a pre-existing lint error unrelated to this branch failed the build: - remote-tunnel-controller.test.ts had two `afterEach` hooks in one describe (lint/suspicious/noDuplicateTestHooks, error severity). Merge them into one — behavior preserved (env reset + console-spy clear both still run after each test). This file is byte-identical to main; the error predated the branch and only surfaced because CI lints the entire tree. Also drop two dead `biome-ignore lint/suspicious/noExplicitAny` suppressions in the testing kit (fixtures.ts, expect-stream.test.ts): `noExplicitAny` is turned off repo-wide in biome.json, so the comments had no effect (suppressions/unused warnings). The invalid-source test now casts through `unknown as never`. Signed-off-by: Galymzhan * fix(appkit): address cross-model review findings in the testing kit Verified and fixed the findings from an independent code review: - #1 (correctness) expectStream dropped the wire `event:` name when the JSON payload carried its own `type` (spread ran after the assignment). Spread the payload first, then set `type = name ?? parsed.type`, so a frame like `event: error` + `data: {"type":"result"}` reports `error`. Regression test added. - #2 (contract) `@databricks/appkit/testing` eagerly loads vitest via fixtures even for `expectStream`, so vitest is a real requirement. Drop the "optional" peerDependenciesMeta and correct the docs sentence. - #6 (OBO fidelity) the fake `asUser` recorded `asUser: true` unconditionally. Enforce the real `Plugin.asUser` token precondition: a request without `x-forwarded-access-token` throws `missingToken` (missing user id throws too), and the resolved `userId` is recorded on each tool call. Tests now assert both directions (well-formed request vs token-less). - #3 (fidelity) attach() now mirrors AppKit core: registerPlugin plus registerToolProvider for real tool providers, without clobbering injected fakes. getPlugins()/getPluginNames()/hasPlugin() behave as in production. - #12 unknown-tool lookup used `tools[name] === undefined`, so a tool named "constructor"/"toString" hit Object.prototype. Use Object.hasOwn. - #5 drop data-less named SSE frames (real clients ignore them). - #7 re-export the PluginContext type from the testing barrel so MockPluginContext.ctx is nameable through the exports map. - #13 correct the docs: mock.telemetry captures the context's executeTool spans, not plugin-level spans (attachContext rebuilds the plugin's own telemetry). - #4 parseSSEResponse now delegates to the same parseSSEBody as expectStream — one parser, no divergence. All 3 analytics.integration call sites still pass. - #8 reformat template/server/example.test.ts with the template's Prettier so a scaffolded app's `npm run format` passes. - #10 fix the package-doc @example (agentsPlugin._handleStream does not exist). - #11 add kit tests that exercise attach() end-to-end (cache seed, isReady, registration, fake-not-clobbered). Build passes attw + publint; full appkit suite 3125 passed / 1 pre-existing skip. Signed-off-by: Galymzhan * chore(appkit): drop knip vitest-ignore now that vitest is a real peer dep With vitest declared as a (non-optional) peerDependency, knip recognizes it as used, so the earlier ignoreDependencies entry is unnecessary. This reverts knip.json to its original state. Signed-off-by: Galymzhan * fix(appkit): make vitest a normal dependency, not a package-wide peer A required peerDependency has no per-subpath scope: it applied to the whole @databricks/appkit package, so every production consumer that never imports the testing kit got an unsatisfied peer (npm 7+ auto-installs vitest into their tree; pnpm warns) — a wider blast radius than the eager-import bug it was meant to fix. Follow appkit's own precedent instead: `vite` backs the ./type-generator subpath as a normal `dependency`, installed for everyone but loaded only by importers of that subpath. Do the same for `vitest` and ./testing. vitest is referenced solely by dist/testing/fixtures.js, never by the main/plugin/core entry, so a consumer importing createApp never loads it. Verified end-to-end: scaffolded an app whose own vitest (4.1.9) differs in major from appkit's dependency (3.2.4), forcing a nested second copy. The testing kit's vi.fn()/vi.spyOn() mocks and expect(...).toHaveBeenCalled() assertions work across the two instances (vi spies carry their own call state), and npm install emits no peer-dep warning. Build passes attw + publint. Also fold in the template example's Prettier formatting (template uses Prettier, not Biome) so a scaffolded app's `npm run format` passes. Signed-off-by: Galymzhan * refactor(appkit): rename mockPluginContext to createTestPluginContext The helper builds the REAL PluginContext with faked edges — it does not mock the context — so the name was misleading. Rename to createTestPluginContext (and the MockPluginContext type to TestPluginContext), matching the create*-for-tests convention, and rename the files to test-plugin-context.ts. Pre-merge and unreleased, so no external consumers are affected. Also finish the #13 doc-accuracy fix in the shipped JSDoc (not just the docs page): the telemetry field comment now states it captures the context's spans (executeTool), not plugin-internal spans — attachContext rebuilds the plugin's this.telemetry from the real TelemetryManager. These comments ship in dist/testing/*.d.ts, so IntelliSense previously showed the unqualified claim. Build passes attw + publint; full appkit suite 3125 passed / 1 pre-existing skip. Signed-off-by: Galymzhan * refactor(appkit): dedupe testing fixtures and tidy test-plugin-context Behavior-preserving cleanups in the testing kit: - createMockRequest reuses createMockWorkspaceClient() instead of an inline copy of the same mock client (verified identical). - createMockServiceContext / createMockUserContext / mockServiceContext inline the createMockWorkspaceClient() call into the `||` fallback, so the mock client is built only when the caller did not supply one. - The fake asUser view spreads `...base` and overrides executeAgentTool rather than re-declaring getAgentTools. - expectStream's isSubsequence breaks once the expected sequence is fully matched. No semantic change; typecheck clean and all kit + migrated tests pass. Signed-off-by: Galymzhan * fix(appkit): resolve third-review findings in the testing kit - #1 (P1) The docs called vitest a peer dependency, but the manifest ships it under `dependencies` (the decision we landed on, matching how appkit ships `vite` for ./type-generator). Correct the docs to match: appkit installs vitest for you, and it loads only when you import ./testing. Manifest and docs now agree. - #2 (P2) expectStream buffered the source eagerly with no bound, so a non-terminating stream hung until the runner's own timeout. Add an optional `{ timeout }` that fails fast with a clear, kit-specific error; document it and cover both directions with tests. - #3 (P2) The fake asUser replicates asUser's token precondition but not the real dev-mode `DEV_OBO_FALLBACK_KEY` OTel marker (a module-private telemetry detail). Narrow the docs and JSDoc to say so and point users at the recorded asUser/userId fields instead of isDevOboFallback(). Build passes attw + publint; full appkit suite 3141 passed / 1 pre-existing skip. Signed-off-by: Galymzhan * test(appkit): dogfood the testing kit on analytics and genie plugins Exercise @databricks/appkit/testing against real core plugins to validate it beyond the two agent proof sites and produce usage references: - analytics.kit.test.ts: cross-plugin executeTool via createTestPluginContext — OBO identity (asUser/userId), token-precondition rejection, and per-call timeout abort. Needs only the kit (no workspace/ServiceContext). - genie.kit.test.ts: drives the real _handleSendMessage SSE stream and asserts event order with expectStream(...).toEmit(...). Both add genuinely new coverage (streamed SSE order + OBO dispatch identity were untested). Full appkit suite 3145 passed / 1 pre-existing skip. Developer-experience notes (kit wins + friction, e.g. createMockResponse doesn't compose with expectStream) captured in internal/ for the milestone review. Signed-off-by: Galymzhan * refactor(appkit): address testing-kit review feedback Resolve the eight review comments on the testing kit: - createMockResponse now captures written SSE bytes and exposes sseResponse(); expectStream reads a captured mock response directly, so streaming-route tests no longer need a hand-rolled bridge. - Ship vitest as an optional peer dependency (+ devDependency) instead of a plain runtime dependency, keeping the test framework out of production installs and deduping to the app's own copy. Ignore it in knip. - Add an obo option to createMockRequest so on-behalf-of tests set the forwarded identity headers with one flag. - Add resetTestCache() to clear the shared cache singleton between tests. - Use the documented attach() instead of an any-cast in the agents dispatch tests. - Drop the unused createMockServiceContext/createMockUserContext builders from the public surface; keep the service-context builder internal. - Pin the previously untested edges: the Object.hasOwn tool-lookup guard, the dev-mode asUser branch, and parseSSEBody's non-object data values. - Add useServiceContextMock() to register the mock lifecycle in one line, returning a live accessor. Dogfood the new helpers in the analytics, genie, and serving suites, and document them in the testing guide. Signed-off-by: Galymzhan * docs(appkit): move the testing guide under Plugins The testing kit is entirely plugin-scoped (createTestPluginContext, attach(plugin), plugin route/tool/SSE assertions), and the page's own cross-links already pointed into plugins/. Move it next to custom-plugins and fix the relative links. Keep the heading as 'Testing'; the Plugins section supplies the context. Signed-off-by: Galymzhan * test(appkit): fold dogfood tests into plugin suites Address round-2 review: the kit should be the default way to test a plugin, not a parallel '*.kit.test.ts' track. - Fold the three cross-plugin executeTool OBO tests into analytics.test.ts and delete analytics.kit.test.ts. - Upgrade genie.test.ts's SSE test to assert event ORDER via expectStream on genie's real event names (message_start, status, message_result, query_result), replacing brittle write.mock.calls substring checks, and delete genie.kit.test.ts. - Trim the heavy comment narration from the folded-in tests. - Re-export createTestPluginContext and expectStream from the test-helpers shim. - Finish the testing-guide move under plugins/ (sidebar position + links). Signed-off-by: Galymzhan * test(appkit): restore toHaveLength(1) on the analytics OBO dispatch test The dogfood fold trimmed expect(mock.toolCalls).toHaveLength(1), so a double-dispatch would no longer fail the happy-path test — and it was inconsistent with the token-less sibling that kept toHaveLength(0). Restore it. Signed-off-by: Galymzhan * test(appkit): re-assert genie SSE payloads after the expectStream swap The toEmit swap pinned event order but dropped the payload values the old substring checks covered (conversationId=new-conv-id, status=ASKING_AI), which aren't asserted elsewhere. Restore them structurally via collect() + toMatchObject — keeping the ordering guarantee without brittle substrings. Signed-off-by: Galymzhan * fix(appkit): drop fabricated workspace-client fields from createMockRequest createMockRequest returned userWorkspaceClient, serviceWorkspaceClient, getWarehouseId and getWorkspaceId — fields no production code reads (plugins resolve those through getWorkspaceClient()/getWarehouseId() from src/context, which mockServiceContext stands in for). Publishing them via @databricks/appkit/testing would make four inert fields a permanent public promise. The two warehouse cold-start tests (analytics + metric) overrode mockReq.serviceWorkspaceClient.warehouses.get, which the route never reads — so they passed on the default RUNNING client without exercising the warehouse path at all. Route the warehouse client through mockServiceContext (the real seam) so the tests are live, and drop the 'mock WorkspaceClient' claim from the testing guide. Signed-off-by: Galymzhan * feat(appkit): add a never-crash mock WorkspaceClient to the testing kit Every core plugin's actual work runs through getWorkspaceClient(), which the testing kit did not fake — so a jobs/genie/serving/files plugin crashed on its first client call and authors hand-rolled nested client literals instead. createMockWorkspaceClient() fakes the whole facade in three layers: - The 9 facade members are explicitly typed, so `client.jbos` is a compile error. The facade is closed and AppKit-owned, so there is no per-service fixture to maintain as the SDK grows. - Each service is a Proxy minting one memoized vi.fn() per method name, keyed by dotted path. `client.jobs.getRun === client.jobs.getRun`, so call assertions work, and the legacy view shares the map so one `responses` entry covers both — including un-faceted services like `legacy.clusters.list()`. - `config` and `apiClient` are seeded objects rather than bare Proxies, because three of their members must not be mocks: `config.host` is a real string that production code builds URLs from and throws on when falsy, `apiClient.userAgent()` must be synchronous (a Promise inside a Headers value stringifies to "[object Promise]"), and `apiClient.request` resolves {} so destructuring its result does not throw. Two guards keep the Proxy safe. Symbol keys delegate to Reflect.get, and a passthrough deny-set answers `undefined`. `then` is the load-bearing entry: without it a service looks thenable, so `await client.jobs` either hangs or resolves to a mock's return value. ownKeys is left at its default so util.inspect and toEqual see {} instead of recursing forever. The three historical canned defaults are byte-identical, because 13 test files reach them implicitly through mockServiceContext. `currentUser.me` is additive and load-bearing: ServiceContext.createContext reads `currentUser.id`, so an unresolved me() is a TypeError and createApp({ client }) cannot boot without it. getMockFn(client, "jobs.getRun") is the typed assertion path — facade accessors are legacy-SDK-typed, so expect(client.jobs.getRun).toHaveBeenCalled() does not typecheck. It mints idempotently, so the handle can be grabbed before the code under test runs. The compile-time block is enforced by tsc, not at runtime. It records one correction to the plan: the SDK types `config.host` as `string | undefined`, so the contract is that it narrows to a string, not that it is non-optional. 4451 tests pass (+38); the 667 tests reaching the default client indirectly through mockServiceContext are unchanged. Co-authored-by: Isaac Signed-off-by: Galymzhan * refactor(appkit): converge the two mock-workspace-client builders fixtures.ts had its own two-service createMockWorkspaceClient, so the shipped fixture and the new never-crash builder were near-duplicates. The fixture now re-exports the builder and the barrel points at its new home. The blast radius is entirely indirect. Nothing in src imports the exported fixture by name (connectors/genie/tests/client.test.ts defines its own local one), but buildServiceContextState calls it as the default client for mockServiceContext, which 13 test files use. The risk therefore lives in the default return value, which is why U1 kept the three canned defaults byte-identical — and why this commit adds the convergence guard that asserts both halves: jobs/genie now resolve instead of throwing "Cannot read properties of undefined", while the SQL path those 13 files depend on still succeeds. createConfigurableMockWorkspaceClient is left byte-for-byte unchanged and only gains a @deprecated notice. Its bare vi.fn()s return undefined *synchronously* whereas the new floor returns Promise, and its one caller (analytics.integration.test.ts) can observe that difference; reimplementing it here would change behaviour for no benefit. It migrates with that suite later. The jobs suite drops its hand-rolled client literal — the seven method mocks plus the config.host/authenticate block — onto the builder, which is the proof the boilerplate actually goes away. Its 57 assertion sites move to a getMockFn handle because facade accessors are legacy-SDK-typed, so .mockResolvedValue on them does not typecheck. The factory needs `await vi.hoisted(async ...)` with a dynamic import, since a hoisted factory runs before the file's imports. 4454 tests pass (+3). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): split the lifecycle exit from the teardown LifecycleManager's shutdown sequence was reachable only by killing the process, so nothing could release AppKit's sockets, timers, pools, cache, and telemetry and keep running. That is what blocks an app handle's close(), and with it any test that wants to boot more than once in a file. The sequence is now a phase runner that returns an exit code, a promise memo, and two thin callers: - shutdown() is the signal path, observably unchanged: it arms the same unref'd 15s force-exit backstop and still exits 0 on completion, 1 on an unexpected throw. The timer stays here deliberately — it is the one thing close() must not inherit, since a programmatic caller wants a logged error when teardown hangs, not a dead process. - close() is the programmatic path: it detaches signal handlers, runs the same phases under a shorter default budget (5s, not the production 15s), logs the phase that was in flight if the budget is spent, and never exits. Replacing the isShuttingDown boolean with a promise memo is a strict improvement. The boolean made a second caller return *immediately* while teardown was still running — harmless for a signal, since the first caller exits the process anyway, but for close() it would resolve before resources were released, which is the difference between a correct handle and a misleading one. The read and the assignment stay in one synchronous statement, preserving the invariant the boolean was there to protect. One production behaviour does shift: a second signal now awaits the first teardown. installSignalHandlers registered anonymous arrows that could never be removed. The [signal, handler] pairs are now retained and detached individually, never via removeAllListeners, so a host embedding AppKit keeps its own handlers. The tests assert that with two managers installed, a.close() leaves b's pair and an unrelated host listener intact, and that counts return to their pre-install baseline — which is what stops repeated boots tripping MaxListenersExceededWarning. The signal-mid-close race is documented rather than papered over: handlers come off before the first await, and if a signal still lands it joins the memo and exits, because it wanted the process dead. The idempotency test is verified by injection — it fails against the old return-immediately semantics and passes against the memo. Its first draft did not: it counted microtask ticks, which cannot distinguish an early return through close()'s raceWithTimeout wrapper. It now asserts that neither caller settles until the plugin hook has actually completed. 4463 tests pass (+9); the 14 pre-existing shutdown tests are untouched. Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): expose close() on the app handle createApp acquired sockets, timers, and pools but returned no way to release them, so the only teardown was killing the process. The LifecycleManager built at the end of _createApp was constructed and immediately discarded; it is now retained on the instance and reachable through the handle. The return type widens from PluginMap to AppHandle, which is PluginMap plus close() and Symbol.asyncDispose. Widening a return type is source-compatible for every existing caller, and the cast that produces the handle already hid instance methods, so close() rides along naturally. onPluginsReady deliberately keeps PluginMap: it runs before the server starts, so handing it a close() would invite a footgun for no gain. The name collision is a real hazard, not a theoretical one. Plugin exports are installed with Object.defineProperty, and an own property shadows a prototype method — so a plugin named `close` would silently replace teardown rather than merely confuse the types. Three layers guard it: Symbol.asyncDispose is unreachable from a manifest name, so `await using` is always safe; createAndRegisterPlugin now throws a ConfigurationError naming the offending plugin; and no plugin in the repo is affected. Coverage is deliberately unmocked, because the claim is about real resources: a boot on an ephemeral port serves /health, close() runs the plugin's shutdown hook, the socket stops accepting, and the SIGTERM listener count returns to its pre-boot baseline. Also covered: idempotency at the app level, a server-less app closing cleanly, `await using` releasing at scope exit, and the reserved name being rejected. Verified by injection — with close() stubbed to a no-op and the reserved-name guard removed, 5 of the 6 fail. 4469 tests pass (+6). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): make the process-wide singletons re-bootable close() released resources but left the singletons pointing at them, so close() followed by createApp() silently reused what the teardown had just torn down. This delivers the actual driver — boot, assert, close, repeat. CacheManager.reset() drops both `instance` and `initPromise`. Clearing only `instance` is insufficient because getInstance() returns `initPromise` when `instance` is null, so the next boot would await a promise resolving to the dead manager. Testing surfaced a third case the plan missed: clearing both is *still* not enough, because an initialization already in flight runs its continuation and re-publishes the very instance being discarded. A generation counter now invalidates that write. The covering test models PersistentStorage rather than using the default in-memory storage. This matters: InMemoryStorage.close() only clears a Map and stays usable, so an in-memory test passes whether or not the reset exists — which is precisely why the bug hid. Against storage whose close() is terminal, the way pool.end() is, the test shows the stale manager throwing "Cannot use a pool after calling end()" and the reset fixing it. One plan claim is corrected rather than implemented. The plan asserted that TelemetryManager's never-cleared `shutdownPromise` made a second shutdown() return a stale promise and skip flushing a re-initialized SDK. It does not: shutdown() only returns the memo after reassigning it for whatever SDK is currently live, so a stale resolved promise can be returned only when there is no SDK to flush. Verified twice — by mocking NodeSDK across three initialize/shutdown cycles, and by running the original implementation in isolation. An earlier draft of this commit added a generation counter here too; it has been reverted, since it fixed nothing and cost a field. What TelemetryManager did need, and now has, is the static reset() that drops the singleton. The resets are wired into close() only, never the signal path, where the process is dying and pointer drops are pure cost. Symmetry is the justification: core initializes all four in _createApp, so core drops all four. This is a semantic expansion, not purely a bug fix — a host that closes and then expects ServiceContext.get() to work will now get an InitializationError. resetAppKitSingletons() is published from @databricks/appkit/testing for tests that hand-roll createApp and would otherwise deep-import ../context/service-context to reach ServiceContext.reset(). Both it and LifecycleManager.close() delegate to one core-side implementation rather than duplicating the list. resetTestCache() is untouched — it calls clear() on the existing cache, a different and still-useful operation. 4480 tests pass (+11). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): add createTestApp — the customer-grade plugin test harness One call boots a real AppKit app with no workspace, no credentials, and no network, and calls it over real HTTP: const app = await createTestApp({ plugins: [myPlugin()] }); const res = await app.post("/api/my-plugin/thing", { body, obo: true }); await expectStream(res).toEmit("status", "result"); await app.close(); Four of the setup steps exist only because of hazards found by reading the boot path, and each has a test that fails without it: - NODE_ENV is pinned away from "development". Not tidiness: dev mode routes the injected `port: 0` through get-port, where portNumbers(0, …) throws a RangeError. "development" is refused outright with an explanation rather than worked around, since dev mode also boots a real Vite server, downgrades resource validation to a warning, and stops filtering dev-only plugins. - DATABRICKS_WORKSPACE_ID is set, short-circuiting the SCIM probe in getWorkspaceId, and internal telemetry is disabled. Both would otherwise fire apiClient.request during boot. A canary test asserts zero calls after boot, so either regression fails loudly. - The cache gets explicit in-memory storage. Without it CacheManager builds its own workspace client — ignoring the injected one — and probes Lakebase over the network, so "no network" would be false. - The server plugin is reached through a lazy `await import()`, because it runs dotenv.config() at module load. A static import would mutate a consumer's process.env merely by importing the testing entry point. process.env is snapshotted wholesale rather than by whitelist, since plugins read vars the harness cannot enumerate, and restored on close() — including deleting keys the harness added and restoring a pre-existing DATABRICKS_HOST to its own value rather than the test default. Teardown also runs from the boot-failure path, or a plugin whose setup() throws would leak env mutations into every later test in the file. Plugin exports live under app.plugins rather than spread onto the handle: `get` and `delete` are plausible plugin names and would collide with the request methods. The request methods return a native Response, so expectStream composes with no bridge — the dogfooding report's top friction, avoided by construction. `obo` reuses createMockRequest's OboOption rather than inventing a second convention. Two corrections to the plan, both found by testing: - A `strictValidation: false` opt-out was specified and has been dropped as a false affordance. enforceValidation computes `shouldThrow = !isDevelopment || strict`, so with NODE_ENV pinned away from "development" validation always throws and the flag cannot do anything. The env var is still set as belt-and-braces, and a test pins the unconditional behaviour. - The error-middleware test initially asserted a redacted body. It is not redacted: errorHandlerMiddleware hides the message only under NODE_ENV=production, and the harness pins "test". Useful for tests — an assertion can name the failure — but it means that response is the dev shape, which the test now says out loud. The HTTP suite's probe plugin registers routes through `this.route()`, the way real plugins do. Registered with raw `router.get()` a rejection escapes forwardAsyncErrors and hangs the request — correct AppKit behaviour, and worth having a representative test rather than a misleading one. 4511 tests pass (+31). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): publish the harness surface, add createTestPlugin, and dogfood both Audit first: the ./testing subpath was already in both exports maps and the tsdown entry, vitest is already an optional peer dep, and attw/publint pass. The real gap was proof that a test needs nothing else, so the integration suites moved onto the public entry point — that migration *is* the audit. A new acceptance suite imports only from @databricks/appkit/testing and boots, requests, asserts a stream, and closes. Self-referencing the package from inside it needed a tsconfig paths entry. Resolving the package's own export map made tsc's project root ambiguous (TS2209), and the alias mirrors how shared and @databricks/lakebase are already mapped. It also makes source resolution deterministic rather than depending on the "development" export condition — verified by marker: the subpath resolves to src, not dist. createTestPlugin(factory, config) closes the last dogfooding footgun. Reaching through a descriptor with `new (genie({}).plugin)(config)` skips DEFAULT_CONFIG and forgets `name`, so the instance under test is configured differently from the one production builds. It mirrors createAndRegisterPlugin's merge order. createTestApp does not subsume it: the harness takes descriptors and builds instances itself, so the unit path needs its own ergonomics. Dogfooding results, reported as measured rather than as hoped: - analytics.integration.test.ts: 300 -> 216 lines. Setup/teardown went 104 -> 55, against the plan's predicted ~30. Its local getListeningPort helper is gone and its 12 mock handles now come from getMockFn. Same 6 tests, same assertions. - getListeningPort is lifted into the kit, and files/plugin.integration.test.ts imports it instead of carrying its own copy. - server.integration.test.ts moves four of its five blocks to ephemeral ports. The fifth keeps its fixed port deliberately, because it asserts the server honours a configured one; a comment says so. The removed sleep-100ms waits are replaced by getListeningPort, which waits on the listening event instead of guessing. One plan claim corrected: the hardcoded TEST_PORT = 9879 said to collide with server.integration was already fixed on this branch — analytics had moved to port: 0. The real fixed ports were the five in server.integration itself, which is what this commit addresses instead. Docs lead with createTestApp: a which-harness comparison table, the dotted-path responses convention, the teardown contract, a "Mocking Databricks services" section carrying the never-crash floor's honest catch (a misspelled *method* returns undefined, and a Lakebase pool built on the fake cannot connect), and an explicit callout that manifest.config.schema is not validated. The PluginContext/ServiceContext boundary note now says the kit covers the data plane. The template example gains a createTestApp test. 4519 tests pass (+8). pnpm docs:build is clean. Co-authored-by: Isaac Signed-off-by: Galymzhan * test(appkit): tighten the mock-client type contract from tarball findings Verified the kit end to end the way the plan prescribes: pnpm pack:sdk, an app scaffolded by `databricks apps init` from this repo's template, the tarballs installed into it, and the suites run with no .env, no credentials, and every non-loopback socket connection hard-blocked. Nine customer-style tests plus the template's three pass, and the scaffolded app typechecks against the shipped .d.ts. That run corrected a claim this branch had been making. Both the plan's risk table and the docs said a misspelled *method* slips through the never-crash floor and only a misspelled *service* is caught. Not so: each facade accessor is typed against the SDK's own service class, so `client.jobs.getRunz` and `client.files.anything` are compile errors too. The compile-time block now asserts that for three services, and the docs say what the real gap is — a method that exists but has no declared response, or a call that bypasses the types with a cast. Also repointed one doc line that told readers to reach the client via `getWorkspaceClient()`. That is right inside this repo but wrong from the published entry, where the name currently resolves to Lakebase's unrelated `getWorkspaceClient(config)`. The docs now use `getExecutionContext().client`, which is exported and works. The mis-export itself is a main-entry defect, outside this branch's scope, and is left for a follow-up. 4519 tests pass. Build, docs:build, attw, and publint are clean, and the packed tarball carries dist/testing/*.js and .d.ts for every new module. Co-authored-by: Isaac Signed-off-by: Galymzhan * fix(appkit): address code-review findings in the testing kit Eight reviewers over the branch diff produced 24 findings; 14 were actionable. The two real defects, both in code this branch added: - **An orphaned teardown could tear down the *next* app's resources.** `close()` races the shutdown phases against a 5s budget, then resets the core singletons and resolves. The phases keep running. A plugin `shutdown()` hook slower than that budget but inside its own 10s per-plugin budget — which the files plugin's drain can be — left phase 5 re-reading the static slots, so it either skipped draining this app's cache pool or closed the *following* app's storage and shut down its OTEL SDK. The comment claiming the phases had "already closed the cache storage" was only true when teardown finished in budget. Phase 5 now uses the instances captured before the first await, and a test drives the exact 5s-to-10s window (verified by reintroducing the bug). - **`process.env` restore did not compose across overlapping boots.** Each app snapshotted independently, so a second boot captured the first's mutations and whichever closed last re-applied them, stranding the harness keys and the first app's `env` entries after both apps were gone. Confirmed by probe, in-repo and from the packed tarball. There is now one reference-counted baseline: the first live app anchors it, the last one to close restores it, and the outcome no longer depends on close order. Also fixed: `server: false` alongside a caller-supplied server plugin is now refused instead of half-honoured (the plugin still bound a socket while the handle denied one existed); `AppHandle.close()` declares the `{ timeoutMs }` the implementation accepts, so the harness no longer casts to reach it; the duplicate `listeningPort` helper in the close integration suite is gone in favour of the kit's (two reviewers flagged it); the analytics suite drops an `as never` that erased `app.plugins` typing; the `clientFns` WeakMap moved above its users; and a comment claiming "9 typed facade members" over a 7-element array is corrected. Two of my own tests were weak and are now stronger: the `authenticate` test wrapped its whole body in `if (mockFn)` and only asserted "was called" — it now asserts the Authorization header it claims to set — and a close-after-signal test proved ordering by counting microtask ticks, which cannot see through `raceWithTimeout`; it uses the same sentinel the sibling test does. A new compile-time assertion pins that `AppHandle` still satisfies a `PluginMap` annotation, so a regression in the widening can't pass silently. Documented rather than changed: a service's methods are callable but not enumerable, so `'getRun' in client.jobs` is false and `Object.keys` is empty. Reporting those keys would make `util.inspect` mint a mock per probe, which is the recursion the default traps exist to avoid. Also documented why `onPluginsReady` keeps the narrower `PluginMap`. One finding rejected as a false positive: project-standards reported CLAUDE.md still documents Biome. It does not — main's own oxlint migration (9538d58e) updated it, and only the pre-merge copy said Biome. Six findings were demoted to residual risks, chiefly the P1 claim that the mock resolving `undefined` for undeclared paths lets a test pass while production is broken. That is the deliberate, documented contract of the never-crash floor, not a defect; an independent reviewer re-deriving it argues the existing caution callout is warranted, not that the design changed. 4524 tests pass. Re-verified end to end from a repacked tarball in the `databricks apps init` app with all non-loopback sockets blocked. Co-authored-by: Isaac Signed-off-by: Galymzhan * docs(appkit): document `close` as a reserved plugin name `createApp()` rejects a plugin whose manifest name is `close`, because plugin exports are installed as own properties and an own property shadows a prototype method — so such a plugin would silently replace the app handle's teardown. The thrown ConfigurationError already names the offending plugin, but nothing told an author the constraint existed before they hit it. Noted beside where custom-plugins.md introduces `static manifest`. Landing this as part of the `feat:` framing for the branch rather than a BREAKING CHANGE footer: the failure is loud and at boot, not a silent runtime change, and no plugin in this repo is affected. Signed-off-by: Galymzhan * chore: drop the **/.claude ignores from knip, oxlint, and oxfmt These were added earlier on this branch to work around a locked agent worktree at .claude/worktrees/, which every tool saw as a second full copy of the repo: knip reported hundreds of phantom unused exports and failed the pre-commit hook outright, and a repo-root oxfmt would have rewritten that other branch's tree. The worktree has since been removed, so the ignores are treating a symptom that no longer exists and are out of scope for this branch. Verified after removal: `pnpm knip` and `pnpm check` both exit 0 at the repo root. .oxfmtrc.json and .oxlintrc.json are now byte-identical to origin/main. The one remaining knip.json difference — ignoreDependencies: ["vitest"] for packages/appkit — predates this work and is required because vitest is an optional peer dependency of the published testing subpath. Signed-off-by: Galymzhan * refactor(appkit): slim the comments added by the testing-kit work The comments on this branch were far past the repo's own density: reset.ts was 85% comment (35 of 41 lines) for a one-line function, create-test-plugin.ts 63%, lifecycle-manager.ts 48%, mock-workspace-client.ts 45%. Much of it restated the code or ran to several paragraphs where a clause would do. Net 364 comment lines removed. Every file now sits at or below the repo baseline (main's own sources run 20-42%): mock-workspace-client 45% -> 21%, create-test-app 36% -> 25%, lifecycle-manager 48% -> 35%, reset.ts 41 -> 13 lines total, create-test-plugin 63 -> 27. What was kept is the non-obvious "why" that a maintainer would otherwise delete and reintroduce a bug: that `then` must stay in the deny-set or `await client.jobs` hangs; that `ownKeys` stays default or util.inspect mints a mock per probe; that config.host must be a real string; that the three canned defaults are byte-identical because 13 suites depend on them; that phase 5 captures its singletons before the first await; and the four boot hazards behind createTestApp's setup. Pre-existing comments in files this branch only touched (fixtures.ts, test-plugin-context.ts, the shutdown() phase list) are left alone — reverting other people's prose is not this change's business. Also dropped an unnecessary `as Any` cast in createTestPlugin: DEFAULT_CONFIG is already declared on PluginConstructor, so the type escape and its explanatory comment both went. 4524 tests pass; lint, format, and typecheck clean. Signed-off-by: Galymzhan * test(appkit): merge duplicate assertions in the mock-client suite The mock-workspace-client suite had 37 tests written against the plan's checklist rather than against behaviours, so eight asserted something a sibling already covered: getRun memoization twice, config.host being a string twice, the 9-member facade twice (one a strict subset of the other), a rejecting function response twice, getMockFn path resolution twice, function-valued responses twice, the canned defaults twice, and two config-option tests that fit in one. 29 tests now, with no assertion lost — where a dropped test had a unique claim it was folded into the survivor. Three describe blocks became empty and were removed; one had only a comment saying its subject could not be tested at runtime, which the compile-time contract block covers properly. Note for anyone reading this as a bundle-size fix: it is not one. Tests do not ship — the packed tarball contains zero test files — and dropping these eight moved the measured bundle by exactly 0 bytes. The comment trimming in the previous commit is what actually helped (+8.1% -> +6.9%), because JSDoc is preserved in the emitted .d.ts. 4516 tests pass. Signed-off-by: Galymzhan * test(appkit): compact the testing-kit suites The new suites were verbose in ways that cost a reviewer without buying coverage. 1837 lines across 8 files -> 1529 across 7. - mock-workspace-client.test.ts: 471 -> 253. The nine-accessor walk and the canned defaults become one test.each table; per-test client construction goes through a short factory alias. Verified by injection: dropping `then` from the deny-set, changing a canned default, and adding an ownKeys trap each still fail the suite, so the compaction did not gut it. - create-test-app-http.test.ts is folded into create-test-app.test.ts. Both tested one unit through two near-identical 100+ line plugin fixtures; there is now one probe plugin and one manifest helper. Six of the nine routes in the old echo fixture were dead — they duplicated the HTTP file's own plugin. - A withApp() helper absorbs the boot/try/finally/close block that appeared 13 times. It is generic over the plugin tuple so app.plugins stays typed. - The three /headers tests differed only in inputs and expectations, so they are one test.each. The cache reset suite's storage double loses its repeated ended-guard boilerplate. 4509 tests pass; typecheck, lint, and format clean. Signed-off-by: Galymzhan * fix(appkit): address four review findings in the testing kit **Restore three upstream mlflow tests dropped by the merge.** Resolving the route-handler-errors conflict with `--ours` took the whole file from this branch, discarding main's non-conflicting additions from PR #477: the `vi.mock("../mlflow")` hoist, the linkTraceToRun/mockTraceId resets, the parameterised seedPlugin(adapter), the seedEchoPlugin/invoke helpers, and the three trace tests. 15 upstream tests, 12 here, and nothing failed to say so. Restored from 9538d58e alongside this branch's createTestPluginContext rewrite of the aliases test — the two changes are independent. **close() now memoizes itself, not just the phases.** runOnce() guaranteed the teardown body ran once, but resetCoreSingletons() sat outside it, so every call reset again: `await a.close(); createApp(); await a.close()` dropped the second app's singletons. The reset is also skipped when the budget expired, because the phases are still running and still own those instances. Only reachable through the raw AppHandle — createTestApp's wrapper memoizes, which is why the harness-level test could not see it and the regression test lives in app-close.integration. **Refcount singleton ownership.** The env baseline was already refcounted so overlapping harness apps compose, while the singleton layer reset on every boot and every close — so booting B rebound A's ServiceContext and CacheManager, and closing A while B was live left B with none at all. claimCoreSingletons/ releaseCoreSingletons now follow the same model as the env baseline: first boot claims, last close drops. **Fake the on-behalf-of client.** The kit promised "no workspace, no credentials, no network", but createApp({ client }) installs only the service principal; an `obo` request reached ServiceContext.createUserContext, which builds a real SDK client from process.env.DATABRICKS_HOST. The harness now stubs that for the app's lifetime and restores it on close, mirroring fixtures.ts's createUserContextSpy. Every fix has a test verified by reintroducing the bug. Two of those tests needed a second attempt: asserting the OBO client's host does not discriminate, since a real client carries the same DATABRICKS_HOST string — the test now asserts the harness's mock recorded the call. The probe plugin gained a route that calls the client under asUser, because the existing /as-user route only reads ctx.userId, which is how this escaped notice. 4517 tests pass; typecheck, lint, and format clean. Signed-off-by: Galymzhan * chore: restore the dev-playground client lockfile Swept into merge 45601bf1 from a dirty working tree; not part of this work. npm had run in apps/dev-playground/client during the local-tarball verification and pruned the `extraneous: true` entries for ../../../packages/appkit-ui. Both parents of that merge carry the same blob, so `git log -- ` reports nothing for this branch even with --full-history, which is why it went unnoticed. Restored to origin/main byte-for-byte. Signed-off-by: Galymzhan * fix(appkit): correct double singleton release and simplify the testing kit createTestApp's failure path released its singleton claim twice when the boot failed after createApp resolved: app.close() already drops the claim, and the catch block released it again, which could pull the singletons out from under a still-live sibling app. It now releases only when nothing was booted. Alongside it, a simplification pass over the branch: - LifecycleManager.closeOnce: drop the timedOut flag in favour of an early return from the catch. - createTestApp: drop the restoreEnv alias for releaseEnvBaseline. - Delete createConfigurableMockWorkspaceClient. Its last caller moved to createMockWorkspaceClient earlier in this branch, so the JSDoc rationale for keeping it byte-for-byte pointed at deleted code, and the ./testing subpath it would have shipped on is new here and unreleased. - Reuse the kit's createMockRequest instead of three local mockReq helpers that re-rolled the same forwarded-identity headers. - genie.test.ts: reuse one expectStream handle rather than parsing the same captured SSE body twice. Signed-off-by: Galymzhan * chore: drop the .claude ignores and restore the client lockfile Both were swept into the oxlint merge from a dirty working tree and corrected later on the branch; folding those corrections in here keeps them out of the follow-up PR. The knip `.claude/**` entry and the `**/.claude` ignorePatterns in oxfmt/oxlint were never needed — nothing in the repo lints or formats that directory. The `packages/appkit` vitest ignoreDependencies entry stays: vitest is a real dependency of the testing entry. apps/dev-playground/client/package-lock.json is restored to origin/main byte-for-byte; npm had run in that directory and pruned its `extraneous: true` entries. Signed-off-by: Galymzhan * refactor(appkit): trim the cache reset test and its comments The 3-line `CacheManager.reset()` carried 30 lines of comment and 135 lines of test. Mutation testing showed what each test was actually worth: with `reset()` stubbed to a no-op, or clearing only `instance`, four tests fail; with the generation guard removed, exactly one does. - Dropped "reset is safe when the cache was never initialized". It survived all three mutations — a body of three assignments cannot throw, so the test could only ever pass. - Replaced the hand-rolled 23-line `CacheStorage` double with a 7-line `InMemoryStorage` subclass overriding just `close()` and `set()`. It also stops claiming `isPersistent() === true`, which had the manager's probabilistic cleanup eligible to fire against ended storage. - Cut the comments to the two facts a maintainer would otherwise remove and reintroduce the bug with: both fields must clear because `getInstance()` falls back to `initPromise`, and a reset is a pointer drop so callers close first. Same treatment for `reset-singletons.ts` and `testing/reset.ts`, which were at 47% and 52% comment lines against a repo baseline of 20-35%. Test file 135 -> 115 lines, production diff +38 -> +23. Mutation coverage is unchanged, re-verified against all three mutations. 4517 tests pass; typecheck, lint, and format clean. Signed-off-by: Galymzhan * docs(appkit): relocate and tighten the app-handle comments 60% of this branch's additions to appkit.ts were comment: 41 lines over 22 lines of code. Trimmed to 35 with no fact dropped — audited claim by claim — and two of them moved somewhere they do more good. **The onPluginsReady note was on the wrong function.** It sat on the internal `_createApp`, which typedoc does not publish, while the public `createApp` that callers actually read carried the same parameter with no explanation. Moved, and the generated API page now renders it (see the Function.createApp.md diff) where before it reached nobody. While moving it, corrected the hazard it described. It said offering `close()` there "would invite tearing down a half-booted app" — but `#lifecycle` is not assigned until after the server starts, so `close()` at that point is a *no-op*, not a teardown. The narrow type and the `#lifecycle?.` optional chain guard against silently skipping cleanup, which is the opposite failure. **RESERVED_PLUGIN_NAMES gained the reasoning for its scope.** It reserves only `close`, which reads like an oversight: `bindExportMethods` and the other prototype methods are equally shadowable, since TS `private` is compile-time only. The distinction is that shadowing those throws `TypeError` on the next plugin's registration, while shadowing `close` fails silently — you call it, get no error, and leak every socket and pool. Only silent breakage needs a guard, and that is now written down. Also dropped the comment above the LifecycleManager construction, which restated the `#lifecycle` field's own JSDoc. 4517 tests pass; typecheck, format, and docs build clean. Signed-off-by: Galymzhan * test(appkit): cover the singleton release on close()'s timeout path Mutation-tested the close-handle suites (33 tests, 8 mutations). One mutation survived: deleting the early `return` in `closeOnce`, so a timed-out close() releases the core singletons while its own phases are still running and still own those instances. That is the exact hazard the code comment warns about, and it was the review fix with no coverage. The test that sounds like it covers this — "phase 5 still closes the app's own cache and telemetry, not the next app's" — cannot: it mocks CacheManager and TelemetryManager wholesale, so whether releaseCoreSingletons() ran is invisible to it. It guards the captured-singleton half of the fix, not the early return. Closed by mocking the one symbol lifecycle-manager imports from reset-singletons and extending two existing tests, so the count stays at 24: the clean-close test now asserts one release, and the hung-teardown test asserts none. Verified in both directions — dropping the early return fails the second, removing releaseCoreSingletons() entirely fails the first. Also corrected a comment in the phase-5 test that predated the review fix. It said close() "already dropped the singletons" on timeout, which is what the fix stopped it from doing. 4517 tests pass; typecheck and format clean. Signed-off-by: Galymzhan * fix(appkit): lowercase mock request header keys, as Express does `createMockRequest` stored header keys exactly as given while `header()` lowercased the lookup, so a mixed-case override was unreachable: createMockRequest({ obo: { userId: "alice" }, headers: { "X-Forwarded-User": "bob" } }); // header("x-forwarded-user") === "alice" Both keys were kept — ["x-forwarded-access-token", "x-forwarded-user", "X-Forwarded-User"] — and the lowercase one obo seeded still answered, which contradicted the "an explicit override wins" contract documented right above it. Keys are now lowercased on the way in, matching what Node's parser hands Express. Thanks @pkosiec. The existing override test passed because it used a lowercase key, so it is now parametrised over both casings, and the case-insensitivity test additionally pins that every stored key is lowercase. Reverting the fix fails both. Signed-off-by: Galymzhan * fix(appkit): lowercase caller header keys in the harness request methods Same root cause as the `createMockRequest` fix on the parent branch (#530, found by @pkosiec), and worse here. `Object.assign(headers, reqOptions.headers)` kept case variants as separate keys, and `Headers` **comma-joins** duplicates rather than replacing them: new Headers({ "x-forwarded-user": "alice", "X-Forwarded-User": "bob" }) // -> x-forwarded-user: "alice, bob" So a mixed-case override did not merely lose, it corrupted the value the server received — the "caller headers last" contract broken in a way that produces a plausible-looking string rather than an error. Keys are lowercased before assignment. The existing `/headers` table gained a mixed-case row rather than a new test; reverting the fix fails it. Signed-off-by: Galymzhan * chore(appkit): acknowledge the testing kit's tarball growth in the baseline The bundle-size gate is the last red check on this PR. The growth is the feature, not slack: the tarball ships no test files at all, so deleting eight tests moved the measured size by exactly 0 bytes. Comment trimming did help, taking it from +8.1% to +6.9%, because JSDoc survives into the `.d.ts`. `@databricks/appkit` packed 860106 -> 919127 (+58 KB), which is the `./testing` entry: 61 KB gz of harness plus its type declarations. `@databricks/appkit-ui` moves -291 bytes, incidentally. Regenerated with `pnpm size:baseline` from a clean build — `rm -rf packages/*/dist packages/*/tmp` first, because tsdown runs with `clean: false` and a stale artifact would be baked into the baseline as real growth. `pnpm size:compare` now reports no change and exits 0. Signed-off-by: Galymzhan * Revert "chore(appkit): acknowledge the testing kit's tarball growth in the baseline" This reverts commit 04839f0ba41608be2894600f5c44651e997a0dbd. * fix(appkit): address Pawel's review — singleton leak, renames, and guards Nine of the seventeen review comments, the ones with no open decision. **The singleton leak (P2, both instances).** On a `close()` timeout the release was skipped deliberately — the phases still own those instances — but nothing released it afterward either, so the refcount never returned to zero and every later boot skipped its reset and inherited a half-closed app. Now the release is attached to the still-running teardown, so it happens once the phases settle. This also fixes the boot-failure report: that path calls `app.close()`, so it routed through the same gap. My own test was pinning the bug. It asserted `releaseCoreSingletons` was never called on the timeout path, when the contract is "not yet, but once teardown settles". Corrected, and verified by reverting the fix. **Renames, while the pre-release window is open** (both symbols are absent from 0.62.0, so this is free now and breaking later): `getMockFn` -> `getMock`, since it returns vitest's `Mock`; `resetAppKitSingletons` -> `resetGlobalState`, since vitest's own resets are `reset` with no product prefix. **`getListeningPort` was `@internal` and exported**, which is a contradiction either way. Made public: three integration suites already consume it through the public entry, and it is a reasonable helper for anyone hand-rolling a server. **The reserved-name list is now exhaustive by construction.** It was a hand-kept `Set(["close"])` with nothing tying it to `AppHandle`, so a future named method there would be silently shadowable by a plugin. A `Record, true>` makes adding one a compile error until it is reserved — verified by adding a `restart()` and watching tsc fail. **Passing both `client` and `responses` now throws** instead of silently ignoring the responses, matching the existing `server: false` conflict. **The two proxy traps share one guard.** `toLegacyWorkspaceClient` had copied the symbol + PASSTHROUGH_DENY logic, so a key added to the set would have covered one trap and missed the other — and missing `then` there is what makes `await client` hang. **Coverage gap closed:** the boot-failure branch that runs *after* `createApp` succeeded had no test. Added one via a plugin that boots and then throws when the harness reaches for its socket. **Comment de-fragilised:** the canned-defaults note named a count ("13 suites") and a position ("the first three"). It now names the three entries. 4521 tests pass; typecheck, lint, format, and docs build clean. Signed-off-by: Galymzhan * fix(appkit): match the real user-context shape in the harness stub Half of Pawel's P2 on the diverging on-behalf-of fakes. The other half is deferred on semver grounds, recorded in internal/testing-kit/README.md. `stubUserContext` set `tokenFingerprint: test-${userId}` — keyed on the *user*, so constant across tokens. Lakebase rotates its pool by comparing that value (`pool-manager.ts:102`), which means rotation could never fire under `createTestApp` while production rotates on every new token. A test exercising pool rotation through the harness would have passed against a code path that cannot run. It now derives `sha256(token)[0:16]`, exactly as the real `createUserContext` does. The rejection also matched only in spirit: a bare `Error` where production throws `AuthenticationError.missingToken`. A handler that catches that class took a different branch under the harness. Now the same class. Two tests, both verified by reintroducing the old behaviour: the fingerprint is stable per token and differs across tokens (a user-keyed value fails it), and the missing-token rejection is asserted on the error class, not the message — a plain `Error` whose message still says "token" fails it. Not touched: `createUserContextSpy` in fixtures.ts, which neither throws nor sets the field. It shipped in 0.62.0, so tightening it would break consumers' existing tests on a minor upgrade; it needs its own PR and a changelog line. The end state is one shared builder, worth extracting once that half is free to move. Signed-off-by: Galymzhan * fix(appkit): share one user-context fake across the testing kit Completes Pawel's P2. Three fakes of `ServiceContext.createUserContext` disagreed and none matched production; they now share `fakeUserContext` in fixtures.ts — production's rejection (`AuthenticationError.missingToken`), production's `sha256(token)[0:16]` fingerprint, and `userEmail`. What the divergence cost: `pool-manager.ts:101` treats a missing `tokenFingerprint` as "not stale", so the old fixtures spy — which discarded the token outright — made Lakebase's drain-and-recreate branch unreachable. A test written to verify pool rotation would have passed while exercising nothing. I deferred this half earlier on semver grounds, arguing that `mockServiceContext` shipped in 0.62.0 so tightening a published test double would break consumers on a minor upgrade. Wrong: `./testing` was not published at all before 0.62.0 (`publishConfig.exports` lists no `./testing` in 0.61.0 or 0.61.1), and 0.62.0 landed four days ago with two releases since. There is no population to break. Checking the API's actual age beats reasoning about semver in the abstract. Harmless in-repo too — the full suite passed with the fake tightened, so nothing here relied on the looseness either. Six tests pin it, three per fake, all verified by restoring the loose version: the fingerprint is stable per token and differs across tokens, a missing token throws production's error class, and userEmail is carried through. 4579 tests pass; typecheck, lint and format clean. Signed-off-by: Galymzhan * docs(appkit): deslop the testing guide and fix its sidebar slot Pawel's P1 on the guide's voice, plus the two smaller docs points. **Voice.** Ran the repo's doc-deslop pass, then a targeted follow-up. No vocabulary slop or hedging to remove — the divergence from sibling pages was punctuation and register: paired em-dashes injecting lists mid-sentence (the textbook tell), British `behaviour` against the siblings' American spelling, and cutesy editorial phrasing. `:::caution The honest catch` is now `:::caution Undeclared methods return undefined`, matching how siblings title admonitions. Worth recording how the "how florid is it" question actually resolved, because my first two measurements were wrong. Raw em-dash counts said this page was 2.7x denser than `analytics.md`. But that counted comments inside code blocks and the `**term** — definition` bullet pattern the siblings use too. Excluding both: 0.286 dashes per prose line here against 0.252 in `analytics.md`, its closest sibling by length and depth. The page was already at parity, so the remaining dashes were left alone rather than purged below the house norm. Four where a dash was doing a colon's job are now colons. **Teardown now leads with `await using`.** It previously said "use try/finally, or let the runtime do it" and then showed only the `await using` example — naming the weaker option first and never demonstrating it. `await using` comes first with the reason (it closes on a thrown error too), and `try/finally` follows as the form you need when the app outlives a block. Both are supported: TypeScript 5.9.3, Node 24, ES2022, and three suites already rely on it. **Sidebar.** `sidebar_position: 8` collided with `jobs.md` and `manifest.md`; `_category_.json` uses `autogenerated`, so frontmatter really does drive order. Moved to 10, which is unused. The pairs at 2, 5, 6, 7, 8 and 9 predate this work and are left for a separate tidy-up. Also documents two things the code now does but the page didn't say: passing `client` and `responses` together is refused rather than ignored, and `getListeningPort` is public API. Signed-off-by: Galymzhan * chore(docs): regenerate the appkit-ui stylesheet Unrelated to the testing kit. `styles.gen.css` is generated by the docs build and main's committed copy is stale: it lacks the `.collapse` utility. Any `pnpm docs:build` reproduces this diff, so it otherwise sits dirty in every working tree. Verified deterministic — reverted, rebuilt, and it came back. Signed-off-by: Galymzhan * feat(appkit): add strict mode to the mock workspace client Pawel's P3 on undeclared paths resolving `undefined`. I first declined this as a feature deserving its own PR; the actual change is eight lines of production code and additive, so that was over-caution. `createMockWorkspaceClient({ strict: true })` throws when a path with no declared response is *called*, naming the path. Off by default — the never-crash floor is what lets a plugin touch services a test does not care about — so nothing about existing behaviour moves. Three details that needed care: - **Throws on call, not on mint.** `getMock(client, path)` mints so a test can hold a handle before the code under test runs; blowing up there would break the normal assertion pattern. Pinned by a test. - **The canned defaults count as declared**, so a harness boot still works — it reads `currentUser.me` through this client. With `defaults: false` they are undeclared again, which is also pinned. - **Plumbed through `createTestApp`.** Without that the option was unreachable from the recommended entry point: the harness builds its own mock, and passing a hand-built client now refuses `responses`. The conflict guard covers `strict` too. This does not overlap with the typed-`responses`-keys idea, contrary to the review's framing that they are two answers to the same problem: typed keys catch a *typo* at compile time, `strict` catches an *omission* at runtime. Typed keys cannot know you forgot to declare a path you actually call. Six tests, and the guide documents it beside the caution that describes the failure mode. 4585 tests pass; typecheck, lint, format and docs build clean. Signed-off-by: Galymzhan * fix(appkit): address the straightforward half of the second review pass Nine of Pawel's thirteen round-2 comments, all verified true first. The two P1s, the template CI wiring, and the code-comment deslop are left: they need the concurrency decision or a wider change. Recorded in internal/testing-kit/review-log.md. **The timeout release is now immediate, reversing last round's fix.** His first pass said to defer it with `.finally()`; his second says that is wrong, and he is right. Holding the refcount makes the next boot skip its reset, so it inherits this app's `CacheManager` — which the orphaned teardown then closes in phase 5, mid-test. Releasing at once is safe because `runPhases` captures its own cache and telemetry before the first await and never re-reads the shared slots. Shorter and safer. My test asserted the deferred behaviour, so it was pinning the wrong contract — the second time a test of mine has done that on this exact code. **A seeded `apiClient.userAgent` no longer returns a Promise.** Every non-function seed was wrapped in `mockResolvedValue`, so `responses: { "apiClient.userAgent": "x" }` produced `"[object Promise]"` in a Headers value — precisely what the synchronous default exists to prevent. Seeds now match each member's own shape; `request` stays async. The `config.*` path never had this because it assigns values directly. **The published-surface test now guards the published surface.** It imported three symbols, so every other export could vanish from the barrel with this file still green. It now asserts twenty names are reachable through the entry — verified by dropping `createTestPluginContext` and watching it fail by name. `tsc` would also catch that via other suites, which softens his "CI still green" wording, but a test claiming to guard the surface should do it. Also: dead `createHash`/`AuthenticationError` imports removed (left by the `fakeUserContext` move); the boot-failure path now honours `closeTimeoutMs`; the two bare `type Any = any` aliases carry the explanation the third one had; the "seven services" count is gone from a comment; the guide now says the client and the OBO stub are process-wide, not per app, and bounds its `app.client` claim to a single open app; and it distinguishes mock from fake in one line. Finally the `nodeEnv` JSDoc records what the option actually changes beyond refusing `development` — `errorHandlerMiddleware` redacts 5xx bodies only under `production`. That is the evidence his round-1 comment proposing to delete the option overlooked, and it was undocumented, which is a fair reading of why he missed it. 4587 tests pass; typecheck, lint, format and docs build clean. Signed-off-by: Galymzhan * fix(appkit): close the leftover halves of three review comments The docs half of the stale-count comment was never done: testing.md said "AppKit owns that 9-member interface" while `WorkspaceClient` declares ten members (seven services, `config`, `apiClient`, `toLegacyWorkspaceClient`). Drop the count and keep the point, which never needed it. The published-surface test asserted `createTestPluginContext` was a key on the namespace but never called it, so a hollowed-out export would still pass. It now drives the context's real tool registry and on-behalf-of path through the entry. Verified by disabling the dispatch recording: the new test fails, the name check still passes. Also reword the three story-voice comments quoted in review that survived the earlier pass; the other three were already rewritten with the timeout-release fix. Signed-off-by: Galymzhan * fix(appkit): allow one harness app at a time Both P1s from the second review pass are consequences of letting two createTestApp apps overlap, which nobody decided to support. AppKit's workspace client, cache, and on-behalf-of fake are process-wide, so a second live app cannot own its own: its handlers resolve the first app's client while `app.client` returns the second mock, and `close()` restores the real `createUserContext` rather than the other app's spy, putting the survivor's on-behalf-of calls back on the network. Refuse the second boot instead. The guard reuses the existing live-app counter rather than adding a parallel flag that could drift, and runs before any mutation so a refused boot leaves the open app untouched — asserted, not assumed. Vitest isolates test files in separate workers, so this only constrains apps within one file. The cost lands on a `describe` that holds an app in `beforeAll`: it can no longer contain a test that boots its own, which is why the concurrency and on-behalf-of tests move to the top level. Two overlapping-boot env tests are replaced by one covering repeated boot/close cycles; the hazard they pinned (a second snapshot capturing the first's mutations) is now unreachable. Verified by disabling the guard: the three refusal tests fail. Signed-off-by: Galymzhan * ci(appkit): run the scaffolded template's tests, and ignore its staging dir template/server/example.test.ts ran nowhere, so the example shipped to anyone scaffolding an app could rot silently. Wiring it up turned out to need one step rather than a new job: pr-template-artifact already builds and packs the branch, rewrites the template's deps to those tarballs, and npm installs with devDependencies, so the test file and vitest are already staged and installed. The step runs before the zip, so a broken example fails the build instead of being published as a downloadable artifact. It is also the only CI check that consumes @databricks/appkit/testing the way a customer does — through a real npm install of a packed tarball, rather than a workspace path. Verified locally end to end: pack, stage, install, npm test (3 passed), and grep on the installed dist confirms it ran against the branch's tarball, not the published version template/package.json pins. The lockfile is byte identical across the test step, so the registry-rewrite step that follows is unaffected. Also gitignore pr-template/ and appkit-template-*.zip, which prepare-template-artifact.ts leaves in the working tree — a ~200MB staging directory that `git add -A` would otherwise commit. Signed-off-by: Galymzhan * refactor(appkit): drop the dead re-exports from the deprecated test-helpers shim tools/test-helpers.ts is marked @deprecated, yet this branch had added eight exports to it: createTestApp, CreateTestAppOptions, getListeningPort, getMock, MockWorkspaceClient, resetGlobalState, TestApp, TestRequestOptions. None had a consumer. The fifteen files still importing the shim use only nine symbols between them, all predating this work, so the eight came along mechanically with the getMockFn -> getMock rename and never had a caller — visible in that they were inserted mid-list, breaking the alphabetical order the rest of the list keeps. Removing them keeps the shim's surface to what an existing importer actually needs, so the harness API has one import path rather than two. Typecheck and the full suite pass unchanged, which is what confirms they were dead. Signed-off-by: Galymzhan * refactor(appkit): drop inert refcount from singleton-reset machinery createTestApp is the only claim site and forbids a second concurrent harness app unconditionally, while production createApp never claims, so `owners` was always 0 at claim and 0 at release: claim, release, and drop were the same operation and the counter arithmetic was a no-op. The memo in LifecycleManager.close() (`this.closed ??= ...`), not the refcount, is what stops a stale handle's second close() from resetting a newer app. - reset-singletons.ts: remove `owners`, `claimCoreSingletons`, `releaseCoreSingletons`; keep only `dropCoreSingletons`. - create-test-app.ts / lifecycle-manager.ts: call `dropCoreSingletons` directly (import from ../core/reset-singletons); close() memo preserved. - reset.ts: drop the `@internal` claim/release passthroughs, keep the public `resetGlobalState`. - liveHarnessApps: collapse the 0-or-1 counter into a boolean. No behaviour change; the no-concurrency guard and the CacheManager entry in the reset list are untouched. Signed-off-by: Galymzhan * refactor(appkit): drop the multi-boot/embeddable surface from the testing kit AppKit will not support multiple apps in one Node process, so the programmatic teardown / embeddable-app surface the testing-kit PR reached for comes out (addresses Mario's #540 review): - drop AppKit.close(), [Symbol.asyncDispose], and the AppHandle return type; createApp() returns PluginMap again, and the close-name-shadow guard goes - fold the LifecycleManager close/runOnce/removeSignalHandlers split back into shutdown() plus an internal dispose() for the harness; the harness boots with an internal installSignalHandlers:false flag rather than removing handlers - createTestApp tears down what it booted (dispose runs the plugin shutdown hooks and closes the server) then drops the core singletons and restores env - relocate dropCoreSingletons() into testing/ (kit-owned; no core caller left) - drop the cache reset-generation guard: no reset-during-init race under single-boot with awaited sequential teardown (invariant documented in-code) - delete testing/reset.ts (resetGlobalState) — no callers - cover the rewired teardown end-to-end: a booted plugin's shutdown() hook fires and the server socket is released after close() Signed-off-by: Galymzhan * refactor(appkit): collapse the shutdown/dispose split into shutdown({ exit }) (#540 review) Mario's M7: the shutdown()/dispose() split existed only so the harness could run the teardown phases without process.exit. Fold dispose() into shutdown({ exit?: boolean }) — exit defaults true (signal path: force-exit backstop + process.exit); { exit: false } is the harness path (await the phases, no exit). AppKit[disposeApp] now calls shutdown({ exit: false }). With the split gone, runPhases()'s capture-before-await of the cache/telemetry managers is provably dead: the orphan race it guarded (a timed-out dispose returning while the phases still ran, then the harness dropping singletons) cannot occur once the harness fully awaits the phases before dropping. Same basis as the 6dffb59d CacheManager reset-generation-guard removal ("no race under single-boot with awaited sequential teardown"). closeCacheStorage() and flushTelemetry() read their manager at phase-5 time again; DISPOSE_TIMEOUT_MS is deleted (the phases are already individually bounded). Also drop the now-inert createTestApp closeTimeoutMs option (0 refs; the harness path has no outer budget of its own any more). Tests: the harness-path describe uses shutdown({ exit: false }); the two tests that only exercised the removed orphan / outer-timeout machinery are replaced by one asserting the harness path is bounded by the internal per-plugin timeout. Local commit for #540 review; not pushed. Signed-off-by: Galymzhan --------- Signed-off-by: Galymzhan --- .github/workflows/ci.yml | 10 + .gitignore | 4 + docs/docs/api/appkit/Function.createApp.md | 18 +- docs/docs/plugins/testing.md | 211 ++++- docs/static/appkit-ui/styles.gen.css | 6 + packages/appkit/src/cache/index.ts | 18 + .../cache/tests/cache-manager-reset.test.ts | 100 +++ packages/appkit/src/core/appkit.ts | 41 +- packages/appkit/src/core/lifecycle-manager.ts | 106 ++- .../src/core/tests/lifecycle-manager.test.ts | 201 +++++ .../agents/tests/dispatch-tool-call.test.ts | 22 +- .../agents/tests/route-handler-errors.test.ts | 103 ++- .../tests/analytics.integration.test.ts | 185 ++-- .../plugins/analytics/tests/analytics.test.ts | 1 - .../files/tests/plugin.integration.test.ts | 29 +- .../src/plugins/genie/tests/genie.test.ts | 5 +- .../src/plugins/jobs/tests/plugin.test.ts | 187 ++-- .../server/tests/server.integration.test.ts | 40 +- .../appkit/src/telemetry/telemetry-manager.ts | 18 + .../tests/telemetry-manager-reset.test.ts | 184 ++++ .../appkit/src/testing/create-test-app.ts | 436 +++++++++ .../appkit/src/testing/create-test-plugin.ts | 27 + packages/appkit/src/testing/fixtures.ts | 137 ++- packages/appkit/src/testing/index.ts | 22 +- .../src/testing/mock-workspace-client.ts | 319 +++++++ .../appkit/src/testing/reset-singletons.ts | 35 + .../src/testing/tests/create-test-app.test.ts | 849 ++++++++++++++++++ .../testing/tests/create-test-plugin.test.ts | 82 ++ .../appkit/src/testing/tests/fixtures.test.ts | 46 + .../tests/mock-workspace-client.test.ts | 307 +++++++ .../published-surface.integration.test.ts | 153 ++++ .../testing/tests/test-plugin-context.test.ts | 7 +- packages/appkit/tsconfig.json | 4 +- template/server/example.test.ts | 44 +- tools/test-helpers.ts | 5 +- 35 files changed, 3507 insertions(+), 455 deletions(-) create mode 100644 packages/appkit/src/cache/tests/cache-manager-reset.test.ts create mode 100644 packages/appkit/src/telemetry/tests/telemetry-manager-reset.test.ts create mode 100644 packages/appkit/src/testing/create-test-app.ts create mode 100644 packages/appkit/src/testing/create-test-plugin.ts create mode 100644 packages/appkit/src/testing/mock-workspace-client.ts create mode 100644 packages/appkit/src/testing/reset-singletons.ts create mode 100644 packages/appkit/src/testing/tests/create-test-app.test.ts create mode 100644 packages/appkit/src/testing/tests/create-test-plugin.test.ts create mode 100644 packages/appkit/src/testing/tests/mock-workspace-client.test.ts create mode 100644 packages/appkit/src/testing/tests/published-surface.integration.test.ts 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..a623d8837 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 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/plugins/testing.md b/docs/docs/plugins/testing.md index 30fc66642..ca02a5aba 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -1,26 +1,151 @@ --- -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. +## 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 | | --- | --- | @@ -58,6 +183,8 @@ 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()`: ```ts @@ -96,7 +223,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. +`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()`. @@ -131,7 +258,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,6 +270,10 @@ 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.) @@ -160,10 +291,70 @@ 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. - `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 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 -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/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/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/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/analytics/tests/analytics.integration.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts index e099c8350..c69a82e83 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,37 @@ import { } from "vitest"; import { AppManager } from "../../../app"; -import { ServiceContext } from "../../../context/service-context"; -import { createApp } from "../../../core"; -import { server as serverPlugin } from "../../server"; import { analytics } from "../index"; const getAppQuerySpy = vi.spyOn(AppManager.prototype, "getAppQuery"); -/** - * Wait for the supplied server to finish binding, then return the OS-assigned - * port. Required when the test passes `port: 0` to `serverPlugin` — - * `app.server.start()` returns as soon as `listen()` is invoked but before the - * bind completes, so `server.address()` returns `null` until the `listening` - * event fires. - */ -async function getListeningPort(server: Server): Promise { - const addr = server.address(); - if (addr && typeof addr === "object" && typeof addr.port === "number") { - return addr.port; - } - await new Promise((resolve, reject) => { - server.once("listening", () => resolve()); - server.once("error", (err) => reject(err)); - }); - const ready = server.address(); - if (!ready || typeof ready !== "object") { - throw new Error("Server is listening but address() returned null"); - } - return ready.port; -} - describe("Analytics Plugin Integration", () => { - let server: Server; - let baseUrl: string; - let serviceContextMock: Awaited>; - let mockClient: ReturnType; + let app: TestApp<[ReturnType]>; + /** The SQL mock the analytics route drives, via the harness's client. */ + let executeStatement: ReturnType; + let getStatement: ReturnType; beforeAll(async () => { - setupDatabricksEnv(); - ServiceContext.reset(); - - mockClient = createConfigurableMockWorkspaceClient(); - serviceContextMock = await mockServiceContext({ - serviceDatabricksClient: mockClient.client, - }); - - const app = await createApp({ - plugins: [ - // port: 0 → OS assigns an ephemeral port. Avoids EADDRINUSE / cross-test - // route bleed when another integration test (e.g. server.integration) - // holds a fixed port concurrently in the shared vitest worker pool. - serverPlugin({ - port: 0, - host: "127.0.0.1", - }), - analytics({}), - ], - }); - - server = app.server.getServer(); - const port = await getListeningPort(server); - baseUrl = `http://127.0.0.1:${port}`; + // The harness owns the env setup, the singleton resets, the mock client, the + // server plugin on an ephemeral port, and the teardown. + app = await createTestApp({ plugins: [analytics({})] }); + executeStatement = getMock( + app.client, + "statementExecution.executeStatement", + ); + getStatement = getMock(app.client, "statementExecution.getStatement"); }); afterAll(async () => { getAppQuerySpy?.mockRestore(); - serviceContextMock?.restore(); - if (server) { - await new Promise((resolve, reject) => { - server.close((err) => { - if (err) reject(err); - else resolve(); - }); - }); - } + await app?.close(); }); beforeEach(() => { - mockClient.mocks.executeStatement.mockReset(); - mockClient.mocks.getStatement.mockReset(); + // Reset drops the built-in canned SUCCEEDED default too, matching the + // "script it yourself" semantics this suite relied on before. + executeStatement.mockReset(); + getStatement.mockReset(); getAppQuerySpy.mockReset(); }); @@ -119,18 +69,13 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockResolvedValueOnce( + executeStatement.mockResolvedValueOnce( createSuccessfulSQLResponse(mockData, mockColumns), ); - const response = await fetch( - `${baseUrl}/api/analytics/query/test_query`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response = await app.post("/api/analytics/query/test_query", { + body: { parameters: {} }, + }); expect(response.status).toBe(200); expect(response.headers.get("Content-Type")).toBe( @@ -144,8 +89,8 @@ describe("Analytics Plugin Integration", () => { { name: "Bob", age: "25" }, ]); - expect(mockClient.mocks.executeStatement).toHaveBeenCalledTimes(1); - expect(mockClient.mocks.executeStatement).toHaveBeenCalledWith( + expect(executeStatement).toHaveBeenCalledTimes(1); + expect(executeStatement).toHaveBeenCalledWith( expect.objectContaining({ statement: testQuery, warehouse_id: "test-warehouse-id", @@ -162,26 +107,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 +134,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 +153,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 +172,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 +193,23 @@ describe("Analytics Plugin Integration", () => { isAsUser: false, }); - mockClient.mocks.executeStatement.mockResolvedValue( + executeStatement.mockResolvedValue( createSuccessfulSQLResponse([["cached_value"]], [{ name: "value" }]), ); - const response1 = await fetch( - `${baseUrl}/api/analytics/query/cache_test`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response1 = await app.post("/api/analytics/query/cache_test", { + body: { parameters: {} }, + }); const data1 = await parseSSEResponse(response1); - const response2 = await fetch( - `${baseUrl}/api/analytics/query/cache_test`, - { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ parameters: {} }), - }, - ); + const response2 = await app.post("/api/analytics/query/cache_test", { + body: { parameters: {} }, + }); const data2 = await parseSSEResponse(response2); expect(data1.data).toEqual([{ value: "cached_value" }]); expect(data2.data).toEqual([{ value: "cached_value" }]); - expect(mockClient.mocks.executeStatement).toHaveBeenCalledTimes(1); + expect(executeStatement).toHaveBeenCalledTimes(1); }); }); }); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 5101f9424..2151103eb 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -1708,7 +1708,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/files/tests/plugin.integration.test.ts b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts index 3be68b315..0134c09e6 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts @@ -1,6 +1,10 @@ import http, { type Server } from "node:http"; -import { mockServiceContext, setupDatabricksEnv } from "@tools/test-helpers"; +import { + getListeningPort, + mockServiceContext, + setupDatabricksEnv, +} from "@databricks/appkit/testing"; import { afterAll, beforeAll, @@ -67,29 +71,6 @@ const MOCK_AUTH_HEADERS = { /** Volume key used in all integration tests. */ const VOL = "files"; -/** - * Wait for the supplied server to finish binding, then return the - * OS-assigned port. Required when tests pass `port: 0` to `serverPlugin` - * — `appkit.server.start()` returns as soon as `listen()` is invoked but - * before the bind completes, so `server.address()` returns `null` until - * the `listening` event fires. - */ -async function getListeningPort(server: Server): Promise { - const addr = server.address(); - if (addr && typeof addr === "object" && typeof addr.port === "number") { - return addr.port; - } - await new Promise((resolve, reject) => { - server.once("listening", () => resolve()); - server.once("error", (err) => reject(err)); - }); - const ready = server.address(); - if (!ready || typeof ready !== "object") { - throw new Error("Server is listening but address() returned null"); - } - return ready.port; -} - describe("Files Plugin Integration", () => { let server: Server; let baseUrl: string; diff --git a/packages/appkit/src/plugins/genie/tests/genie.test.ts b/packages/appkit/src/plugins/genie/tests/genie.test.ts index 2d867d335..2ada2ef3d 100644 --- a/packages/appkit/src/plugins/genie/tests/genie.test.ts +++ b/packages/appkit/src/plugins/genie/tests/genie.test.ts @@ -339,13 +339,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..8933d9eed 100644 --- a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts @@ -12,38 +12,47 @@ import { import { mapParams } from "../params"; import { JobsPlugin, jobs } from "../plugin"; -const { mockClient, mockCacheInstance } = vi.hoisted(() => { - const mockJobsApi = { - runNow: vi.fn(), - submit: vi.fn(), - getRun: vi.fn(), - getRunOutput: vi.fn(), - cancelRun: vi.fn(), - listRuns: vi.fn(), - get: vi.fn(), - }; - - const mockClient = { - jobs: mockJobsApi, - config: { - host: "https://test.databricks.com", - authenticate: vi.fn(), - }, - }; +const { mockClient, jobsApi, mockCacheInstance } = await vi.hoisted( + async () => { + // The testing kit's fake, not a hand-rolled literal: the seven jobs methods, + // `config.host` as a real string, and `config.authenticate` all come for free, + // and any *other* service this plugin grows into resolves instead of throwing. + // Imported inside the hoisted factory because the factory runs before the + // file's own imports are evaluated. + const { createMockWorkspaceClient, 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 mockCacheInstance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async (_key: unknown[], fn: (signal?: AbortSignal) => Promise) => - fn(), - ), - generateKey: vi.fn(), - }; + const mockCacheInstance = { + get: vi.fn(), + set: vi.fn(), + delete: vi.fn(), + getOrExecute: vi.fn( + async ( + _key: unknown[], + fn: (signal?: AbortSignal) => Promise, + ) => fn(), + ), + generateKey: vi.fn(), + }; - return { mockJobsApi, mockClient, mockCacheInstance }; -}); + return { mockClient, jobsApi, mockCacheInstance }; + }, +); vi.mock("../../../workspace-client", async (importOriginal) => { const actual = @@ -290,7 +299,7 @@ describe("JobsPlugin", () => { test("runNow passes configured job_id to connector", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const exported = plugin.exports(); @@ -298,7 +307,7 @@ describe("JobsPlugin", () => { await handle.runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123 }), expect.anything(), ); @@ -307,7 +316,7 @@ describe("JobsPlugin", () => { test("runNow merges user params with configured job_id (no taskType)", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const exported = plugin.exports(); @@ -317,7 +326,7 @@ describe("JobsPlugin", () => { notebook_params: { key: "value" }, }); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123, notebook_params: { key: "value" }, @@ -349,7 +358,7 @@ describe("JobsPlugin", () => { test("runNow maps validated params to SDK fields when taskType is set", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({ jobs: { @@ -363,7 +372,7 @@ describe("JobsPlugin", () => { await handle.runNow({ key: "value" }); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 123, notebook_params: { key: "value" }, @@ -375,7 +384,7 @@ describe("JobsPlugin", () => { test("runNow skips validation when no schema is configured", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -388,7 +397,7 @@ describe("JobsPlugin", () => { test("getRun wraps call in execute", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ + jobsApi.getRun.mockResolvedValue({ run_id: 1, state: { life_cycle_state: "TERMINATED" }, }); @@ -415,7 +424,7 @@ describe("JobsPlugin", () => { test("getJob wraps call in execute", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.get.mockResolvedValue({ job_id: 123 }); + jobsApi.get.mockResolvedValue({ job_id: 123 }); const plugin = new JobsPlugin({}); const executeSpy = vi.spyOn(plugin as any, "execute"); @@ -439,7 +448,7 @@ describe("JobsPlugin", () => { test("listRuns clamps caller-supplied limit before calling the SDK", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockReturnValue((async function* () {})()); + jobsApi.listRuns.mockReturnValue((async function* () {})()); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -447,7 +456,7 @@ describe("JobsPlugin", () => { await handle.listRuns({ limit: 10000 }); // SDK should receive the clamped limit, not the caller-supplied 10000. - expect(mockClient.jobs.listRuns).toHaveBeenCalledWith( + expect(jobsApi.listRuns).toHaveBeenCalledWith( expect.objectContaining({ limit: 100 }), expect.anything(), ); @@ -457,8 +466,8 @@ describe("JobsPlugin", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Pre-flight getRun verifies the run belongs to the configured jobId. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 1, job_id: 123 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 1, job_id: 123 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const executeSpy = vi.spyOn(plugin as any, "execute"); @@ -478,8 +487,8 @@ describe("JobsPlugin", () => { test("runAndWait yields status updates and terminates on TERMINATED", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); - mockClient.jobs.getRun + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.getRun .mockResolvedValueOnce({ run_id: 42, state: { life_cycle_state: "RUNNING" }, @@ -505,7 +514,7 @@ describe("JobsPlugin", () => { test("runAndWait throws when runNow returns no run_id", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({}); + jobsApi.runNow.mockResolvedValue({}); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -521,7 +530,7 @@ describe("JobsPlugin", () => { test("runNow returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockRejectedValue(new Error("API timeout")); + jobsApi.runNow.mockRejectedValue(new Error("API timeout")); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -538,9 +547,7 @@ describe("JobsPlugin", () => { test("cancelRun returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.cancelRun.mockRejectedValue( - new Error("Permission denied"), - ); + jobsApi.cancelRun.mockRejectedValue(new Error("Permission denied")); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -556,9 +563,7 @@ describe("JobsPlugin", () => { test("getRun returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockRejectedValue( - new Error("Internal server error"), - ); + jobsApi.getRun.mockRejectedValue(new Error("Internal server error")); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -574,7 +579,7 @@ describe("JobsPlugin", () => { test("listRuns returns error result on execute failure", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockImplementation(() => { + jobsApi.listRuns.mockImplementation(() => { throw new Error("Auth failure"); }); @@ -594,7 +599,7 @@ describe("JobsPlugin", () => { const error = new Error("Detailed internal failure: db connection reset"); (error as any).statusCode = 403; - mockClient.jobs.getRun.mockRejectedValue(error); + jobsApi.getRun.mockRejectedValue(error); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -611,7 +616,7 @@ describe("JobsPlugin", () => { test("successful operations return ok result with data", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -628,7 +633,7 @@ describe("JobsPlugin", () => { test("getRun returns 404 when run.job_id does not match configured jobId", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -641,8 +646,8 @@ describe("JobsPlugin", () => { test("getRunOutput returns 404 when run belongs to another job", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); - mockClient.jobs.getRunOutput.mockResolvedValue({ logs: "nope" }); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.getRunOutput.mockResolvedValue({ logs: "nope" }); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -651,14 +656,14 @@ describe("JobsPlugin", () => { expect(result.ok).toBe(false); if (!result.ok) expect(result.status).toBe(404); // Should never have called getRunOutput on the upstream SDK - expect(mockClient.jobs.getRunOutput).not.toHaveBeenCalled(); + expect(jobsApi.getRunOutput).not.toHaveBeenCalled(); }); test("cancelRun returns 404 when run belongs to another job", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const handle = plugin.exports()("etl"); @@ -666,13 +671,13 @@ describe("JobsPlugin", () => { const result = await handle.cancelRun(99); expect(result.ok).toBe(false); if (!result.ok) expect(result.status).toBe(404); - expect(mockClient.jobs.cancelRun).not.toHaveBeenCalled(); + expect(jobsApi.cancelRun).not.toHaveBeenCalled(); }); test("getRun succeeds when run.job_id matches configured jobId", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ + jobsApi.getRun.mockResolvedValue({ run_id: 42, job_id: 123, state: { life_cycle_state: "TERMINATED" }, @@ -696,7 +701,7 @@ describe("JobsPlugin", () => { const { JobsConnector } = await import("../../../connectors/jobs"); const connector = new JobsConnector({}); - mockClient.jobs.get.mockResolvedValue({ job_id: 123 }); + jobsApi.get.mockResolvedValue({ job_id: 123 }); const controller = new AbortController(); await connector.getJob( @@ -718,8 +723,8 @@ describe("JobsPlugin", () => { test("runAndWait stops polling when signal is aborted", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); - mockClient.jobs.getRun.mockResolvedValue({ + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.getRun.mockResolvedValue({ run_id: 42, state: { life_cycle_state: "RUNNING" }, }); @@ -829,21 +834,21 @@ describe("JobsPlugin", () => { process.env.DATABRICKS_JOB_ETL = "100"; process.env.DATABRICKS_JOB_ML = "200"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 1 }); + jobsApi.runNow.mockResolvedValue({ run_id: 1 }); const plugin = new JobsPlugin({}); const exported = plugin.exports(); await exported("etl").runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 100 }), expect.anything(), ); - mockClient.jobs.runNow.mockClear(); + jobsApi.runNow.mockClear(); await exported("ml").runNow(); - expect(mockClient.jobs.runNow).toHaveBeenCalledWith( + expect(jobsApi.runNow).toHaveBeenCalledWith( expect.objectContaining({ job_id: 200 }), expect.anything(), ); @@ -1081,7 +1086,7 @@ describe("injectRoutes", () => { test("returns runId on successful non-streaming run", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1203,7 +1208,7 @@ describe("injectRoutes", () => { { run_id: 1, state: { life_cycle_state: "TERMINATED" } }, { run_id: 2, state: { life_cycle_state: "RUNNING" } }, ]; - mockClient.jobs.listRuns.mockReturnValue( + jobsApi.listRuns.mockReturnValue( (async function* () { for (const run of mockRuns) yield run; })(), @@ -1241,7 +1246,7 @@ describe("injectRoutes", () => { test("passes limit query param to listRuns", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockReturnValue((async function* () {})()); + jobsApi.listRuns.mockReturnValue((async function* () {})()); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1268,7 +1273,7 @@ describe("injectRoutes", () => { await handler(mockReq, mockRes); // Verify the connector was called with limit 5 - expect(mockClient.jobs.listRuns).toHaveBeenCalledWith( + expect(jobsApi.listRuns).toHaveBeenCalledWith( expect.objectContaining({ limit: 5 }), expect.anything(), ); @@ -1284,7 +1289,7 @@ describe("injectRoutes", () => { job_id: 123, state: { life_cycle_state: "TERMINATED" }, }; - mockClient.jobs.getRun.mockResolvedValue(mockRun); + jobsApi.getRun.mockResolvedValue(mockRun); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1351,7 +1356,7 @@ describe("injectRoutes", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Run exists upstream but is owned by job 456, not the configured 123. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1393,7 +1398,7 @@ describe("injectRoutes", () => { run_id: 42, state: { life_cycle_state: "TERMINATED" }, }; - mockClient.jobs.listRuns.mockReturnValue( + jobsApi.listRuns.mockReturnValue( (async function* () { yield mockRun; })(), @@ -1432,7 +1437,7 @@ describe("injectRoutes", () => { test("returns null status when no runs exist", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.listRuns.mockReturnValue((async function* () {})()); + jobsApi.listRuns.mockReturnValue((async function* () {})()); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1469,8 +1474,8 @@ describe("injectRoutes", () => { test("cancels run and returns 204", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1540,8 +1545,8 @@ describe("injectRoutes", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Pre-flight getRun reports a run owned by a different job. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); - mockClient.jobs.cancelRun.mockResolvedValue(undefined); + jobsApi.getRun.mockResolvedValue({ run_id: 99, job_id: 456 }); + jobsApi.cancelRun.mockResolvedValue(undefined); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1570,7 +1575,7 @@ describe("injectRoutes", () => { expect(mockRes.status).toHaveBeenCalledWith(404); // Must not fall through to the cancel call or the 204. - expect(mockClient.jobs.cancelRun).not.toHaveBeenCalled(); + expect(jobsApi.cancelRun).not.toHaveBeenCalled(); expect(mockRes.end).not.toHaveBeenCalled(); }); @@ -1722,7 +1727,7 @@ describe("injectRoutes", () => { test("allows exactly MAX_UNVALIDATED_PARAM_KEYS (50) keys without schema", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({ jobs: { etl: { taskType: "notebook" } }, @@ -1760,13 +1765,13 @@ describe("injectRoutes", () => { // 50 keys is under the cap — request proceeds to the SDK. expect(mockRes.json).toHaveBeenCalledWith({ runId: 42 }); - expect(mockClient.jobs.runNow).toHaveBeenCalled(); + expect(jobsApi.runNow).toHaveBeenCalled(); }); test("allows undefined params", async () => { process.env.DATABRICKS_JOB_ETL = "123"; - mockClient.jobs.runNow.mockResolvedValue({ run_id: 42 }); + jobsApi.runNow.mockResolvedValue({ run_id: 42 }); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1807,7 +1812,7 @@ describe("injectRoutes", () => { const error = new Error("Sensitive internal detail: token expired"); (error as any).statusCode = 403; - mockClient.jobs.runNow.mockRejectedValue(error); + jobsApi.runNow.mockRejectedValue(error); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); @@ -1849,7 +1854,7 @@ describe("injectRoutes", () => { const error = new Error("Unauthorized"); (error as any).statusCode = 401; - mockClient.jobs.listRuns.mockImplementation(() => { + jobsApi.listRuns.mockImplementation(() => { throw error; }); @@ -1884,10 +1889,10 @@ describe("injectRoutes", () => { process.env.DATABRICKS_JOB_ETL = "123"; // Pre-flight succeeds so we reach the actual cancel call. - mockClient.jobs.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); + jobsApi.getRun.mockResolvedValue({ run_id: 42, job_id: 123 }); const error = new Error("Forbidden"); (error as any).statusCode = 403; - mockClient.jobs.cancelRun.mockRejectedValue(error); + jobsApi.cancelRun.mockRejectedValue(error); const plugin = new JobsPlugin({}); const routeSpy = vi.spyOn(plugin as any, "route"); diff --git a/packages/appkit/src/plugins/server/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/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..13f79b5d7 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,12 @@ 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 { 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 +125,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", @@ -332,28 +388,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 +430,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, @@ -534,39 +563,3 @@ export function createFailedSQLResponse(errorMessage: string) { statement_id: `stmt-${Date.now()}`, }; } - -/** - * A WorkspaceClient whose `executeStatement`/`getStatement` are bare `vi.fn()`s - * (no default resolution) so a test can script exactly what SQL returns. - * `warehouses.get` defaults to RUNNING. - */ -export function createConfigurableMockWorkspaceClient() { - const executeStatement = vi.fn(); - const getStatement = vi.fn(); - // Analytics route now calls `warehouses.get` before issuing SQL; default to - // RUNNING so callers that don't care about warehouse readiness don't have - // to wire it up. - const warehousesGet = vi.fn().mockResolvedValue({ state: "RUNNING" }); - const warehousesStart = vi.fn().mockResolvedValue(undefined); - - const client = { - statementExecution: { - executeStatement, - getStatement, - }, - warehouses: { - get: warehousesGet, - start: warehousesStart, - }, - }; - - return { - client, - mocks: { - executeStatement, - getStatement, - warehousesGet, - warehousesStart, - }, - }; -} diff --git a/packages/appkit/src/testing/index.ts b/packages/appkit/src/testing/index.ts index 8565e4d5e..797ce6935 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,11 @@ export { type StreamSource, } from "./expect-stream"; export { - createConfigurableMockWorkspaceClient, createFailedSQLResponse, createMockRequest, createMockResponse, createMockRouter, createMockTelemetry, - createMockWorkspaceClient, createSuccessfulSQLResponse, mockServiceContext, type OboOption, @@ -68,6 +75,13 @@ export { type TestContextOptions, useServiceContextMock, } from "./fixtures"; +export { + createMockWorkspaceClient, + type CreateMockWorkspaceClientOptions, + getMock, + type MockWorkspaceClient, +} from "./mock-workspace-client"; +export { createTestPlugin } from "./create-test-plugin"; export { createTestPluginContext, type FakeProvider, 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..7b9298379 --- /dev/null +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -0,0 +1,319 @@ +/** + * 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.get` and `warehouses.start` + * 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.get": { state: "RUNNING" }, + "warehouses.start": 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/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..599fb8151 100644 --- a/packages/appkit/src/testing/tests/fixtures.test.ts +++ b/packages/appkit/src/testing/tests/fixtures.test.ts @@ -3,8 +3,10 @@ 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 { createMockRequest, + mockServiceContext, resetTestCache, useServiceContextMock, } from "../fixtures"; @@ -102,6 +104,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" }); 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..a2cbeda1a --- /dev/null +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -0,0 +1,307 @@ +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", "get", { state: "RUNNING" }], + ["warehouses", "start", 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.get({} as never)).resolves.toEqual({ + state: "RUNNING", + }); + await expect( + client.warehouses.start({} 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..536099425 --- /dev/null +++ b/packages/appkit/src/testing/tests/published-surface.integration.test.ts @@ -0,0 +1,153 @@ +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", + "parseSSEResponse", + "resetTestCache", + "runWithRequestContext", + "setupDatabricksEnv", + "useServiceContextMock", + ]; + 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-plugin-context.test.ts b/packages/appkit/src/testing/tests/test-plugin-context.test.ts index 60b30b917..d84a52937 100644 --- a/packages/appkit/src/testing/tests/test-plugin-context.test.ts +++ b/packages/appkit/src/testing/tests/test-plugin-context.test.ts @@ -4,6 +4,7 @@ import { describe, expect, test } 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", () => { 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/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/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, From 72fb58a7038d93a86e91bbe2d58f1f65f238113a Mon Sep 17 00:00:00 2001 From: IamGalymzhan <62868459+IamGalymzhan@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:01:10 +0500 Subject: [PATCH 02/12] feat(appkit): add testing-kit helpers for env, errors, context, and cache (#555) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(appkit): make PluginContext telemetry injectable The testing kit needs to construct a real PluginContext without a live OpenTelemetry pipeline. Add an optional constructor dependency for the telemetry provider, defaulting to the shared "plugin-context" provider so the production path is unchanged. This is the single production edit required to wrap the real class in tests rather than reimplementing it. Signed-off-by: Galymzhan * feat(appkit): ship @databricks/appkit/testing and migrate first stub Wire the testing kit as a published subpath and prove it against the first of the two hand-rolled context stubs (the design gate): - Add ./testing to both exports maps (dev + publishConfig) following the ./type-generator shape, add src/testing/index.ts to the tsdown entry, and declare vitest as an optional peerDependency. Build passes attw + publint; dist/testing/{index,mock-plugin-context,expect-stream,fixtures}.{js,d.ts} are emitted and vitest stays external to the main entry. - Migrate dispatch-tool-call.test.ts: replace (plugin as any).context = { executeTool } with mockPluginContext. executeTool is now the REAL method, so the forwarded toolCallTimeoutMs is asserted through actual signal composition, the on-behalf-of (asUser) path is verified, and a new test proves the forwarded timeout actually aborts a slow toolkit tool end-to-end. This is the primary win from the plan: executeTool's OBO and timeout paths gain real assertions instead of a stub that proved nothing. Signed-off-by: Galymzhan * test(appkit): migrate route-handler-errors context stub to mockPluginContext Replace the second and final hand-rolled stub — (plugin as any).context = { addRoute } — with the real PluginContext from mockPluginContext. The kit's route recorder captures raw handlers, so the alias assertion (both /invocations and /responses mount the same handler reference) holds against the real class, where forwardAsyncErrors wrapping would otherwise break reference identity. Both context stubs the plan identified are now migrated. Signed-off-by: Galymzhan * docs(appkit): document the testing kit and ship a template example test - Add docs/docs/development/testing.md covering mockPluginContext(), expectStream(), and the fixture helpers, with a full end-to-end example. Cross-links to local-development, custom-plugins, and execution-context. - Add template/server/example.test.ts: a self-contained, plugin-agnostic example that scaffolded apps ship with — it defines a tiny custom plugin and exercises both mockPluginContext (route recording) and expectStream (ordered event assertions), running with no workspace or network. Ships the kit to users, satisfying the plan's acceptance criteria that a docs page exists and the template carries at least one example test. Signed-off-by: Galymzhan * docs(appkit): fix testing-kit examples to instantiate the plugin class Validation by scaffolding a real app with `databricks apps init` surfaced that the examples called the `analytics()`/`toPlugin()` factory and then treated the result as a plugin instance — but a factory returns a { plugin, config, name } descriptor for createApp to construct, so `.attachContext`/handler methods are absent. Rewrite both the template example test and the docs "Full example" to instantiate the plugin class directly (`new GreeterPlugin({})`), matching how the migrated agents suites use the kit. The scaffolded app's `npm test` and `tsc` both pass against the published `@databricks/appkit/testing` subpath with no workspace or network. Signed-off-by: Galymzhan * refactor(appkit): tighten FakeToolResponse so a missing value is a type error Drop `undefined` from the static FakeToolValue union. `resolve()` treats an undefined map entry as "unregistered tool" and throws, so allowing undefined as a declared response made `{ query: undefined }` a confusing runtime error instead of a compile error. A function returning undefined still works for the rare "returns nothing" case. Add a test pinning that a null response is returned as a value, not misread as a missing tool. Signed-off-by: Galymzhan * refactor(appkit): make tools/test-helpers a shim over the shipped testing kit The plan's step 5 was to MOVE the fixtures into the package, not copy them. The shipped kit (src/testing/fixtures.ts) duplicated all 15 exports of tools/test-helpers.ts, which would drift over time. Collapse the original into a thin re-export of @databricks/appkit/testing so src/testing is the single source of truth while the 18 existing @tools/test-helpers importers keep working unchanged. The re-exported mockServiceContext is now synchronous; every call site either awaits it (no-op on a non-promise) or reads it through Awaited>, so all suites pass unchanged (full appkit suite: 3117 passed, 1 pre-existing skip). Signed-off-by: Galymzhan * fix(appkit): normalize CRLF in expectStream SSE parsing; sharpen testing docs Code review follow-ups: - expectStream's parseSSEBody split frames on \n\n, so a spec-compliant SSE stream delimited by \r\n\r\n (from a real server) collapsed into one event. AppKit's own writer uses \n\n so existing tests were unaffected, but expectStream is public API that accepts any Response. Normalize CRLF to LF before splitting; add a CRLF regression test. - Docs: instantiate the plugin CLASS in the attach() snippet (the factory returns a descriptor, not an instance), and note that the cache attach() seeds is a per-process singleton shared by tests within a file. Signed-off-by: Galymzhan * fix(appkit): resolve repo-wide Biome error blocking CI CI's "Lint & Type Check" job runs `pnpm run check` over the whole repo, so a pre-existing lint error unrelated to this branch failed the build: - remote-tunnel-controller.test.ts had two `afterEach` hooks in one describe (lint/suspicious/noDuplicateTestHooks, error severity). Merge them into one — behavior preserved (env reset + console-spy clear both still run after each test). This file is byte-identical to main; the error predated the branch and only surfaced because CI lints the entire tree. Also drop two dead `biome-ignore lint/suspicious/noExplicitAny` suppressions in the testing kit (fixtures.ts, expect-stream.test.ts): `noExplicitAny` is turned off repo-wide in biome.json, so the comments had no effect (suppressions/unused warnings). The invalid-source test now casts through `unknown as never`. Signed-off-by: Galymzhan * fix(appkit): address cross-model review findings in the testing kit Verified and fixed the findings from an independent code review: - #1 (correctness) expectStream dropped the wire `event:` name when the JSON payload carried its own `type` (spread ran after the assignment). Spread the payload first, then set `type = name ?? parsed.type`, so a frame like `event: error` + `data: {"type":"result"}` reports `error`. Regression test added. - #2 (contract) `@databricks/appkit/testing` eagerly loads vitest via fixtures even for `expectStream`, so vitest is a real requirement. Drop the "optional" peerDependenciesMeta and correct the docs sentence. - #6 (OBO fidelity) the fake `asUser` recorded `asUser: true` unconditionally. Enforce the real `Plugin.asUser` token precondition: a request without `x-forwarded-access-token` throws `missingToken` (missing user id throws too), and the resolved `userId` is recorded on each tool call. Tests now assert both directions (well-formed request vs token-less). - #3 (fidelity) attach() now mirrors AppKit core: registerPlugin plus registerToolProvider for real tool providers, without clobbering injected fakes. getPlugins()/getPluginNames()/hasPlugin() behave as in production. - #12 unknown-tool lookup used `tools[name] === undefined`, so a tool named "constructor"/"toString" hit Object.prototype. Use Object.hasOwn. - #5 drop data-less named SSE frames (real clients ignore them). - #7 re-export the PluginContext type from the testing barrel so MockPluginContext.ctx is nameable through the exports map. - #13 correct the docs: mock.telemetry captures the context's executeTool spans, not plugin-level spans (attachContext rebuilds the plugin's own telemetry). - #4 parseSSEResponse now delegates to the same parseSSEBody as expectStream — one parser, no divergence. All 3 analytics.integration call sites still pass. - #8 reformat template/server/example.test.ts with the template's Prettier so a scaffolded app's `npm run format` passes. - #10 fix the package-doc @example (agentsPlugin._handleStream does not exist). - #11 add kit tests that exercise attach() end-to-end (cache seed, isReady, registration, fake-not-clobbered). Build passes attw + publint; full appkit suite 3125 passed / 1 pre-existing skip. Signed-off-by: Galymzhan * chore(appkit): drop knip vitest-ignore now that vitest is a real peer dep With vitest declared as a (non-optional) peerDependency, knip recognizes it as used, so the earlier ignoreDependencies entry is unnecessary. This reverts knip.json to its original state. Signed-off-by: Galymzhan * fix(appkit): make vitest a normal dependency, not a package-wide peer A required peerDependency has no per-subpath scope: it applied to the whole @databricks/appkit package, so every production consumer that never imports the testing kit got an unsatisfied peer (npm 7+ auto-installs vitest into their tree; pnpm warns) — a wider blast radius than the eager-import bug it was meant to fix. Follow appkit's own precedent instead: `vite` backs the ./type-generator subpath as a normal `dependency`, installed for everyone but loaded only by importers of that subpath. Do the same for `vitest` and ./testing. vitest is referenced solely by dist/testing/fixtures.js, never by the main/plugin/core entry, so a consumer importing createApp never loads it. Verified end-to-end: scaffolded an app whose own vitest (4.1.9) differs in major from appkit's dependency (3.2.4), forcing a nested second copy. The testing kit's vi.fn()/vi.spyOn() mocks and expect(...).toHaveBeenCalled() assertions work across the two instances (vi spies carry their own call state), and npm install emits no peer-dep warning. Build passes attw + publint. Also fold in the template example's Prettier formatting (template uses Prettier, not Biome) so a scaffolded app's `npm run format` passes. Signed-off-by: Galymzhan * refactor(appkit): rename mockPluginContext to createTestPluginContext The helper builds the REAL PluginContext with faked edges — it does not mock the context — so the name was misleading. Rename to createTestPluginContext (and the MockPluginContext type to TestPluginContext), matching the create*-for-tests convention, and rename the files to test-plugin-context.ts. Pre-merge and unreleased, so no external consumers are affected. Also finish the #13 doc-accuracy fix in the shipped JSDoc (not just the docs page): the telemetry field comment now states it captures the context's spans (executeTool), not plugin-internal spans — attachContext rebuilds the plugin's this.telemetry from the real TelemetryManager. These comments ship in dist/testing/*.d.ts, so IntelliSense previously showed the unqualified claim. Build passes attw + publint; full appkit suite 3125 passed / 1 pre-existing skip. Signed-off-by: Galymzhan * refactor(appkit): dedupe testing fixtures and tidy test-plugin-context Behavior-preserving cleanups in the testing kit: - createMockRequest reuses createMockWorkspaceClient() instead of an inline copy of the same mock client (verified identical). - createMockServiceContext / createMockUserContext / mockServiceContext inline the createMockWorkspaceClient() call into the `||` fallback, so the mock client is built only when the caller did not supply one. - The fake asUser view spreads `...base` and overrides executeAgentTool rather than re-declaring getAgentTools. - expectStream's isSubsequence breaks once the expected sequence is fully matched. No semantic change; typecheck clean and all kit + migrated tests pass. Signed-off-by: Galymzhan * fix(appkit): resolve third-review findings in the testing kit - #1 (P1) The docs called vitest a peer dependency, but the manifest ships it under `dependencies` (the decision we landed on, matching how appkit ships `vite` for ./type-generator). Correct the docs to match: appkit installs vitest for you, and it loads only when you import ./testing. Manifest and docs now agree. - #2 (P2) expectStream buffered the source eagerly with no bound, so a non-terminating stream hung until the runner's own timeout. Add an optional `{ timeout }` that fails fast with a clear, kit-specific error; document it and cover both directions with tests. - #3 (P2) The fake asUser replicates asUser's token precondition but not the real dev-mode `DEV_OBO_FALLBACK_KEY` OTel marker (a module-private telemetry detail). Narrow the docs and JSDoc to say so and point users at the recorded asUser/userId fields instead of isDevOboFallback(). Build passes attw + publint; full appkit suite 3141 passed / 1 pre-existing skip. Signed-off-by: Galymzhan * test(appkit): dogfood the testing kit on analytics and genie plugins Exercise @databricks/appkit/testing against real core plugins to validate it beyond the two agent proof sites and produce usage references: - analytics.kit.test.ts: cross-plugin executeTool via createTestPluginContext — OBO identity (asUser/userId), token-precondition rejection, and per-call timeout abort. Needs only the kit (no workspace/ServiceContext). - genie.kit.test.ts: drives the real _handleSendMessage SSE stream and asserts event order with expectStream(...).toEmit(...). Both add genuinely new coverage (streamed SSE order + OBO dispatch identity were untested). Full appkit suite 3145 passed / 1 pre-existing skip. Developer-experience notes (kit wins + friction, e.g. createMockResponse doesn't compose with expectStream) captured in internal/ for the milestone review. Signed-off-by: Galymzhan * refactor(appkit): address testing-kit review feedback Resolve the eight review comments on the testing kit: - createMockResponse now captures written SSE bytes and exposes sseResponse(); expectStream reads a captured mock response directly, so streaming-route tests no longer need a hand-rolled bridge. - Ship vitest as an optional peer dependency (+ devDependency) instead of a plain runtime dependency, keeping the test framework out of production installs and deduping to the app's own copy. Ignore it in knip. - Add an obo option to createMockRequest so on-behalf-of tests set the forwarded identity headers with one flag. - Add resetTestCache() to clear the shared cache singleton between tests. - Use the documented attach() instead of an any-cast in the agents dispatch tests. - Drop the unused createMockServiceContext/createMockUserContext builders from the public surface; keep the service-context builder internal. - Pin the previously untested edges: the Object.hasOwn tool-lookup guard, the dev-mode asUser branch, and parseSSEBody's non-object data values. - Add useServiceContextMock() to register the mock lifecycle in one line, returning a live accessor. Dogfood the new helpers in the analytics, genie, and serving suites, and document them in the testing guide. Signed-off-by: Galymzhan * docs(appkit): move the testing guide under Plugins The testing kit is entirely plugin-scoped (createTestPluginContext, attach(plugin), plugin route/tool/SSE assertions), and the page's own cross-links already pointed into plugins/. Move it next to custom-plugins and fix the relative links. Keep the heading as 'Testing'; the Plugins section supplies the context. Signed-off-by: Galymzhan * test(appkit): fold dogfood tests into plugin suites Address round-2 review: the kit should be the default way to test a plugin, not a parallel '*.kit.test.ts' track. - Fold the three cross-plugin executeTool OBO tests into analytics.test.ts and delete analytics.kit.test.ts. - Upgrade genie.test.ts's SSE test to assert event ORDER via expectStream on genie's real event names (message_start, status, message_result, query_result), replacing brittle write.mock.calls substring checks, and delete genie.kit.test.ts. - Trim the heavy comment narration from the folded-in tests. - Re-export createTestPluginContext and expectStream from the test-helpers shim. - Finish the testing-guide move under plugins/ (sidebar position + links). Signed-off-by: Galymzhan * test(appkit): restore toHaveLength(1) on the analytics OBO dispatch test The dogfood fold trimmed expect(mock.toolCalls).toHaveLength(1), so a double-dispatch would no longer fail the happy-path test — and it was inconsistent with the token-less sibling that kept toHaveLength(0). Restore it. Signed-off-by: Galymzhan * test(appkit): re-assert genie SSE payloads after the expectStream swap The toEmit swap pinned event order but dropped the payload values the old substring checks covered (conversationId=new-conv-id, status=ASKING_AI), which aren't asserted elsewhere. Restore them structurally via collect() + toMatchObject — keeping the ordering guarantee without brittle substrings. Signed-off-by: Galymzhan * fix(appkit): drop fabricated workspace-client fields from createMockRequest createMockRequest returned userWorkspaceClient, serviceWorkspaceClient, getWarehouseId and getWorkspaceId — fields no production code reads (plugins resolve those through getWorkspaceClient()/getWarehouseId() from src/context, which mockServiceContext stands in for). Publishing them via @databricks/appkit/testing would make four inert fields a permanent public promise. The two warehouse cold-start tests (analytics + metric) overrode mockReq.serviceWorkspaceClient.warehouses.get, which the route never reads — so they passed on the default RUNNING client without exercising the warehouse path at all. Route the warehouse client through mockServiceContext (the real seam) so the tests are live, and drop the 'mock WorkspaceClient' claim from the testing guide. Signed-off-by: Galymzhan * feat(appkit): add a never-crash mock WorkspaceClient to the testing kit Every core plugin's actual work runs through getWorkspaceClient(), which the testing kit did not fake — so a jobs/genie/serving/files plugin crashed on its first client call and authors hand-rolled nested client literals instead. createMockWorkspaceClient() fakes the whole facade in three layers: - The 9 facade members are explicitly typed, so `client.jbos` is a compile error. The facade is closed and AppKit-owned, so there is no per-service fixture to maintain as the SDK grows. - Each service is a Proxy minting one memoized vi.fn() per method name, keyed by dotted path. `client.jobs.getRun === client.jobs.getRun`, so call assertions work, and the legacy view shares the map so one `responses` entry covers both — including un-faceted services like `legacy.clusters.list()`. - `config` and `apiClient` are seeded objects rather than bare Proxies, because three of their members must not be mocks: `config.host` is a real string that production code builds URLs from and throws on when falsy, `apiClient.userAgent()` must be synchronous (a Promise inside a Headers value stringifies to "[object Promise]"), and `apiClient.request` resolves {} so destructuring its result does not throw. Two guards keep the Proxy safe. Symbol keys delegate to Reflect.get, and a passthrough deny-set answers `undefined`. `then` is the load-bearing entry: without it a service looks thenable, so `await client.jobs` either hangs or resolves to a mock's return value. ownKeys is left at its default so util.inspect and toEqual see {} instead of recursing forever. The three historical canned defaults are byte-identical, because 13 test files reach them implicitly through mockServiceContext. `currentUser.me` is additive and load-bearing: ServiceContext.createContext reads `currentUser.id`, so an unresolved me() is a TypeError and createApp({ client }) cannot boot without it. getMockFn(client, "jobs.getRun") is the typed assertion path — facade accessors are legacy-SDK-typed, so expect(client.jobs.getRun).toHaveBeenCalled() does not typecheck. It mints idempotently, so the handle can be grabbed before the code under test runs. The compile-time block is enforced by tsc, not at runtime. It records one correction to the plan: the SDK types `config.host` as `string | undefined`, so the contract is that it narrows to a string, not that it is non-optional. 4451 tests pass (+38); the 667 tests reaching the default client indirectly through mockServiceContext are unchanged. Co-authored-by: Isaac Signed-off-by: Galymzhan * refactor(appkit): converge the two mock-workspace-client builders fixtures.ts had its own two-service createMockWorkspaceClient, so the shipped fixture and the new never-crash builder were near-duplicates. The fixture now re-exports the builder and the barrel points at its new home. The blast radius is entirely indirect. Nothing in src imports the exported fixture by name (connectors/genie/tests/client.test.ts defines its own local one), but buildServiceContextState calls it as the default client for mockServiceContext, which 13 test files use. The risk therefore lives in the default return value, which is why U1 kept the three canned defaults byte-identical — and why this commit adds the convergence guard that asserts both halves: jobs/genie now resolve instead of throwing "Cannot read properties of undefined", while the SQL path those 13 files depend on still succeeds. createConfigurableMockWorkspaceClient is left byte-for-byte unchanged and only gains a @deprecated notice. Its bare vi.fn()s return undefined *synchronously* whereas the new floor returns Promise, and its one caller (analytics.integration.test.ts) can observe that difference; reimplementing it here would change behaviour for no benefit. It migrates with that suite later. The jobs suite drops its hand-rolled client literal — the seven method mocks plus the config.host/authenticate block — onto the builder, which is the proof the boilerplate actually goes away. Its 57 assertion sites move to a getMockFn handle because facade accessors are legacy-SDK-typed, so .mockResolvedValue on them does not typecheck. The factory needs `await vi.hoisted(async ...)` with a dynamic import, since a hoisted factory runs before the file's imports. 4454 tests pass (+3). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): split the lifecycle exit from the teardown LifecycleManager's shutdown sequence was reachable only by killing the process, so nothing could release AppKit's sockets, timers, pools, cache, and telemetry and keep running. That is what blocks an app handle's close(), and with it any test that wants to boot more than once in a file. The sequence is now a phase runner that returns an exit code, a promise memo, and two thin callers: - shutdown() is the signal path, observably unchanged: it arms the same unref'd 15s force-exit backstop and still exits 0 on completion, 1 on an unexpected throw. The timer stays here deliberately — it is the one thing close() must not inherit, since a programmatic caller wants a logged error when teardown hangs, not a dead process. - close() is the programmatic path: it detaches signal handlers, runs the same phases under a shorter default budget (5s, not the production 15s), logs the phase that was in flight if the budget is spent, and never exits. Replacing the isShuttingDown boolean with a promise memo is a strict improvement. The boolean made a second caller return *immediately* while teardown was still running — harmless for a signal, since the first caller exits the process anyway, but for close() it would resolve before resources were released, which is the difference between a correct handle and a misleading one. The read and the assignment stay in one synchronous statement, preserving the invariant the boolean was there to protect. One production behaviour does shift: a second signal now awaits the first teardown. installSignalHandlers registered anonymous arrows that could never be removed. The [signal, handler] pairs are now retained and detached individually, never via removeAllListeners, so a host embedding AppKit keeps its own handlers. The tests assert that with two managers installed, a.close() leaves b's pair and an unrelated host listener intact, and that counts return to their pre-install baseline — which is what stops repeated boots tripping MaxListenersExceededWarning. The signal-mid-close race is documented rather than papered over: handlers come off before the first await, and if a signal still lands it joins the memo and exits, because it wanted the process dead. The idempotency test is verified by injection — it fails against the old return-immediately semantics and passes against the memo. Its first draft did not: it counted microtask ticks, which cannot distinguish an early return through close()'s raceWithTimeout wrapper. It now asserts that neither caller settles until the plugin hook has actually completed. 4463 tests pass (+9); the 14 pre-existing shutdown tests are untouched. Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): expose close() on the app handle createApp acquired sockets, timers, and pools but returned no way to release them, so the only teardown was killing the process. The LifecycleManager built at the end of _createApp was constructed and immediately discarded; it is now retained on the instance and reachable through the handle. The return type widens from PluginMap to AppHandle, which is PluginMap plus close() and Symbol.asyncDispose. Widening a return type is source-compatible for every existing caller, and the cast that produces the handle already hid instance methods, so close() rides along naturally. onPluginsReady deliberately keeps PluginMap: it runs before the server starts, so handing it a close() would invite a footgun for no gain. The name collision is a real hazard, not a theoretical one. Plugin exports are installed with Object.defineProperty, and an own property shadows a prototype method — so a plugin named `close` would silently replace teardown rather than merely confuse the types. Three layers guard it: Symbol.asyncDispose is unreachable from a manifest name, so `await using` is always safe; createAndRegisterPlugin now throws a ConfigurationError naming the offending plugin; and no plugin in the repo is affected. Coverage is deliberately unmocked, because the claim is about real resources: a boot on an ephemeral port serves /health, close() runs the plugin's shutdown hook, the socket stops accepting, and the SIGTERM listener count returns to its pre-boot baseline. Also covered: idempotency at the app level, a server-less app closing cleanly, `await using` releasing at scope exit, and the reserved name being rejected. Verified by injection — with close() stubbed to a no-op and the reserved-name guard removed, 5 of the 6 fail. 4469 tests pass (+6). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): make the process-wide singletons re-bootable close() released resources but left the singletons pointing at them, so close() followed by createApp() silently reused what the teardown had just torn down. This delivers the actual driver — boot, assert, close, repeat. CacheManager.reset() drops both `instance` and `initPromise`. Clearing only `instance` is insufficient because getInstance() returns `initPromise` when `instance` is null, so the next boot would await a promise resolving to the dead manager. Testing surfaced a third case the plan missed: clearing both is *still* not enough, because an initialization already in flight runs its continuation and re-publishes the very instance being discarded. A generation counter now invalidates that write. The covering test models PersistentStorage rather than using the default in-memory storage. This matters: InMemoryStorage.close() only clears a Map and stays usable, so an in-memory test passes whether or not the reset exists — which is precisely why the bug hid. Against storage whose close() is terminal, the way pool.end() is, the test shows the stale manager throwing "Cannot use a pool after calling end()" and the reset fixing it. One plan claim is corrected rather than implemented. The plan asserted that TelemetryManager's never-cleared `shutdownPromise` made a second shutdown() return a stale promise and skip flushing a re-initialized SDK. It does not: shutdown() only returns the memo after reassigning it for whatever SDK is currently live, so a stale resolved promise can be returned only when there is no SDK to flush. Verified twice — by mocking NodeSDK across three initialize/shutdown cycles, and by running the original implementation in isolation. An earlier draft of this commit added a generation counter here too; it has been reverted, since it fixed nothing and cost a field. What TelemetryManager did need, and now has, is the static reset() that drops the singleton. The resets are wired into close() only, never the signal path, where the process is dying and pointer drops are pure cost. Symmetry is the justification: core initializes all four in _createApp, so core drops all four. This is a semantic expansion, not purely a bug fix — a host that closes and then expects ServiceContext.get() to work will now get an InitializationError. resetAppKitSingletons() is published from @databricks/appkit/testing for tests that hand-roll createApp and would otherwise deep-import ../context/service-context to reach ServiceContext.reset(). Both it and LifecycleManager.close() delegate to one core-side implementation rather than duplicating the list. resetTestCache() is untouched — it calls clear() on the existing cache, a different and still-useful operation. 4480 tests pass (+11). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): add createTestApp — the customer-grade plugin test harness One call boots a real AppKit app with no workspace, no credentials, and no network, and calls it over real HTTP: const app = await createTestApp({ plugins: [myPlugin()] }); const res = await app.post("/api/my-plugin/thing", { body, obo: true }); await expectStream(res).toEmit("status", "result"); await app.close(); Four of the setup steps exist only because of hazards found by reading the boot path, and each has a test that fails without it: - NODE_ENV is pinned away from "development". Not tidiness: dev mode routes the injected `port: 0` through get-port, where portNumbers(0, …) throws a RangeError. "development" is refused outright with an explanation rather than worked around, since dev mode also boots a real Vite server, downgrades resource validation to a warning, and stops filtering dev-only plugins. - DATABRICKS_WORKSPACE_ID is set, short-circuiting the SCIM probe in getWorkspaceId, and internal telemetry is disabled. Both would otherwise fire apiClient.request during boot. A canary test asserts zero calls after boot, so either regression fails loudly. - The cache gets explicit in-memory storage. Without it CacheManager builds its own workspace client — ignoring the injected one — and probes Lakebase over the network, so "no network" would be false. - The server plugin is reached through a lazy `await import()`, because it runs dotenv.config() at module load. A static import would mutate a consumer's process.env merely by importing the testing entry point. process.env is snapshotted wholesale rather than by whitelist, since plugins read vars the harness cannot enumerate, and restored on close() — including deleting keys the harness added and restoring a pre-existing DATABRICKS_HOST to its own value rather than the test default. Teardown also runs from the boot-failure path, or a plugin whose setup() throws would leak env mutations into every later test in the file. Plugin exports live under app.plugins rather than spread onto the handle: `get` and `delete` are plausible plugin names and would collide with the request methods. The request methods return a native Response, so expectStream composes with no bridge — the dogfooding report's top friction, avoided by construction. `obo` reuses createMockRequest's OboOption rather than inventing a second convention. Two corrections to the plan, both found by testing: - A `strictValidation: false` opt-out was specified and has been dropped as a false affordance. enforceValidation computes `shouldThrow = !isDevelopment || strict`, so with NODE_ENV pinned away from "development" validation always throws and the flag cannot do anything. The env var is still set as belt-and-braces, and a test pins the unconditional behaviour. - The error-middleware test initially asserted a redacted body. It is not redacted: errorHandlerMiddleware hides the message only under NODE_ENV=production, and the harness pins "test". Useful for tests — an assertion can name the failure — but it means that response is the dev shape, which the test now says out loud. The HTTP suite's probe plugin registers routes through `this.route()`, the way real plugins do. Registered with raw `router.get()` a rejection escapes forwardAsyncErrors and hangs the request — correct AppKit behaviour, and worth having a representative test rather than a misleading one. 4511 tests pass (+31). Co-authored-by: Isaac Signed-off-by: Galymzhan * feat(appkit): publish the harness surface, add createTestPlugin, and dogfood both Audit first: the ./testing subpath was already in both exports maps and the tsdown entry, vitest is already an optional peer dep, and attw/publint pass. The real gap was proof that a test needs nothing else, so the integration suites moved onto the public entry point — that migration *is* the audit. A new acceptance suite imports only from @databricks/appkit/testing and boots, requests, asserts a stream, and closes. Self-referencing the package from inside it needed a tsconfig paths entry. Resolving the package's own export map made tsc's project root ambiguous (TS2209), and the alias mirrors how shared and @databricks/lakebase are already mapped. It also makes source resolution deterministic rather than depending on the "development" export condition — verified by marker: the subpath resolves to src, not dist. createTestPlugin(factory, config) closes the last dogfooding footgun. Reaching through a descriptor with `new (genie({}).plugin)(config)` skips DEFAULT_CONFIG and forgets `name`, so the instance under test is configured differently from the one production builds. It mirrors createAndRegisterPlugin's merge order. createTestApp does not subsume it: the harness takes descriptors and builds instances itself, so the unit path needs its own ergonomics. Dogfooding results, reported as measured rather than as hoped: - analytics.integration.test.ts: 300 -> 216 lines. Setup/teardown went 104 -> 55, against the plan's predicted ~30. Its local getListeningPort helper is gone and its 12 mock handles now come from getMockFn. Same 6 tests, same assertions. - getListeningPort is lifted into the kit, and files/plugin.integration.test.ts imports it instead of carrying its own copy. - server.integration.test.ts moves four of its five blocks to ephemeral ports. The fifth keeps its fixed port deliberately, because it asserts the server honours a configured one; a comment says so. The removed sleep-100ms waits are replaced by getListeningPort, which waits on the listening event instead of guessing. One plan claim corrected: the hardcoded TEST_PORT = 9879 said to collide with server.integration was already fixed on this branch — analytics had moved to port: 0. The real fixed ports were the five in server.integration itself, which is what this commit addresses instead. Docs lead with createTestApp: a which-harness comparison table, the dotted-path responses convention, the teardown contract, a "Mocking Databricks services" section carrying the never-crash floor's honest catch (a misspelled *method* returns undefined, and a Lakebase pool built on the fake cannot connect), and an explicit callout that manifest.config.schema is not validated. The PluginContext/ServiceContext boundary note now says the kit covers the data plane. The template example gains a createTestApp test. 4519 tests pass (+8). pnpm docs:build is clean. Co-authored-by: Isaac Signed-off-by: Galymzhan * test(appkit): tighten the mock-client type contract from tarball findings Verified the kit end to end the way the plan prescribes: pnpm pack:sdk, an app scaffolded by `databricks apps init` from this repo's template, the tarballs installed into it, and the suites run with no .env, no credentials, and every non-loopback socket connection hard-blocked. Nine customer-style tests plus the template's three pass, and the scaffolded app typechecks against the shipped .d.ts. That run corrected a claim this branch had been making. Both the plan's risk table and the docs said a misspelled *method* slips through the never-crash floor and only a misspelled *service* is caught. Not so: each facade accessor is typed against the SDK's own service class, so `client.jobs.getRunz` and `client.files.anything` are compile errors too. The compile-time block now asserts that for three services, and the docs say what the real gap is — a method that exists but has no declared response, or a call that bypasses the types with a cast. Also repointed one doc line that told readers to reach the client via `getWorkspaceClient()`. That is right inside this repo but wrong from the published entry, where the name currently resolves to Lakebase's unrelated `getWorkspaceClient(config)`. The docs now use `getExecutionContext().client`, which is exported and works. The mis-export itself is a main-entry defect, outside this branch's scope, and is left for a follow-up. 4519 tests pass. Build, docs:build, attw, and publint are clean, and the packed tarball carries dist/testing/*.js and .d.ts for every new module. Co-authored-by: Isaac Signed-off-by: Galymzhan * fix(appkit): address code-review findings in the testing kit Eight reviewers over the branch diff produced 24 findings; 14 were actionable. The two real defects, both in code this branch added: - **An orphaned teardown could tear down the *next* app's resources.** `close()` races the shutdown phases against a 5s budget, then resets the core singletons and resolves. The phases keep running. A plugin `shutdown()` hook slower than that budget but inside its own 10s per-plugin budget — which the files plugin's drain can be — left phase 5 re-reading the static slots, so it either skipped draining this app's cache pool or closed the *following* app's storage and shut down its OTEL SDK. The comment claiming the phases had "already closed the cache storage" was only true when teardown finished in budget. Phase 5 now uses the instances captured before the first await, and a test drives the exact 5s-to-10s window (verified by reintroducing the bug). - **`process.env` restore did not compose across overlapping boots.** Each app snapshotted independently, so a second boot captured the first's mutations and whichever closed last re-applied them, stranding the harness keys and the first app's `env` entries after both apps were gone. Confirmed by probe, in-repo and from the packed tarball. There is now one reference-counted baseline: the first live app anchors it, the last one to close restores it, and the outcome no longer depends on close order. Also fixed: `server: false` alongside a caller-supplied server plugin is now refused instead of half-honoured (the plugin still bound a socket while the handle denied one existed); `AppHandle.close()` declares the `{ timeoutMs }` the implementation accepts, so the harness no longer casts to reach it; the duplicate `listeningPort` helper in the close integration suite is gone in favour of the kit's (two reviewers flagged it); the analytics suite drops an `as never` that erased `app.plugins` typing; the `clientFns` WeakMap moved above its users; and a comment claiming "9 typed facade members" over a 7-element array is corrected. Two of my own tests were weak and are now stronger: the `authenticate` test wrapped its whole body in `if (mockFn)` and only asserted "was called" — it now asserts the Authorization header it claims to set — and a close-after-signal test proved ordering by counting microtask ticks, which cannot see through `raceWithTimeout`; it uses the same sentinel the sibling test does. A new compile-time assertion pins that `AppHandle` still satisfies a `PluginMap` annotation, so a regression in the widening can't pass silently. Documented rather than changed: a service's methods are callable but not enumerable, so `'getRun' in client.jobs` is false and `Object.keys` is empty. Reporting those keys would make `util.inspect` mint a mock per probe, which is the recursion the default traps exist to avoid. Also documented why `onPluginsReady` keeps the narrower `PluginMap`. One finding rejected as a false positive: project-standards reported CLAUDE.md still documents Biome. It does not — main's own oxlint migration (9538d58e) updated it, and only the pre-merge copy said Biome. Six findings were demoted to residual risks, chiefly the P1 claim that the mock resolving `undefined` for undeclared paths lets a test pass while production is broken. That is the deliberate, documented contract of the never-crash floor, not a defect; an independent reviewer re-deriving it argues the existing caution callout is warranted, not that the design changed. 4524 tests pass. Re-verified end to end from a repacked tarball in the `databricks apps init` app with all non-loopback sockets blocked. Co-authored-by: Isaac Signed-off-by: Galymzhan * docs(appkit): document `close` as a reserved plugin name `createApp()` rejects a plugin whose manifest name is `close`, because plugin exports are installed as own properties and an own property shadows a prototype method — so such a plugin would silently replace the app handle's teardown. The thrown ConfigurationError already names the offending plugin, but nothing told an author the constraint existed before they hit it. Noted beside where custom-plugins.md introduces `static manifest`. Landing this as part of the `feat:` framing for the branch rather than a BREAKING CHANGE footer: the failure is loud and at boot, not a silent runtime change, and no plugin in this repo is affected. Signed-off-by: Galymzhan * chore: drop the **/.claude ignores from knip, oxlint, and oxfmt These were added earlier on this branch to work around a locked agent worktree at .claude/worktrees/, which every tool saw as a second full copy of the repo: knip reported hundreds of phantom unused exports and failed the pre-commit hook outright, and a repo-root oxfmt would have rewritten that other branch's tree. The worktree has since been removed, so the ignores are treating a symptom that no longer exists and are out of scope for this branch. Verified after removal: `pnpm knip` and `pnpm check` both exit 0 at the repo root. .oxfmtrc.json and .oxlintrc.json are now byte-identical to origin/main. The one remaining knip.json difference — ignoreDependencies: ["vitest"] for packages/appkit — predates this work and is required because vitest is an optional peer dependency of the published testing subpath. Signed-off-by: Galymzhan * refactor(appkit): slim the comments added by the testing-kit work The comments on this branch were far past the repo's own density: reset.ts was 85% comment (35 of 41 lines) for a one-line function, create-test-plugin.ts 63%, lifecycle-manager.ts 48%, mock-workspace-client.ts 45%. Much of it restated the code or ran to several paragraphs where a clause would do. Net 364 comment lines removed. Every file now sits at or below the repo baseline (main's own sources run 20-42%): mock-workspace-client 45% -> 21%, create-test-app 36% -> 25%, lifecycle-manager 48% -> 35%, reset.ts 41 -> 13 lines total, create-test-plugin 63 -> 27. What was kept is the non-obvious "why" that a maintainer would otherwise delete and reintroduce a bug: that `then` must stay in the deny-set or `await client.jobs` hangs; that `ownKeys` stays default or util.inspect mints a mock per probe; that config.host must be a real string; that the three canned defaults are byte-identical because 13 suites depend on them; that phase 5 captures its singletons before the first await; and the four boot hazards behind createTestApp's setup. Pre-existing comments in files this branch only touched (fixtures.ts, test-plugin-context.ts, the shutdown() phase list) are left alone — reverting other people's prose is not this change's business. Also dropped an unnecessary `as Any` cast in createTestPlugin: DEFAULT_CONFIG is already declared on PluginConstructor, so the type escape and its explanatory comment both went. 4524 tests pass; lint, format, and typecheck clean. Signed-off-by: Galymzhan * test(appkit): merge duplicate assertions in the mock-client suite The mock-workspace-client suite had 37 tests written against the plan's checklist rather than against behaviours, so eight asserted something a sibling already covered: getRun memoization twice, config.host being a string twice, the 9-member facade twice (one a strict subset of the other), a rejecting function response twice, getMockFn path resolution twice, function-valued responses twice, the canned defaults twice, and two config-option tests that fit in one. 29 tests now, with no assertion lost — where a dropped test had a unique claim it was folded into the survivor. Three describe blocks became empty and were removed; one had only a comment saying its subject could not be tested at runtime, which the compile-time contract block covers properly. Note for anyone reading this as a bundle-size fix: it is not one. Tests do not ship — the packed tarball contains zero test files — and dropping these eight moved the measured bundle by exactly 0 bytes. The comment trimming in the previous commit is what actually helped (+8.1% -> +6.9%), because JSDoc is preserved in the emitted .d.ts. 4516 tests pass. Signed-off-by: Galymzhan * test(appkit): compact the testing-kit suites The new suites were verbose in ways that cost a reviewer without buying coverage. 1837 lines across 8 files -> 1529 across 7. - mock-workspace-client.test.ts: 471 -> 253. The nine-accessor walk and the canned defaults become one test.each table; per-test client construction goes through a short factory alias. Verified by injection: dropping `then` from the deny-set, changing a canned default, and adding an ownKeys trap each still fail the suite, so the compaction did not gut it. - create-test-app-http.test.ts is folded into create-test-app.test.ts. Both tested one unit through two near-identical 100+ line plugin fixtures; there is now one probe plugin and one manifest helper. Six of the nine routes in the old echo fixture were dead — they duplicated the HTTP file's own plugin. - A withApp() helper absorbs the boot/try/finally/close block that appeared 13 times. It is generic over the plugin tuple so app.plugins stays typed. - The three /headers tests differed only in inputs and expectations, so they are one test.each. The cache reset suite's storage double loses its repeated ended-guard boilerplate. 4509 tests pass; typecheck, lint, and format clean. Signed-off-by: Galymzhan * fix(appkit): address four review findings in the testing kit **Restore three upstream mlflow tests dropped by the merge.** Resolving the route-handler-errors conflict with `--ours` took the whole file from this branch, discarding main's non-conflicting additions from PR #477: the `vi.mock("../mlflow")` hoist, the linkTraceToRun/mockTraceId resets, the parameterised seedPlugin(adapter), the seedEchoPlugin/invoke helpers, and the three trace tests. 15 upstream tests, 12 here, and nothing failed to say so. Restored from 9538d58e alongside this branch's createTestPluginContext rewrite of the aliases test — the two changes are independent. **close() now memoizes itself, not just the phases.** runOnce() guaranteed the teardown body ran once, but resetCoreSingletons() sat outside it, so every call reset again: `await a.close(); createApp(); await a.close()` dropped the second app's singletons. The reset is also skipped when the budget expired, because the phases are still running and still own those instances. Only reachable through the raw AppHandle — createTestApp's wrapper memoizes, which is why the harness-level test could not see it and the regression test lives in app-close.integration. **Refcount singleton ownership.** The env baseline was already refcounted so overlapping harness apps compose, while the singleton layer reset on every boot and every close — so booting B rebound A's ServiceContext and CacheManager, and closing A while B was live left B with none at all. claimCoreSingletons/ releaseCoreSingletons now follow the same model as the env baseline: first boot claims, last close drops. **Fake the on-behalf-of client.** The kit promised "no workspace, no credentials, no network", but createApp({ client }) installs only the service principal; an `obo` request reached ServiceContext.createUserContext, which builds a real SDK client from process.env.DATABRICKS_HOST. The harness now stubs that for the app's lifetime and restores it on close, mirroring fixtures.ts's createUserContextSpy. Every fix has a test verified by reintroducing the bug. Two of those tests needed a second attempt: asserting the OBO client's host does not discriminate, since a real client carries the same DATABRICKS_HOST string — the test now asserts the harness's mock recorded the call. The probe plugin gained a route that calls the client under asUser, because the existing /as-user route only reads ctx.userId, which is how this escaped notice. 4517 tests pass; typecheck, lint, and format clean. Signed-off-by: Galymzhan * chore: restore the dev-playground client lockfile Swept into merge 45601bf1 from a dirty working tree; not part of this work. npm had run in apps/dev-playground/client during the local-tarball verification and pruned the `extraneous: true` entries for ../../../packages/appkit-ui. Both parents of that merge carry the same blob, so `git log -- ` reports nothing for this branch even with --full-history, which is why it went unnoticed. Restored to origin/main byte-for-byte. Signed-off-by: Galymzhan * fix(appkit): correct double singleton release and simplify the testing kit createTestApp's failure path released its singleton claim twice when the boot failed after createApp resolved: app.close() already drops the claim, and the catch block released it again, which could pull the singletons out from under a still-live sibling app. It now releases only when nothing was booted. Alongside it, a simplification pass over the branch: - LifecycleManager.closeOnce: drop the timedOut flag in favour of an early return from the catch. - createTestApp: drop the restoreEnv alias for releaseEnvBaseline. - Delete createConfigurableMockWorkspaceClient. Its last caller moved to createMockWorkspaceClient earlier in this branch, so the JSDoc rationale for keeping it byte-for-byte pointed at deleted code, and the ./testing subpath it would have shipped on is new here and unreleased. - Reuse the kit's createMockRequest instead of three local mockReq helpers that re-rolled the same forwarded-identity headers. - genie.test.ts: reuse one expectStream handle rather than parsing the same captured SSE body twice. Signed-off-by: Galymzhan * chore: drop the .claude ignores and restore the client lockfile Both were swept into the oxlint merge from a dirty working tree and corrected later on the branch; folding those corrections in here keeps them out of the follow-up PR. The knip `.claude/**` entry and the `**/.claude` ignorePatterns in oxfmt/oxlint were never needed — nothing in the repo lints or formats that directory. The `packages/appkit` vitest ignoreDependencies entry stays: vitest is a real dependency of the testing entry. apps/dev-playground/client/package-lock.json is restored to origin/main byte-for-byte; npm had run in that directory and pruned its `extraneous: true` entries. Signed-off-by: Galymzhan * refactor(appkit): trim the cache reset test and its comments The 3-line `CacheManager.reset()` carried 30 lines of comment and 135 lines of test. Mutation testing showed what each test was actually worth: with `reset()` stubbed to a no-op, or clearing only `instance`, four tests fail; with the generation guard removed, exactly one does. - Dropped "reset is safe when the cache was never initialized". It survived all three mutations — a body of three assignments cannot throw, so the test could only ever pass. - Replaced the hand-rolled 23-line `CacheStorage` double with a 7-line `InMemoryStorage` subclass overriding just `close()` and `set()`. It also stops claiming `isPersistent() === true`, which had the manager's probabilistic cleanup eligible to fire against ended storage. - Cut the comments to the two facts a maintainer would otherwise remove and reintroduce the bug with: both fields must clear because `getInstance()` falls back to `initPromise`, and a reset is a pointer drop so callers close first. Same treatment for `reset-singletons.ts` and `testing/reset.ts`, which were at 47% and 52% comment lines against a repo baseline of 20-35%. Test file 135 -> 115 lines, production diff +38 -> +23. Mutation coverage is unchanged, re-verified against all three mutations. 4517 tests pass; typecheck, lint, and format clean. Signed-off-by: Galymzhan * docs(appkit): relocate and tighten the app-handle comments 60% of this branch's additions to appkit.ts were comment: 41 lines over 22 lines of code. Trimmed to 35 with no fact dropped — audited claim by claim — and two of them moved somewhere they do more good. **The onPluginsReady note was on the wrong function.** It sat on the internal `_createApp`, which typedoc does not publish, while the public `createApp` that callers actually read carried the same parameter with no explanation. Moved, and the generated API page now renders it (see the Function.createApp.md diff) where before it reached nobody. While moving it, corrected the hazard it described. It said offering `close()` there "would invite tearing down a half-booted app" — but `#lifecycle` is not assigned until after the server starts, so `close()` at that point is a *no-op*, not a teardown. The narrow type and the `#lifecycle?.` optional chain guard against silently skipping cleanup, which is the opposite failure. **RESERVED_PLUGIN_NAMES gained the reasoning for its scope.** It reserves only `close`, which reads like an oversight: `bindExportMethods` and the other prototype methods are equally shadowable, since TS `private` is compile-time only. The distinction is that shadowing those throws `TypeError` on the next plugin's registration, while shadowing `close` fails silently — you call it, get no error, and leak every socket and pool. Only silent breakage needs a guard, and that is now written down. Also dropped the comment above the LifecycleManager construction, which restated the `#lifecycle` field's own JSDoc. 4517 tests pass; typecheck, format, and docs build clean. Signed-off-by: Galymzhan * test(appkit): cover the singleton release on close()'s timeout path Mutation-tested the close-handle suites (33 tests, 8 mutations). One mutation survived: deleting the early `return` in `closeOnce`, so a timed-out close() releases the core singletons while its own phases are still running and still own those instances. That is the exact hazard the code comment warns about, and it was the review fix with no coverage. The test that sounds like it covers this — "phase 5 still closes the app's own cache and telemetry, not the next app's" — cannot: it mocks CacheManager and TelemetryManager wholesale, so whether releaseCoreSingletons() ran is invisible to it. It guards the captured-singleton half of the fix, not the early return. Closed by mocking the one symbol lifecycle-manager imports from reset-singletons and extending two existing tests, so the count stays at 24: the clean-close test now asserts one release, and the hung-teardown test asserts none. Verified in both directions — dropping the early return fails the second, removing releaseCoreSingletons() entirely fails the first. Also corrected a comment in the phase-5 test that predated the review fix. It said close() "already dropped the singletons" on timeout, which is what the fix stopped it from doing. 4517 tests pass; typecheck and format clean. Signed-off-by: Galymzhan * fix(appkit): lowercase mock request header keys, as Express does `createMockRequest` stored header keys exactly as given while `header()` lowercased the lookup, so a mixed-case override was unreachable: createMockRequest({ obo: { userId: "alice" }, headers: { "X-Forwarded-User": "bob" } }); // header("x-forwarded-user") === "alice" Both keys were kept — ["x-forwarded-access-token", "x-forwarded-user", "X-Forwarded-User"] — and the lowercase one obo seeded still answered, which contradicted the "an explicit override wins" contract documented right above it. Keys are now lowercased on the way in, matching what Node's parser hands Express. Thanks @pkosiec. The existing override test passed because it used a lowercase key, so it is now parametrised over both casings, and the case-insensitivity test additionally pins that every stored key is lowercase. Reverting the fix fails both. Signed-off-by: Galymzhan * fix(appkit): lowercase caller header keys in the harness request methods Same root cause as the `createMockRequest` fix on the parent branch (#530, found by @pkosiec), and worse here. `Object.assign(headers, reqOptions.headers)` kept case variants as separate keys, and `Headers` **comma-joins** duplicates rather than replacing them: new Headers({ "x-forwarded-user": "alice", "X-Forwarded-User": "bob" }) // -> x-forwarded-user: "alice, bob" So a mixed-case override did not merely lose, it corrupted the value the server received — the "caller headers last" contract broken in a way that produces a plausible-looking string rather than an error. Keys are lowercased before assignment. The existing `/headers` table gained a mixed-case row rather than a new test; reverting the fix fails it. Signed-off-by: Galymzhan * chore(appkit): acknowledge the testing kit's tarball growth in the baseline The bundle-size gate is the last red check on this PR. The growth is the feature, not slack: the tarball ships no test files at all, so deleting eight tests moved the measured size by exactly 0 bytes. Comment trimming did help, taking it from +8.1% to +6.9%, because JSDoc survives into the `.d.ts`. `@databricks/appkit` packed 860106 -> 919127 (+58 KB), which is the `./testing` entry: 61 KB gz of harness plus its type declarations. `@databricks/appkit-ui` moves -291 bytes, incidentally. Regenerated with `pnpm size:baseline` from a clean build — `rm -rf packages/*/dist packages/*/tmp` first, because tsdown runs with `clean: false` and a stale artifact would be baked into the baseline as real growth. `pnpm size:compare` now reports no change and exits 0. Signed-off-by: Galymzhan * Revert "chore(appkit): acknowledge the testing kit's tarball growth in the baseline" This reverts commit 04839f0ba41608be2894600f5c44651e997a0dbd. * fix(appkit): address Pawel's review — singleton leak, renames, and guards Nine of the seventeen review comments, the ones with no open decision. **The singleton leak (P2, both instances).** On a `close()` timeout the release was skipped deliberately — the phases still own those instances — but nothing released it afterward either, so the refcount never returned to zero and every later boot skipped its reset and inherited a half-closed app. Now the release is attached to the still-running teardown, so it happens once the phases settle. This also fixes the boot-failure report: that path calls `app.close()`, so it routed through the same gap. My own test was pinning the bug. It asserted `releaseCoreSingletons` was never called on the timeout path, when the contract is "not yet, but once teardown settles". Corrected, and verified by reverting the fix. **Renames, while the pre-release window is open** (both symbols are absent from 0.62.0, so this is free now and breaking later): `getMockFn` -> `getMock`, since it returns vitest's `Mock`; `resetAppKitSingletons` -> `resetGlobalState`, since vitest's own resets are `reset` with no product prefix. **`getListeningPort` was `@internal` and exported**, which is a contradiction either way. Made public: three integration suites already consume it through the public entry, and it is a reasonable helper for anyone hand-rolling a server. **The reserved-name list is now exhaustive by construction.** It was a hand-kept `Set(["close"])` with nothing tying it to `AppHandle`, so a future named method there would be silently shadowable by a plugin. A `Record, true>` makes adding one a compile error until it is reserved — verified by adding a `restart()` and watching tsc fail. **Passing both `client` and `responses` now throws** instead of silently ignoring the responses, matching the existing `server: false` conflict. **The two proxy traps share one guard.** `toLegacyWorkspaceClient` had copied the symbol + PASSTHROUGH_DENY logic, so a key added to the set would have covered one trap and missed the other — and missing `then` there is what makes `await client` hang. **Coverage gap closed:** the boot-failure branch that runs *after* `createApp` succeeded had no test. Added one via a plugin that boots and then throws when the harness reaches for its socket. **Comment de-fragilised:** the canned-defaults note named a count ("13 suites") and a position ("the first three"). It now names the three entries. 4521 tests pass; typecheck, lint, format, and docs build clean. Signed-off-by: Galymzhan * fix(appkit): match the real user-context shape in the harness stub Half of Pawel's P2 on the diverging on-behalf-of fakes. The other half is deferred on semver grounds, recorded in internal/testing-kit/README.md. `stubUserContext` set `tokenFingerprint: test-${userId}` — keyed on the *user*, so constant across tokens. Lakebase rotates its pool by comparing that value (`pool-manager.ts:102`), which means rotation could never fire under `createTestApp` while production rotates on every new token. A test exercising pool rotation through the harness would have passed against a code path that cannot run. It now derives `sha256(token)[0:16]`, exactly as the real `createUserContext` does. The rejection also matched only in spirit: a bare `Error` where production throws `AuthenticationError.missingToken`. A handler that catches that class took a different branch under the harness. Now the same class. Two tests, both verified by reintroducing the old behaviour: the fingerprint is stable per token and differs across tokens (a user-keyed value fails it), and the missing-token rejection is asserted on the error class, not the message — a plain `Error` whose message still says "token" fails it. Not touched: `createUserContextSpy` in fixtures.ts, which neither throws nor sets the field. It shipped in 0.62.0, so tightening it would break consumers' existing tests on a minor upgrade; it needs its own PR and a changelog line. The end state is one shared builder, worth extracting once that half is free to move. Signed-off-by: Galymzhan * fix(appkit): share one user-context fake across the testing kit Completes Pawel's P2. Three fakes of `ServiceContext.createUserContext` disagreed and none matched production; they now share `fakeUserContext` in fixtures.ts — production's rejection (`AuthenticationError.missingToken`), production's `sha256(token)[0:16]` fingerprint, and `userEmail`. What the divergence cost: `pool-manager.ts:101` treats a missing `tokenFingerprint` as "not stale", so the old fixtures spy — which discarded the token outright — made Lakebase's drain-and-recreate branch unreachable. A test written to verify pool rotation would have passed while exercising nothing. I deferred this half earlier on semver grounds, arguing that `mockServiceContext` shipped in 0.62.0 so tightening a published test double would break consumers on a minor upgrade. Wrong: `./testing` was not published at all before 0.62.0 (`publishConfig.exports` lists no `./testing` in 0.61.0 or 0.61.1), and 0.62.0 landed four days ago with two releases since. There is no population to break. Checking the API's actual age beats reasoning about semver in the abstract. Harmless in-repo too — the full suite passed with the fake tightened, so nothing here relied on the looseness either. Six tests pin it, three per fake, all verified by restoring the loose version: the fingerprint is stable per token and differs across tokens, a missing token throws production's error class, and userEmail is carried through. 4579 tests pass; typecheck, lint and format clean. Signed-off-by: Galymzhan * docs(appkit): deslop the testing guide and fix its sidebar slot Pawel's P1 on the guide's voice, plus the two smaller docs points. **Voice.** Ran the repo's doc-deslop pass, then a targeted follow-up. No vocabulary slop or hedging to remove — the divergence from sibling pages was punctuation and register: paired em-dashes injecting lists mid-sentence (the textbook tell), British `behaviour` against the siblings' American spelling, and cutesy editorial phrasing. `:::caution The honest catch` is now `:::caution Undeclared methods return undefined`, matching how siblings title admonitions. Worth recording how the "how florid is it" question actually resolved, because my first two measurements were wrong. Raw em-dash counts said this page was 2.7x denser than `analytics.md`. But that counted comments inside code blocks and the `**term** — definition` bullet pattern the siblings use too. Excluding both: 0.286 dashes per prose line here against 0.252 in `analytics.md`, its closest sibling by length and depth. The page was already at parity, so the remaining dashes were left alone rather than purged below the house norm. Four where a dash was doing a colon's job are now colons. **Teardown now leads with `await using`.** It previously said "use try/finally, or let the runtime do it" and then showed only the `await using` example — naming the weaker option first and never demonstrating it. `await using` comes first with the reason (it closes on a thrown error too), and `try/finally` follows as the form you need when the app outlives a block. Both are supported: TypeScript 5.9.3, Node 24, ES2022, and three suites already rely on it. **Sidebar.** `sidebar_position: 8` collided with `jobs.md` and `manifest.md`; `_category_.json` uses `autogenerated`, so frontmatter really does drive order. Moved to 10, which is unused. The pairs at 2, 5, 6, 7, 8 and 9 predate this work and are left for a separate tidy-up. Also documents two things the code now does but the page didn't say: passing `client` and `responses` together is refused rather than ignored, and `getListeningPort` is public API. Signed-off-by: Galymzhan * chore(docs): regenerate the appkit-ui stylesheet Unrelated to the testing kit. `styles.gen.css` is generated by the docs build and main's committed copy is stale: it lacks the `.collapse` utility. Any `pnpm docs:build` reproduces this diff, so it otherwise sits dirty in every working tree. Verified deterministic — reverted, rebuilt, and it came back. Signed-off-by: Galymzhan * feat(appkit): add strict mode to the mock workspace client Pawel's P3 on undeclared paths resolving `undefined`. I first declined this as a feature deserving its own PR; the actual change is eight lines of production code and additive, so that was over-caution. `createMockWorkspaceClient({ strict: true })` throws when a path with no declared response is *called*, naming the path. Off by default — the never-crash floor is what lets a plugin touch services a test does not care about — so nothing about existing behaviour moves. Three details that needed care: - **Throws on call, not on mint.** `getMock(client, path)` mints so a test can hold a handle before the code under test runs; blowing up there would break the normal assertion pattern. Pinned by a test. - **The canned defaults count as declared**, so a harness boot still works — it reads `currentUser.me` through this client. With `defaults: false` they are undeclared again, which is also pinned. - **Plumbed through `createTestApp`.** Without that the option was unreachable from the recommended entry point: the harness builds its own mock, and passing a hand-built client now refuses `responses`. The conflict guard covers `strict` too. This does not overlap with the typed-`responses`-keys idea, contrary to the review's framing that they are two answers to the same problem: typed keys catch a *typo* at compile time, `strict` catches an *omission* at runtime. Typed keys cannot know you forgot to declare a path you actually call. Six tests, and the guide documents it beside the caution that describes the failure mode. 4585 tests pass; typecheck, lint, format and docs build clean. Signed-off-by: Galymzhan * fix(appkit): address the straightforward half of the second review pass Nine of Pawel's thirteen round-2 comments, all verified true first. The two P1s, the template CI wiring, and the code-comment deslop are left: they need the concurrency decision or a wider change. Recorded in internal/testing-kit/review-log.md. **The timeout release is now immediate, reversing last round's fix.** His first pass said to defer it with `.finally()`; his second says that is wrong, and he is right. Holding the refcount makes the next boot skip its reset, so it inherits this app's `CacheManager` — which the orphaned teardown then closes in phase 5, mid-test. Releasing at once is safe because `runPhases` captures its own cache and telemetry before the first await and never re-reads the shared slots. Shorter and safer. My test asserted the deferred behaviour, so it was pinning the wrong contract — the second time a test of mine has done that on this exact code. **A seeded `apiClient.userAgent` no longer returns a Promise.** Every non-function seed was wrapped in `mockResolvedValue`, so `responses: { "apiClient.userAgent": "x" }` produced `"[object Promise]"` in a Headers value — precisely what the synchronous default exists to prevent. Seeds now match each member's own shape; `request` stays async. The `config.*` path never had this because it assigns values directly. **The published-surface test now guards the published surface.** It imported three symbols, so every other export could vanish from the barrel with this file still green. It now asserts twenty names are reachable through the entry — verified by dropping `createTestPluginContext` and watching it fail by name. `tsc` would also catch that via other suites, which softens his "CI still green" wording, but a test claiming to guard the surface should do it. Also: dead `createHash`/`AuthenticationError` imports removed (left by the `fakeUserContext` move); the boot-failure path now honours `closeTimeoutMs`; the two bare `type Any = any` aliases carry the explanation the third one had; the "seven services" count is gone from a comment; the guide now says the client and the OBO stub are process-wide, not per app, and bounds its `app.client` claim to a single open app; and it distinguishes mock from fake in one line. Finally the `nodeEnv` JSDoc records what the option actually changes beyond refusing `development` — `errorHandlerMiddleware` redacts 5xx bodies only under `production`. That is the evidence his round-1 comment proposing to delete the option overlooked, and it was undocumented, which is a fair reading of why he missed it. 4587 tests pass; typecheck, lint, format and docs build clean. Signed-off-by: Galymzhan * fix(appkit): close the leftover halves of three review comments The docs half of the stale-count comment was never done: testing.md said "AppKit owns that 9-member interface" while `WorkspaceClient` declares ten members (seven services, `config`, `apiClient`, `toLegacyWorkspaceClient`). Drop the count and keep the point, which never needed it. The published-surface test asserted `createTestPluginContext` was a key on the namespace but never called it, so a hollowed-out export would still pass. It now drives the context's real tool registry and on-behalf-of path through the entry. Verified by disabling the dispatch recording: the new test fails, the name check still passes. Also reword the three story-voice comments quoted in review that survived the earlier pass; the other three were already rewritten with the timeout-release fix. Signed-off-by: Galymzhan * fix(appkit): allow one harness app at a time Both P1s from the second review pass are consequences of letting two createTestApp apps overlap, which nobody decided to support. AppKit's workspace client, cache, and on-behalf-of fake are process-wide, so a second live app cannot own its own: its handlers resolve the first app's client while `app.client` returns the second mock, and `close()` restores the real `createUserContext` rather than the other app's spy, putting the survivor's on-behalf-of calls back on the network. Refuse the second boot instead. The guard reuses the existing live-app counter rather than adding a parallel flag that could drift, and runs before any mutation so a refused boot leaves the open app untouched — asserted, not assumed. Vitest isolates test files in separate workers, so this only constrains apps within one file. The cost lands on a `describe` that holds an app in `beforeAll`: it can no longer contain a test that boots its own, which is why the concurrency and on-behalf-of tests move to the top level. Two overlapping-boot env tests are replaced by one covering repeated boot/close cycles; the hazard they pinned (a second snapshot capturing the first's mutations) is now unreachable. Verified by disabling the guard: the three refusal tests fail. Signed-off-by: Galymzhan * ci(appkit): run the scaffolded template's tests, and ignore its staging dir template/server/example.test.ts ran nowhere, so the example shipped to anyone scaffolding an app could rot silently. Wiring it up turned out to need one step rather than a new job: pr-template-artifact already builds and packs the branch, rewrites the template's deps to those tarballs, and npm installs with devDependencies, so the test file and vitest are already staged and installed. The step runs before the zip, so a broken example fails the build instead of being published as a downloadable artifact. It is also the only CI check that consumes @databricks/appkit/testing the way a customer does — through a real npm install of a packed tarball, rather than a workspace path. Verified locally end to end: pack, stage, install, npm test (3 passed), and grep on the installed dist confirms it ran against the branch's tarball, not the published version template/package.json pins. The lockfile is byte identical across the test step, so the registry-rewrite step that follows is unaffected. Also gitignore pr-template/ and appkit-template-*.zip, which prepare-template-artifact.ts leaves in the working tree — a ~200MB staging directory that `git add -A` would otherwise commit. Signed-off-by: Galymzhan * refactor(appkit): drop the dead re-exports from the deprecated test-helpers shim tools/test-helpers.ts is marked @deprecated, yet this branch had added eight exports to it: createTestApp, CreateTestAppOptions, getListeningPort, getMock, MockWorkspaceClient, resetGlobalState, TestApp, TestRequestOptions. None had a consumer. The fifteen files still importing the shim use only nine symbols between them, all predating this work, so the eight came along mechanically with the getMockFn -> getMock rename and never had a caller — visible in that they were inserted mid-list, breaking the alphabetical order the rest of the list keeps. Removing them keeps the shim's surface to what an existing importer actually needs, so the harness API has one import path rather than two. Typecheck and the full suite pass unchanged, which is what confirms they were dead. Signed-off-by: Galymzhan * feat(appkit): add withEnv scoped-environment test helper withEnv(vars, fn) sets env for the duration of fn and restores each key's prior state on exit — including deleting keys that were previously unset, rather than the blanket delete the current test pattern uses. Sync and async forms; nested calls restore LIFO. Exported from @databricks/appkit/testing. Signed-off-by: Galymzhan * feat(appkit): add createApiError test-error factory createApiError({ statusCode, message, errorCode }) returns a genuine ApiError instance, so error-path tests can assert real instanceof ApiError checks instead of hand-rolled look-alikes that pass name but fail instanceof. Exported from @databricks/appkit/testing. Signed-off-by: Galymzhan * feat(appkit): add composition options to createTestPluginContext An optional second parameter { responses, env, strict } composes a plugin unit-test in one call: a mock workspace client seeded from responses, an installed service context wired to it, and scoped env — all auto-restored via an afterEach hook, with an idempotent restore() escape hatch. Synchronous and non-breaking: no-options callers are unchanged. Exports TestPluginContextOptions. Signed-off-by: Galymzhan * test(appkit): adopt withEnv for inline env handling in plugin suites Replaces 11 hand-rolled process.env set/try/finally/delete blocks with withEnv across the files, jobs, and ai-search test suites. Behavior-preserving; the vi.mock/vi.hoisted preludes are untouched (the broader internal-mock migration is deferred). Signed-off-by: Galymzhan * docs(appkit): document the new testing-kit helpers Documents withEnv, createApiError, and the createTestPluginContext options parameter in the testing guide, and surfaces the already-shipped createMockRouter fixture. The composition example is shown synchronously (no await on the factory). Signed-off-by: Galymzhan * refactor(appkit): drop unused import and bindings in the composition tests Removes the unused ServiceContextMock import and three side-effect-only mock bindings, and stops capturing an unused priorInitialized — clearing the lint warnings the composition-options work introduced. No behavior change. Signed-off-by: Galymzhan * refactor(appkit): share env capture/restore and guard withEnv error paths Addresses two code-review findings. Extracts a shared applyEnv() helper so withEnv and the createTestPluginContext options path no longer duplicate the env capture/restore loop. Guards withEnv's error paths so a failing restore can no longer mask the caller's original error (a restore failure still surfaces on the success path). Behavior-preserving for normal use; all tests pass. Signed-off-by: Galymzhan * feat(appkit): add useTestCache helper to the testing kit Boots AppKit's real in-memory cache singleton and clears it before each test, exposing the real CacheManager so plugin tests assert cache-key behaviour through production's generateKey instead of mocking the internal cache module. Signed-off-by: Galymzhan * test(appkit): add internal createCacheMock passthrough fake Shared, unexported factory for the passthrough cache instance the non-behavioural suites mock; adopted across those suites in a follow-up unit. Signed-off-by: Galymzhan * test(appkit): migrate analytics suite off the cache mock to useTestCache Drops the hand-rolled functional cache fake; the suite now runs against the real in-memory cache. Reworks the abort-fallback test to model a real client disconnect (response close aborts the handler signal) rather than spying the signal executeStatement happens to receive. Signed-off-by: Galymzhan * test(appkit): migrate metric suite off the cache mock to useTestCache Runs against the real in-memory cache; the cache-key-equality test now spies the real getOrExecute to capture the composed key parts instead of a fake. Signed-off-by: Galymzhan * test(appkit): migrate ai-search suite off the cache mock to useTestCache Runs against the real in-memory cache. Completes the telemetry mock span (adds end/addEvent/etc.) since the real getOrExecute opens a span the old functional cache fake never did. Signed-off-by: Galymzhan * test(appkit): adopt shared createCacheMock in files/plugin suite Replaces the copy-pasted passthrough cache fake with createCacheMock(), pulled in via an async vi.hoisted + dynamic import (require can't resolve the TS helper inside a hoisted block). Signed-off-by: Galymzhan * docs(appkit): document useTestCache for asserting cache behaviour Signed-off-by: Galymzhan * test(appkit): adopt shared createCacheMock across passthrough cache suites Replaces the copy-pasted passthrough cache fake in 16 suites (files/*, genie, jobs, serving, lakebase, connectors/lakebase, analytics.readonly) with createCacheMock(). server.test.ts keeps its own mock (close(), not a passthrough getOrExecute). Signed-off-by: Galymzhan * fix(appkit): give createCacheMock getOrExecute the real userKey arity The fake omitted the third (userKey) argument the real getOrExecute takes, so calling it with three args failed typecheck. Signed-off-by: Galymzhan * docs(appkit): tighten testing-kit guide prose Removes reasoning-trace and over-explained wording across the testing guide (content from #530/#540/#555) and the reserved-name callout, matching the terser house style. No facts or examples changed. Signed-off-by: Galymzhan * feat(appkit): add useTestApp to wire the harness hooks for a suite The kit shipped two hook-wiring helpers — useServiceContextMock and, in this stack, useTestCache — and none for the app itself, so a suite needing an app per test hand-wired beforeEach/afterEach and had to remember close(). That gap widened when useTestCache adopted the pattern for the cache and skipped the app. It matters more since the harness began allowing one open app at a time. A forgotten close() used to leak a listener quietly; now the next boot throws and takes the rest of the suite with it. await using is still shorter for a single test, but it cannot carry an app from a beforeEach into the test body, and a describe holding one in beforeAll cannot contain a test that boots its own — so beforeEach/afterEach is the remaining pattern for per-test apps, and it was the one without a helper. Mirrors useTestCache: same { current } accessor, same call-it-in-a-describe rule, same self-explaining error when read outside a test. The afterEach clears the handle before awaiting close, so a close that throws cannot leave a stale app readable by the next test. The one-app-at-a-time guard makes the tests discriminating: verified by removing the close, which fails four of them with "a harness app is already open" rather than passing quietly. Removing the barrel export fails the published-surface assertion by name. Signed-off-by: Galymzhan * test(appkit): fake the workspace client through the kit, not module mocks Fourteen suites hand-rolled a workspace client and an ApiError look-alike, then patched `../../workspace-client` and `../../context` so production code would pick them up. None of that expressed anything the shipped kit cannot: the client is `createMockWorkspaceClient`, the error is `createApiError`, and injecting the client through `mockServiceContext` makes `getWorkspaceClient()` resolve to it via the real ServiceContext — so nineteen of twenty-four module mocks go away. `setupTestEnv` in the files suite now takes the client and passes it through, which is what let nine files drop their `context` mock at once. Strictness is preserved deliberately, not incidentally. A hand-rolled object literal threw a TypeError on any undeclared call; `strict: true` keeps that loudness with a sentence instead of a stack trace. Verified by counting rather than asserting: 306 tests before and after, 558 expects before and 559 after. One class of assertion gets stronger. `connectors/files` asserted `toBeInstanceOf(MockApiError)` — circular, since it could only pass while the module was patched to install that fake. It now asserts the real `ApiError`, and a mutation that throws a plain Error fails it, which the old form could not catch. Two mocks are kept on purpose. jobs keeps a three-line `Context: vi.fn()` — the SDK cancellation-token class, an unrelated concern the old mock also served. The nine files that mock `workspace-client` to stop a real client being built have no injection seam (they test ServiceContext, CacheManager, or the type-generator), so module mocking is correct there. Not migrated: files/plugin.test.ts (3978 lines, 49 client refs, 7 local getRouteHandler copies) wants a reviewed diff rather than a scripted one, and ai-search pins getCurrentUserId constant on purpose so cache-key scoping is driven only by executorKey — the real path would change the keys its tests assert on. Signed-off-by: Galymzhan * chore: gitignore the .codex-tmp agent scratch dir Signed-off-by: Galymzhan * test(appkit): replace createCacheMock with useTestCache across suites Migrate 17 plugin/connector test suites off the internal vi.mock("../../../cache") + createCacheMock fake onto the published useTestCache(), which boots AppKit's real in-memory CacheManager — the seam a third-party plugin author can actually reach. Delete createCacheMock and its unit test; nothing imports them. - Passthrough suites (files shutdown/volume-config/download-endpoint/ error-handling/path-validation/raw-endpoint, serving, genie, analytics.readonly, lakebase-agent-tool, jobs) drop the hoisted mock for a top-level useTestCache(). - Connector suites (routing-pool, pool-manager) never touched the cache; the mock was vestigial and is removed outright. - Invalidation and cache-behaviour assertions (files delete/mkdir/upload and files/plugin.test.ts OBO + invalidation-key tests) now spy on the real cache via vi.spyOn(testCache.current, ...), so the keys asserted are production's generateKey, not a re-implemented copy. Two jobs error tests rejected a plain Error carrying a statusCode; the passthrough mock let that status leak through, but the real cache's getOrExecute only preserves an ApiError status and wraps anything else to 500. Switched them to throw a genuine ApiError via createApiError, which is what the SDK does. The context/workspace-client module mocks in files/plugin.test.ts and ai-search.test.ts remain: their OBO tests toggle getWorkspaceClient / getCurrentUserId between the mocked and real impl, which requires the module mocked — removing them is gated on the ServiceContext-as-instance refactor. Signed-off-by: Galymzhan * test(appkit): drop the stale EXPERIMENT label from mkdir test The no-module-mock approach it flagged is now the default across every files/* suite; the rationale lives in setupTestEnv. Comment-only. Signed-off-by: Galymzhan * fix(appkit): make createTestPluginContext options auto-restore actually fire The options overload of `createTestPluginContext(fakes, options)` registered its cleanup with `afterEach(() => restore())` at call time. It is documented and used from inside a test body, where a runtime-registered `afterEach` never runs for that test (Vitest collects `afterEach` before the test executes) — so the advertised auto-restore silently no-op'd and the mocked ServiceContext spies / env could leak past the creating test. Switch to `onTestFinished`, the hook built for runtime registration, which runs after the creating test; fall back to `afterEach` when called at collection scope (where `onTestFinished` throws). `restore()` stays idempotent, so manual restores still work. The test that should have caught this ("...auto-restored after test via afterEach") asserted nothing after cleanup — it only checked the mock was active, then restored manually. Replace it with an ordered two-test pair that creates the context with NO manual restore and asserts in the next test that the spy was removed (`vi.isMockFunction(...) === false`). Verified: the new test fails against the old `afterEach` implementation and passes with `onTestFinished`. Signed-off-by: Galymzhan * refactor(appkit): drop inert refcount from singleton-reset machinery createTestApp is the only claim site and forbids a second concurrent harness app unconditionally, while production createApp never claims, so `owners` was always 0 at claim and 0 at release: claim, release, and drop were the same operation and the counter arithmetic was a no-op. The memo in LifecycleManager.close() (`this.closed ??= ...`), not the refcount, is what stops a stale handle's second close() from resetting a newer app. - reset-singletons.ts: remove `owners`, `claimCoreSingletons`, `releaseCoreSingletons`; keep only `dropCoreSingletons`. - create-test-app.ts / lifecycle-manager.ts: call `dropCoreSingletons` directly (import from ../core/reset-singletons); close() memo preserved. - reset.ts: drop the `@internal` claim/release passthroughs, keep the public `resetGlobalState`. - liveHarnessApps: collapse the 0-or-1 counter into a boolean. No behaviour change; the no-concurrency guard and the CacheManager entry in the reset list are untouched. Signed-off-by: Galymzhan * refactor(appkit): drop the multi-boot/embeddable surface from the testing kit AppKit will not support multiple apps in one Node process, so the programmatic teardown / embeddable-app surface the testing-kit PR reached for comes out (addresses Mario's #540 review): - drop AppKit.close(), [Symbol.asyncDispose], and the AppHandle return type; createApp() returns PluginMap again, and the close-name-shadow guard goes - fold the LifecycleManager close/runOnce/removeSignalHandlers split back into shutdown() plus an internal dispose() for the harness; the harness boots with an internal installSignalHandlers:false flag rather than removing handlers - createTestApp tears down what it booted (dispose runs the plugin shutdown hooks and closes the server) then drops the core singletons and restores env - relocate dropCoreSingletons() into testing/ (kit-owned; no core caller left) - drop the cache reset-generation guard: no reset-during-init race under single-boot with awaited sequential teardown (invariant documented in-code) - delete testing/reset.ts (resetGlobalState) — no callers - cover the rewired teardown end-to-end: a booted plugin's shutdown() hook fires and the server socket is released after close() Signed-off-by: Galymzhan * refactor(appkit): collapse the shutdown/dispose split into shutdown({ exit }) (#540 review) Mario's M7: the shutdown()/dispose() split existed only so the harness could run the teardown phases without process.exit. Fold dispose() into shutdown({ exit?: boolean }) — exit defaults true (signal path: force-exit backstop + process.exit); { exit: false } is the harness path (await the phases, no exit). AppKit[disposeApp] now calls shutdown({ exit: false }). With the split gone, runPhases()'s capture-before-await of the cache/telemetry managers is provably dead: the orphan race it guarded (a timed-out dispose returning while the phases still ran, then the harness dropping singletons) cannot occur once the harness fully awaits the phases before dropping. Same basis as the 6dffb59d CacheManager reset-generation-guard removal ("no race under single-boot with awaited sequential teardown"). closeCacheStorage() and flushTelemetry() read their manager at phase-5 time again; DISPOSE_TIMEOUT_MS is deleted (the phases are already individually bounded). Also drop the now-inert createTestApp closeTimeoutMs option (0 refs; the harness path has no outer budget of its own any more). Tests: the harness-path describe uses shutdown({ exit: false }); the two tests that only exercised the removed orphan / outer-timeout machinery are replaced by one asserting the harness path is bounded by the internal per-plugin timeout. Local commit for #540 review; not pushed. Signed-off-by: Galymzhan * docs(appkit): document the public withEnv helper The withEnv JSDoc (the sync/async @example block) was stranded above applyEnv, where applyEnv's own second doc block won and the first was ignored, leaving the public `export function withEnv` with no doc in IDE hover or TypeDoc. Move it directly above `export function withEnv`; applyEnv keeps only its own (correct) doc. Addresses Jorge's review comment. Co-authored-by: Isaac Signed-off-by: Galymzhan --------- Signed-off-by: Galymzhan Co-authored-by: Isaac --- .gitignore | 2 + docs/docs/plugins/testing.md | 220 +++++++++++++++++- .../src/connectors/files/tests/client.test.ts | 86 +++---- .../lakebase/tests/pool-manager.test.ts | 15 -- .../lakebase/tests/routing-pool.test.ts | 15 -- .../src/plugin/tests/asUser-proxy.test.ts | 11 - .../appkit/src/plugin/tests/plugin.test.ts | 58 ++--- .../plugins/ai-search/tests/ai-search.test.ts | 61 ++--- .../tests/analytics.readonly.test.ts | 19 +- .../plugins/analytics/tests/analytics.test.ts | 68 ++---- .../plugins/analytics/tests/metric.test.ts | 50 +--- .../src/plugins/files/tests/_test-helpers.ts | 10 +- .../src/plugins/files/tests/delete.test.ts | 102 +++----- .../files/tests/download-endpoint.test.ts | 97 +++----- .../files/tests/error-handling.test.ts | 177 +++++++------- .../src/plugins/files/tests/mkdir.test.ts | 94 +++----- .../files/tests/path-validation.test.ts | 99 +++----- .../files/tests/plugin.integration.test.ts | 41 ++-- .../src/plugins/files/tests/plugin.test.ts | 120 ++++------ .../plugins/files/tests/raw-endpoint.test.ts | 101 +++----- .../src/plugins/files/tests/shutdown.test.ts | 79 ++----- .../src/plugins/files/tests/upload.test.ts | 86 ++----- .../plugins/files/tests/volume-config.test.ts | 124 ++++------ .../src/plugins/genie/tests/genie.test.ts | 30 +-- .../src/plugins/jobs/tests/plugin.test.ts | 140 +++++------ .../tests/lakebase-agent-tool.test.ts | 20 +- .../src/plugins/serving/tests/serving.test.ts | 29 +-- packages/appkit/src/testing/fixtures.ts | 130 +++++++++++ packages/appkit/src/testing/index.ts | 5 + packages/appkit/src/testing/test-app.ts | 83 +++++++ packages/appkit/src/testing/test-cache.ts | 90 +++++++ .../appkit/src/testing/test-plugin-context.ts | 112 ++++++++- .../appkit/src/testing/tests/fixtures.test.ts | 184 +++++++++++++++ .../published-surface.integration.test.ts | 4 + .../appkit/src/testing/tests/test-app.test.ts | 87 +++++++ .../src/testing/tests/test-cache.test.ts | 73 ++++++ .../testing/tests/test-plugin-context.test.ts | 212 ++++++++++++++++- 37 files changed, 1760 insertions(+), 1174 deletions(-) create mode 100644 packages/appkit/src/testing/test-app.ts create mode 100644 packages/appkit/src/testing/test-cache.ts create mode 100644 packages/appkit/src/testing/tests/test-app.test.ts create mode 100644 packages/appkit/src/testing/tests/test-cache.test.ts diff --git a/.gitignore b/.gitignore index a623d8837..b70e0c2e8 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,5 @@ coverage internal .isaac/ + +.codex-tmp/ diff --git a/docs/docs/plugins/testing.md b/docs/docs/plugins/testing.md index ca02a5aba..6d06d355b 100644 --- a/docs/docs/plugins/testing.md +++ b/docs/docs/plugins/testing.md @@ -17,7 +17,148 @@ The kit has three entry points plus a set of fixture helpers: - **`expectStream(...).toEmit(...)`** — assert the ordered event types a stream emits. - **Fixtures** — `createMockRequest`, `createMockResponse`, `createMockWorkspaceClient`, `mockServiceContext`, and SQL response builders. -The kit uses [Vitest](https://vitest.dev)'s `vi` for its mocks, so `vitest` is an **optional peer dependency**: you already have it (you're writing Vitest tests), and the kit resolves to your copy rather than bundling a second one. Because it's optional, it is not installed into apps that never import `@databricks/appkit/testing` — production installs stay free of the test framework. Any Vitest v3 or v4 works. +`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 @@ -153,7 +294,7 @@ The harness validates that required resources' **environment variables are prese | 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 @@ -183,9 +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 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. +### Seeding with workspace responses and environment -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()`: +`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"; @@ -197,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: @@ -223,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 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. +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()`. @@ -272,11 +464,12 @@ await expectStream(handler.stream(req), { timeout: 1000 }).toEmit("result"); 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 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 @@ -290,6 +483,17 @@ 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`). @@ -331,9 +535,9 @@ 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. +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 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. +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. ::: 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/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/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/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 2151103eb..9cb9b2036 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(); }); @@ -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); }); diff --git a/packages/appkit/src/plugins/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index bf721feee..111ee0a95 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(); }); @@ -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 0134c09e6..2fd955297 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.integration.test.ts @@ -17,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(), @@ -42,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 = { @@ -246,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( @@ -677,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( @@ -696,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 2ada2ef3d..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(); diff --git a/packages/appkit/src/plugins/jobs/tests/plugin.test.ts b/packages/appkit/src/plugins/jobs/tests/plugin.test.ts index 8933d9eed..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,72 +13,42 @@ import { import { mapParams } from "../params"; import { JobsPlugin, jobs } from "../plugin"; -const { mockClient, jobsApi, mockCacheInstance } = await vi.hoisted( - async () => { - // The testing kit's fake, not a hand-rolled literal: the seven jobs methods, - // `config.host` as a real string, and `config.authenticate` all come for free, - // and any *other* service this plugin grows into resolves instead of throwing. - // Imported inside the hoisted factory because the factory runs before the - // file's own imports are evaluated. - const { createMockWorkspaceClient, 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 mockCacheInstance = { - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - getOrExecute: vi.fn( - async ( - _key: unknown[], - fn: (signal?: AbortSignal) => Promise, - ) => fn(), - ), - generateKey: vi.fn(), - }; +const { 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"), + }; - return { mockClient, jobsApi, mockCacheInstance }; - }, -); + return { mockClient, jobsApi }; +}); vi.mock("../../../workspace-client", async (importOriginal) => { const actual = await importOriginal(); - return { - ...actual, - createWorkspaceClient: () => mockClient, - Context: vi.fn(), - }; -}); - -vi.mock("../../../context", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getWorkspaceClient: vi.fn(() => mockClient), - isInUserContext: vi.fn(() => true), - }; + // Only `Context` — the client itself is injected through ServiceContext. + return { ...actual, Context: vi.fn() }; }); -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>; @@ -86,7 +57,10 @@ describe("JobsPlugin", () => { vi.clearAllMocks(); setupDatabricksEnv(); ServiceContext.reset(); - serviceContextMock = await mockServiceContext(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: mockClient, + userDatabricksClient: mockClient, + }); }); afterEach(() => { @@ -163,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", () => { @@ -597,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; - jobsApi.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"); @@ -934,7 +905,10 @@ describe("injectRoutes", () => { vi.clearAllMocks(); setupDatabricksEnv(); ServiceContext.reset(); - serviceContextMock = await mockServiceContext(); + serviceContextMock = await mockServiceContext({ + serviceDatabricksClient: mockClient, + userDatabricksClient: mockClient, + }); }); afterEach(() => { @@ -1852,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; jobsApi.listRuns.mockImplementation(() => { - throw error; + throw createApiError({ + statusCode: 401, + message: "Unauthorized", + errorCode: "UNAUTHENTICATED", + }); }); const plugin = new JobsPlugin({}); 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/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/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 13f79b5d7..97d93f8b7 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -9,6 +9,7 @@ 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 @@ -340,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. @@ -563,3 +660,36 @@ export function createFailedSQLResponse(errorMessage: string) { statement_id: `stmt-${Date.now()}`, }; } + +/** + * 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 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 797ce6935..165f81ae7 100644 --- a/packages/appkit/src/testing/index.ts +++ b/packages/appkit/src/testing/index.ts @@ -60,6 +60,7 @@ export { type StreamSource, } from "./expect-stream"; export { + createApiError, createFailedSQLResponse, createMockRequest, createMockResponse, @@ -74,6 +75,7 @@ export { setupDatabricksEnv, type TestContextOptions, useServiceContextMock, + withEnv, } from "./fixtures"; export { createMockWorkspaceClient, @@ -82,6 +84,8 @@ export { 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, @@ -90,4 +94,5 @@ export { type RecordedRoute, type RecordedToolCall, type TestPluginContext, + type TestPluginContextOptions, } from "./test-plugin-context"; 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/fixtures.test.ts b/packages/appkit/src/testing/tests/fixtures.test.ts index 599fb8151..9f7a906f8 100644 --- a/packages/appkit/src/testing/tests/fixtures.test.ts +++ b/packages/appkit/src/testing/tests/fixtures.test.ts @@ -4,11 +4,14 @@ 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", () => { @@ -184,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/published-surface.integration.test.ts b/packages/appkit/src/testing/tests/published-surface.integration.test.ts index 536099425..52f5f1925 100644 --- a/packages/appkit/src/testing/tests/published-surface.integration.test.ts +++ b/packages/appkit/src/testing/tests/published-surface.integration.test.ts @@ -77,11 +77,15 @@ describe("@databricks/appkit/testing as a standalone surface", () => { "createMockTelemetry", "createSuccessfulSQLResponse", "createFailedSQLResponse", + "createApiError", "parseSSEResponse", "resetTestCache", "runWithRequestContext", "setupDatabricksEnv", "useServiceContextMock", + "useTestApp", + "useTestCache", + "withEnv", ]; const missing = expected.filter( (name) => 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 d84a52937..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,5 +1,5 @@ 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"; @@ -324,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" }); + }); +}); From 49f699efbd3f776c8a9fc28b70b1fffaad6ba524 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 1 Sep 2026 18:20:19 +0200 Subject: [PATCH 03/12] feat(appkit): migrate analytics to the modular @databricks/sdk-* MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Migrate the analytics stack (SQLWarehouseConnector + type-generator) off the legacy monolithic @databricks/sdk-experimental onto the new modular per-service @databricks/sdk-* SDK (v0.46.0, ESM-only), behind the existing workspace-client facade seam. The two services analytics depends on — warehouses and statementExecution — move together; every other service still routes through the legacy client (mixed state by design). - New packages/shared/src/workspace-client/modular.ts is the sole importer of @databricks/sdk-* (oxlint no-restricted-imports boundary), mirroring legacy.ts. Builds per-service WarehousesClient / StatementExecutionClient; maps wrapper options -> ClientOptions (host scheme-normalization, PAT empty-token guard, profile); stamps process-global client-info (sanitized, best-effort). - Connector + type-generator rewritten to the modular API: method renames (getStatement -> getStatementResult, getStatementResultChunkN -> getResultData), camelCase response model, CallOptions { signal }. - statementExecution relies on a pinned pnpm patch that restores the undocumented Reyden `attachment` field the SDK's unmarshal transform would otherwise strip. - Coerce the SDK's bigint row/byte counts back to number at the connector boundary so INLINE + ARROW_STREAM results stay JSON-serializable (cache / SSE frames). - Read the modular ApiError's `.code` (not only the legacy `.errorCode`) so the arrow disposition/format capability-rejection fallback still fires. Verified against live warehouses (standard + Reyden serverless): JSON and arrow (INLINE attachment + EXTERNAL_LINKS), OBO, warehouse auto-start, and metric views. Full appkit + shared suite green (3877 tests). Signed-off-by: MarioCadenas --- .oxlintrc.json | 7 +- .../api/appkit/Interface.WorkspaceClient.md | 8 +- package.json | 3 + .../connectors/sql-warehouse/arrow-schema.ts | 6 +- .../src/connectors/sql-warehouse/client.ts | 257 ++++++++++-------- .../src/connectors/sql-warehouse/defaults.ts | 14 +- .../sql-warehouse/tests/arrow-schema.test.ts | 26 +- .../sql-warehouse/tests/client.test.ts | 221 +++++++-------- .../sql-warehouse/warehouse-status-emitter.ts | 6 +- .../connectors/tests/sql-warehouse.test.ts | 215 +++++++++++++-- packages/appkit/src/evals/dataset.ts | 2 +- .../appkit/src/evals/tests/dataset.test.ts | 2 +- .../appkit/src/plugins/analytics/analytics.ts | 2 +- .../appkit/src/plugins/analytics/query.ts | 8 +- .../src/plugins/analytics/result-delivery.ts | 6 +- .../tests/analytics.integration.test.ts | 15 +- .../plugins/analytics/tests/analytics.test.ts | 10 +- .../tests/arrow-delivery.integration.test.ts | 6 +- .../plugins/analytics/tests/metric.test.ts | 4 +- .../src/stream/arrow-stream-processor.ts | 24 +- .../tests/arrow-stream-processor.test.ts | 15 +- packages/appkit/src/testing/fixtures.ts | 10 +- .../src/type-generator/statement-result.ts | 58 +++- .../tests/generate-queries.test.ts | 31 ++- .../src/type-generator/tests/index.test.ts | 74 +++-- .../type-generator/tests/mv-registry.test.ts | 52 +++- .../tests/statement-result.test.ts | 32 ++- .../tests/unreachable-warehouse-gate.test.ts | 2 +- .../tests/warehouse-status.test.ts | 10 +- .../src/type-generator/warehouse-status.ts | 6 +- packages/appkit/src/workspace-client/index.ts | 17 +- packages/shared/package.json | 5 + .../shared/src/workspace-client/client.ts | 24 +- packages/shared/src/workspace-client/index.ts | 2 + .../shared/src/workspace-client/modular.ts | 159 +++++++++++ .../workspace-client/tests/modular.test.ts | 113 ++++++++ packages/shared/src/workspace-client/types.ts | 19 +- ...ricks__sdk-statementexecution@0.46.0.patch | 34 +++ pnpm-lock.yaml | 86 ++++++ 39 files changed, 1173 insertions(+), 418 deletions(-) create mode 100644 packages/shared/src/workspace-client/modular.ts create mode 100644 packages/shared/src/workspace-client/tests/modular.test.ts create mode 100644 patches/@databricks__sdk-statementexecution@0.46.0.patch 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/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/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/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/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/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index dc3543be4..9ec29131e 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -1069,7 +1069,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 c69a82e83..0265feeea 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.integration.test.ts @@ -26,7 +26,7 @@ describe("Analytics Plugin Integration", () => { let app: TestApp<[ReturnType]>; /** The SQL mock the analytics route drives, via the harness's client. */ let executeStatement: ReturnType; - let getStatement: ReturnType; + let getStatementResult: ReturnType; beforeAll(async () => { // The harness owns the env setup, the singleton resets, the mock client, the @@ -36,7 +36,10 @@ describe("Analytics Plugin Integration", () => { app.client, "statementExecution.executeStatement", ); - getStatement = getMock(app.client, "statementExecution.getStatement"); + getStatementResult = getMock( + app.client, + "statementExecution.getStatementResult", + ); }); afterAll(async () => { @@ -48,7 +51,7 @@ describe("Analytics Plugin Integration", () => { // Reset drops the built-in canned SUCCEEDED default too, matching the // "script it yourself" semantics this suite relied on before. executeStatement.mockReset(); - getStatement.mockReset(); + getStatementResult.mockReset(); getAppQuerySpy.mockReset(); }); @@ -60,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({ @@ -93,7 +96,7 @@ describe("Analytics Plugin Integration", () => { expect(executeStatement).toHaveBeenCalledWith( expect.objectContaining({ statement: testQuery, - warehouse_id: "test-warehouse-id", + warehouseId: "test-warehouse-id", }), expect.anything(), ); diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index 9cb9b2036..a30e4be8e 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -137,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), ); @@ -205,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), ); @@ -602,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), ); @@ -637,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", }), @@ -1648,7 +1648,7 @@ describe("Analytics Plugin", () => { result: { data: [] }, }), }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ 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 111ee0a95..63cb7ea27 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -700,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), ); @@ -748,7 +748,7 @@ describe("analytics metric route", () => { result: { data: [] }, }), }, - warehouses: { get: warehouseGet, start: vi.fn() }, + warehouses: { getWarehouse: warehouseGet, startWarehouse: vi.fn() }, }, }); const mockReq = createMockRequest({ 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/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 97d93f8b7..eaa8d8eed 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -629,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", })), }, }, @@ -657,7 +657,7 @@ export function createFailedSQLResponse(errorMessage: string) { message: errorMessage, }, }, - statement_id: `stmt-${Date.now()}`, + statementId: `stmt-${Date.now()}`, }; } diff --git a/packages/appkit/src/type-generator/statement-result.ts b/packages/appkit/src/type-generator/statement-result.ts index 7ae091aaf..24620988b 100644 --- a/packages/appkit/src/type-generator/statement-result.ts +++ b/packages/appkit/src/type-generator/statement-result.ts @@ -1,5 +1,5 @@ import { createLogger } from "../logging/logger"; -import type { WorkspaceClient } from "../workspace-client"; +import type { StatementResponse, WorkspaceClient } from "../workspace-client"; import { getErrorMessage } from "./errors"; import type { DatabricksStatementExecutionResponse } from "./types"; @@ -147,6 +147,42 @@ function isFormatRejection( ); } +/** + * Adapt the modular SDK's camelCase {@link StatementResponse} onto the + * type-generator's own snake_case {@link DatabricksStatementExecutionResponse} + * — the shape every downstream DESCRIBE parser (and every mocked test) reads. + * Keeping the boundary here means only this mapper touches the SDK shape; + * {@link normalizeResultRows} and the parsers stay unchanged. `attachment` + * survives thanks to the pinned pnpm patch on `@databricks/sdk-statementexecution`. + */ +function toDescribeResponse( + r: StatementResponse, +): DatabricksStatementExecutionResponse { + return { + statement_id: r.statementId ?? "", + status: { + state: r.status?.state ?? "", + error: r.status?.error + ? { + error_code: r.status.error.errorCode, + message: r.status.error.message, + } + : undefined, + }, + manifest: r.manifest ? { format: r.manifest.format } : undefined, + result: r.result + ? { + // DESCRIBE rows are always string/null cells. Local key stays + // snake_case (`data_array`); value is the SDK's camelCase `dataArray`. + data_array: r.result.dataArray as (string | null)[][] | undefined, + attachment: r.result.attachment, + next_chunk_index: r.result.nextChunkIndex, + next_chunk_internal_link: r.result.nextChunkInternalLink, + } + : undefined, + }; +} + /** * Run a DESCRIBE and return a response whose rows are readable via * `result.data_array`, adapting to the warehouse's result-format capability. @@ -175,15 +211,17 @@ export async function describeAdaptive( let lastError: unknown; for (const format of formats) { try { - const response = (await client.statementExecution.executeStatement({ - statement, - warehouse_id: 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", - format, - disposition: "INLINE", - })) as DatabricksStatementExecutionResponse; + const response = toDescribeResponse( + await client.statementExecution.executeStatement({ + statement, + warehouseId, + // Synchronous wait: without it the call can return PENDING/RUNNING with + // no rows, which downstream misreads as a no-result degrade. + waitTimeout: "30s", + format, + disposition: "INLINE", + }), + ); const normalized = await normalizeResultRows(response); if ( normalized.status?.state === "FAILED" && 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..48a35ed1a 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,9 +85,9 @@ 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 }, }; } @@ -108,7 +111,7 @@ 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 @@ -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..6fa33ed48 100644 --- a/packages/appkit/src/type-generator/tests/index.test.ts +++ b/packages/appkit/src/type-generator/tests/index.test.ts @@ -11,8 +11,37 @@ import { vi, } from "vitest"; +import type { StatementResponse } from "../../workspace-client"; import type { DatabricksStatementExecutionResponse } from "../types"; +/** + * Adapt a local snake_case describe fixture to the modular SDK's camelCase + * `StatementResponse` — the shape the mocked `executeStatement` now returns. + * `describeAdaptive` maps it back to the local shape via `toDescribeResponse`, + * so fixtures stay authored in the type-generator's own domain shape. + */ +function asSdkResponse( + r: DatabricksStatementExecutionResponse, +): StatementResponse { + return { + statementId: r.statement_id, + status: r.status && { + state: r.status.state, + error: r.status.error && { + errorCode: r.status.error.error_code, + message: r.status.error.message, + }, + }, + manifest: r.manifest && { format: r.manifest.format }, + result: r.result && { + dataArray: r.result.data_array, + attachment: r.result.attachment, + nextChunkIndex: r.result.next_chunk_index, + nextChunkInternalLink: r.result.next_chunk_internal_link, + }, + } as unknown as StatementResponse; +} + const mocks = vi.hoisted(() => ({ generateQueriesFromDescribe: vi.fn(), getWarehouseState: vi.fn(), @@ -553,7 +582,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING warehouse: DESCRIBEs run and land full schemas", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -568,7 +597,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"); @@ -582,7 +611,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("blocking + RUNNING: one preflight probe, no start/wait, DESCRIBEs run", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -756,7 +785,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { mocks.getWarehouseState.mockResolvedValue("STOPPED"); mocks.startWarehouse.mockResolvedValue(undefined); mocks.waitUntilRunning.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -1015,7 +1044,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING with the default fetcher: probe and DESCRIBEs share exactly one client", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ @@ -1154,24 +1183,23 @@ describe("generateFromEntryPoint — metric cache section", () => { const outFile = path.join(cacheTestDir, "generated", "analytics.d.ts"); const metricFile = path.join(cacheTestDir, "generated", "metric-views.d.ts"); - const describeResponseFor = ( - measure: string, - ): DatabricksStatementExecutionResponse => ({ - statement_id: "stmt-mock", - status: { state: "SUCCEEDED" }, - result: { - data_array: [ - [ - JSON.stringify({ - columns: [ - { name: measure, type: "DECIMAL(38,2)", is_measure: true }, - { name: "region", type: "STRING", is_measure: false }, - ], - }), + const describeResponseFor = (measure: string): StatementResponse => + asSdkResponse({ + statement_id: "stmt-mock", + status: { state: "SUCCEEDED" }, + result: { + data_array: [ + [ + JSON.stringify({ + columns: [ + { name: measure, type: "DECIMAL(38,2)", is_measure: true }, + { name: "region", type: "STRING", is_measure: false }, + ], + }), + ], ], - ], - }, - }); + }, + }); const writeConfig = ( metricViews: Record< @@ -2081,7 +2109,7 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { }; mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(describeResponse); + mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); await expect( generateFromEntryPoint({ 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..e6d90bb75 100644 --- a/packages/appkit/src/type-generator/tests/mv-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/mv-registry.test.ts @@ -10,6 +10,7 @@ import { afterEach, beforeEach, describe, expect, test } from "vitest"; // imports it from there. import { quoteFqnForSql } from "../../../../shared/src/schemas/metric-fqn"; import { metricSourceSchema } from "../../../../shared/src/schemas/metric-source"; +import type { StatementResponse } from "../../workspace-client"; import { readMetricConfig, resolveMetricConfig } from "../mv-registry/config"; import { createWorkspaceDescribeFetcher, @@ -53,6 +54,35 @@ function mockDescribeResponse( }; } +/** + * Adapt a local snake_case describe fixture to the modular SDK's camelCase + * `StatementResponse` — the shape a mocked `executeStatement` (consumed by the + * real `createWorkspaceDescribeFetcher` → `describeAdaptive`) now returns. + * Direct `syncMetrics(resolution, fetcher)` fixtures stay in the local snake + * shape (they bypass the SDK), so only the executeStatement mocks wrap with this. + */ +function asSdkResponse( + r: DatabricksStatementExecutionResponse, +): StatementResponse { + return { + statementId: r.statement_id, + status: r.status && { + state: r.status.state, + error: r.status.error && { + errorCode: r.status.error.error_code, + message: r.status.error.message, + }, + }, + manifest: r.manifest && { format: r.manifest.format }, + result: r.result && { + dataArray: r.result.data_array, + attachment: r.result.attachment, + nextChunkIndex: r.result.next_chunk_index, + nextChunkInternalLink: r.result.next_chunk_internal_link, + }, + } as unknown as StatementResponse; +} + /** * Real Arrow IPC attachment captured live from dogfood: * DESCRIBE TABLE EXTENDED `appkit_demo`.`public`.`revenue_metrics` AS JSON @@ -326,9 +356,11 @@ describe("resolveMetricConfig — FQN naming (UC-accurate)", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return mockDescribeResponse({ - columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], - }); + return asSdkResponse( + mockDescribeResponse({ + columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], + }), + ); }, }, } as unknown as Parameters[0]; @@ -664,7 +696,7 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return mockDescribeResponse(payload); + return asSdkResponse(mockDescribeResponse(payload)); }, }, } as unknown as Parameters[0]; @@ -680,8 +712,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", @@ -700,13 +732,13 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return { + return asSdkResponse({ statement_id: "stmt-arrow", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, // Only an attachment — no data_array (the bug's trigger condition). result: { attachment: ARROW_ATTACHMENT_B64 }, - } as DatabricksStatementExecutionResponse; + }); }, }, } as unknown as Parameters[0]; @@ -1016,7 +1048,7 @@ describe("syncMetrics", () => { const client = { statementExecution: { executeStatement: async () => - ({ + asSdkResponse({ statement_id: "stmt-chunked", status: { state: "SUCCEEDED" }, manifest: { format: "ARROW_STREAM" }, @@ -1024,7 +1056,7 @@ describe("syncMetrics", () => { attachment: ARROW_ATTACHMENT_B64, next_chunk_index: 1, }, - }) as DatabricksStatementExecutionResponse, + }), }, } as unknown as Parameters[0]; const fetcher = createWorkspaceDescribeFetcher(client, "wh-1"); 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..d545e49ce 100644 --- a/packages/appkit/src/type-generator/tests/statement-result.test.ts +++ b/packages/appkit/src/type-generator/tests/statement-result.test.ts @@ -3,7 +3,10 @@ import path from "node:path"; import { describe, expect, test } from "vitest"; -import type { WorkspaceClient } from "../../workspace-client"; +import type { + StatementResponse, + WorkspaceClient, +} from "../../workspace-client"; import { type DescribeFormatMemo, describeAdaptive, @@ -270,6 +273,31 @@ describe("describeAdaptive", () => { | DatabricksStatementExecutionResponse | Promise; + // Adapt a local snake_case fixture to the modular SDK's camelCase + // StatementResponse — the shape executeStatement now returns; describeAdaptive + // maps it back to the local shape via toDescribeResponse. + function asSdkResponse( + r: DatabricksStatementExecutionResponse, + ): StatementResponse { + return { + statementId: r.statement_id, + status: r.status && { + state: r.status.state, + error: r.status.error && { + errorCode: r.status.error.error_code, + message: r.status.error.message, + }, + }, + manifest: r.manifest && { format: r.manifest.format }, + result: r.result && { + dataArray: r.result.data_array, + attachment: r.result.attachment, + nextChunkIndex: r.result.next_chunk_index, + nextChunkInternalLink: r.result.next_chunk_internal_link, + }, + } as unknown as StatementResponse; + } + // Minimal WorkspaceClient stub: records the formats requested and delegates // each executeStatement to behavior(format), which may resolve or throw. function stubClient(behavior: StubBehavior) { @@ -278,7 +306,7 @@ describe("describeAdaptive", () => { statementExecution: { executeStatement: async (req: { format: string }) => { formats.push(req.format); - return behavior(req.format); + return asSdkResponse(await behavior(req.format)); }, }, } as unknown as WorkspaceClient; 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..ef0b81519 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() }, }), }; }); 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/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/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..a9cdef32b --- /dev/null +++ b/packages/shared/src/workspace-client/modular.ts @@ -0,0 +1,159 @@ +/** + * 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 { newPatCredentials } from "@databricks/sdk-auth/credentials"; +import { addToDefault, setProduct } from "@databricks/sdk-core/clientinfo"; +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`. Mirrors + * `buildLegacyWorkspaceClient`'s auth resolution verbatim, including 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 authenticating as the service + * principal via the default chain (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) { + clientOptions.credentials = newPatCredentials(opts.token); + } else if (opts.profile) { + clientOptions.profileOptions = { profile: opts.profile }; + } + // Neither token nor profile → leave credentials unset so the SDK walks its + // default auth chain (env vars + ~/.databrickscfg), matching the legacy `{}` case. + return clientOptions; +} + +// The modular SDK has no per-client User-Agent option; product/client-info is a +// process-global set once via `setProduct`/`addToDefault` before any client is +// built. The AppKit product/version/userAgentExtra arrive on `opts.clientOptions` +// (from `getClientOptions()`); build-time callers omit them and are left unstamped, +// preserving the legacy behavior where build-time clients carry no AppKit UA. The +// flag latches only once we actually stamp, so a first (unstamped) build-time +// client never blocks a later runtime client from stamping. +let clientInfoStamped = false; + +/** + * Coerce an arbitrary string into a valid client-info segment. The modular SDK + * validates keys as simple tokens and throws `ClientInfoError` on anything else, + * so the legacy product name `@databricks/appkit` (with `@` and `/`) is rejected + * — collapse invalid runs to `-` and trim the ends (`@databricks/appkit` → + * `databricks-appkit`). + */ +function toClientInfoKey(value: string): string { + return value.replace(/[^A-Za-z0-9._-]+/g, "-").replace(/^-+|-+$/g, ""); +} + +function ensureClientInfo(opts: WorkspaceClientOptions): void { + if (clientInfoStamped) { + return; + } + const co = opts.clientOptions; + if (!co?.product || !co?.productVersion) { + return; + } + // User-Agent stamping is best-effort: a value the SDK's client-info validator + // rejects must NEVER break client construction (the legacy SDK stamped the UA + // without validating). On failure the outbound request just carries the SDK's + // default User-Agent. + try { + setProduct(toClientInfoKey(co.product), co.productVersion); + if (co.userAgentExtra) { + for (const [key, value] of Object.entries(co.userAgentExtra)) { + addToDefault(toClientInfoKey(key), String(value)); + } + } + clientInfoStamped = true; + } catch { + clientInfoStamped = true; + } +} + +/** Build a modular Warehouses client from wrapper options. */ +export function buildWarehousesClient( + opts: WorkspaceClientOptions, +): WarehousesClient { + ensureClientInfo(opts); + return new WarehousesClient(mapToClientOptions(opts)); +} + +/** Build a modular Statement Execution client from wrapper options. */ +export function buildStatementExecutionClient( + opts: WorkspaceClientOptions, +): StatementExecutionClient { + ensureClientInfo(opts); + 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..d6d092bcd --- /dev/null +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -0,0 +1,113 @@ +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, productCalls } = vi.hoisted(() => ({ + ctorOpts: [] as Array>, + patTokens: [] as string[], + productCalls: [] as Array<[string, string]>, +})); + +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 }; + }), +})); +vi.mock("@databricks/sdk-core/clientinfo", () => ({ + setProduct: vi.fn((name: string, version: string) => { + // Mirror the real SDK: reject client-info keys that aren't simple tokens. + if (/[^A-Za-z0-9._-]/.test(name)) { + throw new Error(`Invalid key: ${name}.`); + } + productCalls.push([name, version]); + }), + addToDefault: vi.fn(), +})); + +import { buildWarehousesClient } from "../modular"; + +describe("modular mapToClientOptions (via buildWarehousesClient)", () => { + const originalHost = process.env.DATABRICKS_HOST; + + beforeEach(() => { + ctorOpts.length = 0; + patTokens.length = 0; + productCalls.length = 0; + delete process.env.DATABRICKS_HOST; + }); + + afterEach(() => { + if (originalHost === undefined) delete process.env.DATABRICKS_HOST; + else process.env.DATABRICKS_HOST = originalHost; + }); + + 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("client-info: sanitizes an invalid product name (e.g. @databricks/appkit) rather than crashing the client build", () => { + // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` + // (INVALID_KEY), which the legacy SDK accepted. UA stamping must be + // best-effort — a bad product string must never break client construction. + const client = buildWarehousesClient({ + clientOptions: { + product: "@databricks/appkit", + productVersion: "0.64.0", + userAgentExtra: { mode: "dev" }, + }, + } as never); + expect(client).toBeDefined(); + expect(productCalls[0]).toEqual(["databricks-appkit", "0.64.0"]); + }); +}); 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..374003ea4 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: .: @@ -573,9 +578,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 +1955,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 +1971,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 +2649,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 +8712,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 +14262,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 +14290,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 +15511,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 +22128,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 From 8d3b3cba4f29f13a3d858dac0470ce2a99656f69 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 11:40:05 +0200 Subject: [PATCH 04/12] chore: fixup --- packages/appkit/src/plugins/analytics/analytics.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 9ec29131e..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 { From e62b7274a5a7e1c1df7ff1009143486f45fd7666 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 12:32:51 +0200 Subject: [PATCH 05/12] refactor(appkit): camelCase type-generator domain type; drop SDK response mappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the analytics SDK migration: DatabricksStatementExecutionResponse is now camelCase (aligned with the modular SDK's StatementResponse), so the translation shims the migration introduced are no longer needed. - Delete `toDescribeResponse` — `describeAdaptive` narrows the SDK response onto the domain type directly (one cast: the SDK types DESCRIBE cells as JsonValue[][], but for a DESCRIBE they are always string/null). - Delete the three duplicated `asSdkResponse` test helpers; fixtures are now authored in the camelCase domain shape and feed both the SDK-mock and the snake-free parsers directly. - `error_code` stays snake ONLY where it reads the raw Databricks error wire shape (errors.ts auth classification, `{"error_code":...}` JSON bodies) — a network payload, not our type. Internal type only (not exported, not in API docs) — no breaking change. Net -117 LOC. Verified: typecheck, 4355 tests, build (attw + publint), lint, format. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../type-generator/mv-registry/describe.ts | 2 +- .../src/type-generator/query-registry.ts | 10 +- .../src/type-generator/statement-result.ts | 86 +++------- .../tests/generate-queries.test.ts | 6 +- .../src/type-generator/tests/index.test.ts | 110 +++++-------- .../type-generator/tests/mv-registry.test.ts | 113 +++++-------- .../tests/query-registry.test.ts | 28 ++-- .../tests/statement-result.test.ts | 153 ++++++++---------- .../tests/sync-metric-views-types.test.ts | 6 +- .../tests/unreachable-warehouse-gate.test.ts | 4 +- packages/appkit/src/type-generator/types.ts | 27 ++-- 11 files changed, 214 insertions(+), 331 deletions(-) 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 24620988b..9f5f791fe 100644 --- a/packages/appkit/src/type-generator/statement-result.ts +++ b/packages/appkit/src/type-generator/statement-result.ts @@ -1,5 +1,5 @@ import { createLogger } from "../logging/logger"; -import type { StatementResponse, WorkspaceClient } from "../workspace-client"; +import type { WorkspaceClient } from "../workspace-client"; import { getErrorMessage } from "./errors"; import type { DatabricksStatementExecutionResponse } from "./types"; @@ -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) { @@ -147,45 +147,9 @@ function isFormatRejection( ); } -/** - * Adapt the modular SDK's camelCase {@link StatementResponse} onto the - * type-generator's own snake_case {@link DatabricksStatementExecutionResponse} - * — the shape every downstream DESCRIBE parser (and every mocked test) reads. - * Keeping the boundary here means only this mapper touches the SDK shape; - * {@link normalizeResultRows} and the parsers stay unchanged. `attachment` - * survives thanks to the pinned pnpm patch on `@databricks/sdk-statementexecution`. - */ -function toDescribeResponse( - r: StatementResponse, -): DatabricksStatementExecutionResponse { - return { - statement_id: r.statementId ?? "", - status: { - state: r.status?.state ?? "", - error: r.status?.error - ? { - error_code: r.status.error.errorCode, - message: r.status.error.message, - } - : undefined, - }, - manifest: r.manifest ? { format: r.manifest.format } : undefined, - result: r.result - ? { - // DESCRIBE rows are always string/null cells. Local key stays - // snake_case (`data_array`); value is the SDK's camelCase `dataArray`. - data_array: r.result.dataArray as (string | null)[][] | undefined, - attachment: r.result.attachment, - next_chunk_index: r.result.nextChunkIndex, - next_chunk_internal_link: r.result.nextChunkInternalLink, - } - : undefined, - }; -} - /** * 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 @@ -211,23 +175,25 @@ export async function describeAdaptive( let lastError: unknown; for (const format of formats) { try { - const response = toDescribeResponse( - await client.statementExecution.executeStatement({ - statement, - warehouseId, - // Synchronous wait: without it the call can return PENDING/RUNNING with - // no rows, which downstream misreads as a no-result degrade. - waitTimeout: "30s", - format, - disposition: "INLINE", - }), - ); + // 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, + warehouseId, + // Synchronous wait: without it the call can return PENDING/RUNNING with + // no rows, which downstream misreads as a no-result degrade. + waitTimeout: "30s", + format, + disposition: "INLINE", + })) as DatabricksStatementExecutionResponse; const normalized = await normalizeResultRows(response); if ( 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 48a35ed1a..d29b12cd7 100644 --- a/packages/appkit/src/type-generator/tests/generate-queries.test.ts +++ b/packages/appkit/src/type-generator/tests/generate-queries.test.ts @@ -93,7 +93,7 @@ function succeededResult(columns: [string, string, string | null][]) { /** * 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. @@ -114,7 +114,7 @@ async function succeededArrowAttachmentResult( 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 }, }; @@ -160,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"]); diff --git a/packages/appkit/src/type-generator/tests/index.test.ts b/packages/appkit/src/type-generator/tests/index.test.ts index 6fa33ed48..088b2bcf5 100644 --- a/packages/appkit/src/type-generator/tests/index.test.ts +++ b/packages/appkit/src/type-generator/tests/index.test.ts @@ -11,37 +11,8 @@ import { vi, } from "vitest"; -import type { StatementResponse } from "../../workspace-client"; import type { DatabricksStatementExecutionResponse } from "../types"; -/** - * Adapt a local snake_case describe fixture to the modular SDK's camelCase - * `StatementResponse` — the shape the mocked `executeStatement` now returns. - * `describeAdaptive` maps it back to the local shape via `toDescribeResponse`, - * so fixtures stay authored in the type-generator's own domain shape. - */ -function asSdkResponse( - r: DatabricksStatementExecutionResponse, -): StatementResponse { - return { - statementId: r.statement_id, - status: r.status && { - state: r.status.state, - error: r.status.error && { - errorCode: r.status.error.error_code, - message: r.status.error.message, - }, - }, - manifest: r.manifest && { format: r.manifest.format }, - result: r.result && { - dataArray: r.result.data_array, - attachment: r.result.attachment, - nextChunkIndex: r.result.next_chunk_index, - nextChunkInternalLink: r.result.next_chunk_internal_link, - }, - } as unknown as StatementResponse; -} - const mocks = vi.hoisted(() => ({ generateQueriesFromDescribe: vi.fn(), getWarehouseState: vi.fn(), @@ -327,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: [ @@ -489,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" }, }), }), @@ -582,7 +553,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING warehouse: DESCRIBEs run and land full schemas", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -611,7 +582,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("blocking + RUNNING: one preflight probe, no start/wait, DESCRIBEs run", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -737,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" }, }), }), @@ -785,7 +756,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { mocks.getWarehouseState.mockResolvedValue("STOPPED"); mocks.startWarehouse.mockResolvedValue(undefined); mocks.waitUntilRunning.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -922,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(() => {}); @@ -1044,7 +1015,7 @@ describe("generateFromEntryPoint — metric-view emission", () => { test("non-blocking + RUNNING with the default fetcher: probe and DESCRIBEs share exactly one client", async () => { writeMetricConfig(); mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ @@ -1183,23 +1154,24 @@ describe("generateFromEntryPoint — metric cache section", () => { const outFile = path.join(cacheTestDir, "generated", "analytics.d.ts"); const metricFile = path.join(cacheTestDir, "generated", "metric-views.d.ts"); - const describeResponseFor = (measure: string): StatementResponse => - asSdkResponse({ - statement_id: "stmt-mock", - status: { state: "SUCCEEDED" }, - result: { - data_array: [ - [ - JSON.stringify({ - columns: [ - { name: measure, type: "DECIMAL(38,2)", is_measure: true }, - { name: "region", type: "STRING", is_measure: false }, - ], - }), - ], + const describeResponseFor = ( + measure: string, + ): DatabricksStatementExecutionResponse => ({ + statementId: "stmt-mock", + status: { state: "SUCCEEDED" }, + result: { + dataArray: [ + [ + JSON.stringify({ + columns: [ + { name: measure, type: "DECIMAL(38,2)", is_measure: true }, + { name: "region", type: "STRING", is_measure: false }, + ], + }), ], - }, - }); + ], + }, + }); const writeConfig = ( metricViews: Record< @@ -1522,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"); }, @@ -1601,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(() => {}); @@ -1638,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(() => {}); @@ -1655,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( @@ -1693,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(() => {}); @@ -1735,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(() => {}); @@ -2041,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" }, }), }), @@ -2088,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: [ @@ -2109,7 +2081,7 @@ describe("generateFromEntryPoint — anti-clobber for blocking mode", () => { }; mocks.getWarehouseState.mockResolvedValue("RUNNING"); - mocks.executeStatement.mockResolvedValue(asSdkResponse(describeResponse)); + mocks.executeStatement.mockResolvedValue(describeResponse); await expect( generateFromEntryPoint({ 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 e6d90bb75..50c69fe1c 100644 --- a/packages/appkit/src/type-generator/tests/mv-registry.test.ts +++ b/packages/appkit/src/type-generator/tests/mv-registry.test.ts @@ -10,7 +10,6 @@ import { afterEach, beforeEach, describe, expect, test } from "vitest"; // imports it from there. import { quoteFqnForSql } from "../../../../shared/src/schemas/metric-fqn"; import { metricSourceSchema } from "../../../../shared/src/schemas/metric-source"; -import type { StatementResponse } from "../../workspace-client"; import { readMetricConfig, resolveMetricConfig } from "../mv-registry/config"; import { createWorkspaceDescribeFetcher, @@ -46,43 +45,14 @@ 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)]], }, }; } -/** - * Adapt a local snake_case describe fixture to the modular SDK's camelCase - * `StatementResponse` — the shape a mocked `executeStatement` (consumed by the - * real `createWorkspaceDescribeFetcher` → `describeAdaptive`) now returns. - * Direct `syncMetrics(resolution, fetcher)` fixtures stay in the local snake - * shape (they bypass the SDK), so only the executeStatement mocks wrap with this. - */ -function asSdkResponse( - r: DatabricksStatementExecutionResponse, -): StatementResponse { - return { - statementId: r.statement_id, - status: r.status && { - state: r.status.state, - error: r.status.error && { - errorCode: r.status.error.error_code, - message: r.status.error.message, - }, - }, - manifest: r.manifest && { format: r.manifest.format }, - result: r.result && { - dataArray: r.result.data_array, - attachment: r.result.attachment, - nextChunkIndex: r.result.next_chunk_index, - nextChunkInternalLink: r.result.next_chunk_internal_link, - }, - } as unknown as StatementResponse; -} - /** * Real Arrow IPC attachment captured live from dogfood: * DESCRIBE TABLE EXTENDED `appkit_demo`.`public`.`revenue_metrics` AS JSON @@ -356,11 +326,9 @@ describe("resolveMetricConfig — FQN naming (UC-accurate)", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return asSdkResponse( - mockDescribeResponse({ - columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], - }), - ); + return mockDescribeResponse({ + columns: [{ name: "arr", type: "DECIMAL", is_measure: true }], + }); }, }, } as unknown as Parameters[0]; @@ -391,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); @@ -588,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/); @@ -597,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/); }); @@ -607,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/); }); @@ -696,7 +664,7 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return asSdkResponse(mockDescribeResponse(payload)); + return mockDescribeResponse(payload); }, }, } as unknown as Parameters[0]; @@ -723,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. @@ -732,13 +700,13 @@ describe("createWorkspaceDescribeFetcher", () => { statementExecution: { executeStatement: async (req: Record) => { statements.push(req); - return asSdkResponse({ - statement_id: "stmt-arrow", + return { + 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 unknown as Parameters[0]; @@ -747,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. @@ -1037,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 @@ -1047,16 +1015,15 @@ describe("syncMetrics", () => { }); const client = { statementExecution: { - executeStatement: async () => - asSdkResponse({ - statement_id: "stmt-chunked", - status: { state: "SUCCEEDED" }, - manifest: { format: "ARROW_STREAM" }, - result: { - attachment: ARROW_ATTACHMENT_B64, - next_chunk_index: 1, - }, - }), + 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"); @@ -1166,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 })], @@ -1230,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 }, }); @@ -1254,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" } }, }); @@ -1273,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( @@ -1462,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: [ @@ -1634,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 }], @@ -1799,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 d545e49ce..1d1f43a6e 100644 --- a/packages/appkit/src/type-generator/tests/statement-result.test.ts +++ b/packages/appkit/src/type-generator/tests/statement-result.test.ts @@ -3,10 +3,7 @@ import path from "node:path"; import { describe, expect, test } from "vitest"; -import type { - StatementResponse, - WorkspaceClient, -} from "../../workspace-client"; +import type { WorkspaceClient } from "../../workspace-client"; import { type DescribeFormatMemo, describeAdaptive, @@ -46,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 }, @@ -57,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. @@ -69,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 }, @@ -79,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: {}, }; @@ -126,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" }, }; @@ -144,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 }, @@ -163,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!!!" }, @@ -185,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( @@ -252,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 }, @@ -260,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"], ]); }); @@ -273,40 +269,17 @@ describe("describeAdaptive", () => { | DatabricksStatementExecutionResponse | Promise; - // Adapt a local snake_case fixture to the modular SDK's camelCase - // StatementResponse — the shape executeStatement now returns; describeAdaptive - // maps it back to the local shape via toDescribeResponse. - function asSdkResponse( - r: DatabricksStatementExecutionResponse, - ): StatementResponse { - return { - statementId: r.statement_id, - status: r.status && { - state: r.status.state, - error: r.status.error && { - errorCode: r.status.error.error_code, - message: r.status.error.message, - }, - }, - manifest: r.manifest && { format: r.manifest.format }, - result: r.result && { - dataArray: r.result.data_array, - attachment: r.result.attachment, - nextChunkIndex: r.result.next_chunk_index, - nextChunkInternalLink: r.result.next_chunk_internal_link, - }, - } as unknown as StatementResponse; - } - // 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 asSdkResponse(await behavior(req.format)); + return await behavior(req.format); }, }, } as unknown as WorkspaceClient; @@ -316,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 () => { @@ -335,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"]); }); @@ -356,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"]); }); @@ -366,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; @@ -381,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"]); }); @@ -403,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]" }, @@ -438,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", }, }, @@ -461,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"]); }); @@ -474,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", }, @@ -496,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 ef0b81519..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 @@ -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/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; }; } From 20a97d6c205aa077116e02d8265a80814b5bcf73 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 15:22:46 +0200 Subject: [PATCH 06/12] fix(appkit): point #540 never-crash mock at the modular warehouse API Rebasing the analytics SDK migration onto #540 (createTestApp + never-crash mock client) surfaced that #540's mock hardcodes the LEGACY warehouse method names. Point them at the modular API so warehouse-readiness calls resolve: - mock-workspace-client.ts DEFAULT_RESPONSES: warehouses.get/start -> getWarehouse/startWarehouse. - mock-workspace-client.test.ts: never-crash table + convergence asserts. (fixtures.ts + analytics.integration.test.ts adaptations landed inside the migration commit during rebase conflict resolution.) Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/testing/mock-workspace-client.ts | 9 +++++---- .../src/testing/tests/mock-workspace-client.test.ts | 10 ++++++---- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/packages/appkit/src/testing/mock-workspace-client.ts b/packages/appkit/src/testing/mock-workspace-client.ts index 7b9298379..59ec97635 100644 --- a/packages/appkit/src/testing/mock-workspace-client.ts +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -43,8 +43,9 @@ export type MockWorkspaceClient = WorkspaceClient; /** * Applied beneath caller-supplied `responses`. * - * `statementExecution.executeStatement`, `warehouses.get` and `warehouses.start` - * must stay byte-identical to the old `fixtures.ts` values — suites reach them + * `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. @@ -54,8 +55,8 @@ const DEFAULT_RESPONSES: Record = { status: { state: "SUCCEEDED" }, result: { data: [] }, }, - "warehouses.get": { state: "RUNNING" }, - "warehouses.start": undefined, + "warehouses.getWarehouse": { state: "RUNNING" }, + "warehouses.startWarehouse": undefined, "currentUser.me": { id: "test-service-user", userName: "test-service-user", diff --git a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts index a2cbeda1a..5ec90a0d9 100644 --- a/packages/appkit/src/testing/tests/mock-workspace-client.test.ts +++ b/packages/appkit/src/testing/tests/mock-workspace-client.test.ts @@ -24,8 +24,8 @@ describe("createMockWorkspaceClient", () => { ["genie", "getMessage", undefined], ["jobs", "getRun", undefined], ["servingEndpoints", "get", undefined], - ["warehouses", "get", { state: "RUNNING" }], - ["warehouses", "start", undefined], + ["warehouses", "getWarehouse", { state: "RUNNING" }], + ["warehouses", "startWarehouse", undefined], ["statementExecution", "executeStatement", SUCCEEDED], ["currentUser", "me", TEST_USER], ])("%s.%s resolves its default", async (service, method, expected) => { @@ -248,11 +248,13 @@ describe("createMockWorkspaceClient", () => { await expect( client.statementExecution.executeStatement({} as never), ).resolves.toEqual(SUCCEEDED); - await expect(client.warehouses.get({} as never)).resolves.toEqual({ + await expect( + client.warehouses.getWarehouse({} as never), + ).resolves.toEqual({ state: "RUNNING", }); await expect( - client.warehouses.start({} as never), + client.warehouses.startWarehouse({} as never), ).resolves.toBeUndefined(); }); From 829b8783643953ac44a42d0883ef2b17af1d45c0 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 15:51:53 +0200 Subject: [PATCH 07/12] fix(appkit): declare modular @databricks/sdk-* packages as dependencies appkit bundles `shared` inline, whose code imports the modular @databricks/sdk-* packages, but only @databricks/sdk-experimental was declared. A published `@databricks/appkit` install (npm, incl. the Databricks Apps runtime) would therefore fail at runtime with "Cannot find module @databricks/sdk-statementexecution" on the analytics path. Declare sdk-auth/core/options/warehouses/statementexecution at 0.46.0 (mirroring shared) and exempt them from knip's unused check like sdk-experimental (they are imported only by the inlined `shared` workspace, which knip does not analyze). Known follow-up: the pnpm patch restoring Reyden's stripped `attachment` field is a workspace-only mechanism, does not reach consumers, and cannot be bundled under tsdown `unbundle` mode. Standard warehouses (dataArray + external_links) are unaffected; Reyden INLINE+ARROW_STREAM needs the upstream SDK `attachment` fix. See SDK_MIGRATION_GAPS.md. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- knip.json | 10 +++++++++- packages/appkit/package.json | 5 +++++ 2 files changed, 14 insertions(+), 1 deletion(-) 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/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", From 8392ee3cce07130f40d0156c3577e7428b40e74d Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 18:41:24 +0200 Subject: [PATCH 08/12] fix(shared): resolve service-principal credentials from env in modular SDK client The modular @databricks/sdk-* default credential chain resolves auth only from a ~/.databrickscfg profile and reads no DATABRICKS_* env vars, unlike the legacy sdk-experimental. The Databricks Apps runtime injects the app's service-principal credentials via env vars only (no config file), so a deployed app built a modular client with no credentials and every request failed ("Warehouse readiness check failed"). Resolve the service principal from the environment in mapToClientOptions when no explicit token/profile is given: M2M from DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET (what Apps injects), else PAT from DATABRICKS_TOKEN, else fall through to the profile default chain for local dev. The explicit token path (asUser OBO) is unchanged and still guarded by token !== undefined so an empty/invalid OBO token fails loudly instead of silently falling through to the service principal. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 44 ++++++++-- .../workspace-client/tests/modular.test.ts | 84 +++++++++++++++++-- 2 files changed, 114 insertions(+), 14 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index a9cdef32b..2714b8dbd 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -16,7 +16,10 @@ * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. */ -import { newPatCredentials } from "@databricks/sdk-auth/credentials"; +import { + newM2mCredentials, + newPatCredentials, +} from "@databricks/sdk-auth/credentials"; import { addToDefault, setProduct } from "@databricks/sdk-core/clientinfo"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; @@ -37,12 +40,13 @@ function normalizeHost(host: string | undefined): string | undefined { } /** - * Map wrapper options onto the modular SDK's `ClientOptions`. Mirrors - * `buildLegacyWorkspaceClient`'s auth resolution verbatim, including 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 authenticating as the service - * principal via the default chain (which would be an OBO privilege escalation). + * 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 = {}; @@ -57,12 +61,34 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { 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 way the legacy SDK did. The modular SDK's default auth + // chain resolves ONLY from a `~/.databrickscfg` profile — it reads no + // `DATABRICKS_*` env vars — so on the Databricks Apps runtime (which injects + // the app's SP credentials via env, with no config file) it would find no + // credentials and every request would fail. Resolve them here instead: + // M2M (client id + secret, what Apps injects) first, then a PAT, else fall + // through to the default chain for local dev with a config file. + 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`). } - // Neither token nor profile → leave credentials unset so the SDK walks its - // default auth chain (env vars + ~/.databrickscfg), matching the legacy `{}` case. return clientOptions; } diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index d6d092bcd..07af9ef07 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,9 +3,10 @@ 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, productCalls } = vi.hoisted(() => ({ +const { ctorOpts, patTokens, m2mOpts, productCalls } = vi.hoisted(() => ({ ctorOpts: [] as Array>, patTokens: [] as string[], + m2mOpts: [] as Array>, productCalls: [] as Array<[string, string]>, })); @@ -23,6 +24,10 @@ vi.mock("@databricks/sdk-auth/credentials", () => ({ patTokens.push(token); return { kind: "pat", token }; }), + newM2mCredentials: vi.fn((opts: Record) => { + m2mOpts.push(opts); + return { kind: "m2m", ...opts }; + }), })); vi.mock("@databricks/sdk-core/clientinfo", () => ({ setProduct: vi.fn((name: string, version: string) => { @@ -38,18 +43,32 @@ vi.mock("@databricks/sdk-core/clientinfo", () => ({ import { buildWarehousesClient } from "../modular"; describe("modular mapToClientOptions (via buildWarehousesClient)", () => { - const originalHost = process.env.DATABRICKS_HOST; + // 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; productCalls.length = 0; - delete process.env.DATABRICKS_HOST; + for (const key of AUTH_ENV) { + originalEnv[key] = process.env[key]; + delete process.env[key]; + } }); afterEach(() => { - if (originalHost === undefined) delete process.env.DATABRICKS_HOST; - else process.env.DATABRICKS_HOST = originalHost; + 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", () => { @@ -96,6 +115,61 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { 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 this way + // (env only, no config file). The modular SDK's default chain reads no env, + // so we must map them to M2M credentials ourselves. + 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("client-info: sanitizes an invalid product name (e.g. @databricks/appkit) rather than crashing the client build", () => { // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` // (INVALID_KEY), which the legacy SDK accepted. UA stamping must be From a4829765f6d53b2fe22e35280740677e2720b842 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 14 Sep 2026 18:48:15 +0200 Subject: [PATCH 09/12] chore(deps): sync pnpm-lock with appkit modular SDK dependencies The 5 modular @databricks/sdk-* deps were added to packages/appkit/package.json but the appkit importer in pnpm-lock.yaml was never regenerated, so CI's `pnpm install --frozen-lockfile` failed with ERR_PNPM_OUTDATED_LOCKFILE. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- pnpm-lock.yaml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 374003ea4..f14b66053 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -263,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) From 647b566c2eb44cbffadfb877825a897c05a977bc Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 15 Sep 2026 14:49:53 +0200 Subject: [PATCH 10/12] fix(shared): preserve the @databricks/appkit User-Agent on modular clients MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The modular SDK's client-info `setProduct` validates the product as a token and rejects `@databricks/appkit` (the `@`/`/`), so the migrated clients sent a sanitized `databricks-appkit` User-Agent. Databricks-side dashboards filter AppKit's analytics/warehouse traffic on the literal `@databricks/appkit` UA, so that silently dropped AppKit out of them. Set the User-Agent on a per-client `ClientOptions.httpClient` transport wrapper (`buildHttpClient`) that prepends `@databricks/appkit/` (+ userAgentExtra) to each request, replacing the process-global `setProduct` sanitization/latch. This keeps the exact legacy string and ships inside appkit's bundled dist, so it reaches deployed apps — a pnpm patch of the validator would not (pnpm patches apply only at workspace install, not to the npm-installed tarball). Also corrects the service-principal auth comments: the SDK's default chain does read DATABRICKS_* env; the deployed failure was its M2M strategy using the raw scheme-less DATABRICKS_HOST for OAuth discovery, which the normalized-host resolution already works around. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../shared/src/workspace-client/modular.ts | 95 ++++++++++--------- .../workspace-client/tests/modular.test.ts | 81 ++++++++++++---- 2 files changed, 109 insertions(+), 67 deletions(-) diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 2714b8dbd..44a31fbec 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -20,7 +20,7 @@ import { newM2mCredentials, newPatCredentials, } from "@databricks/sdk-auth/credentials"; -import { addToDefault, setProduct } from "@databricks/sdk-core/clientinfo"; +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"; @@ -67,13 +67,16 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { clientOptions.profileOptions = { profile: opts.profile }; } else { // No token, no profile: authenticate as the service principal from the - // environment, the way the legacy SDK did. The modular SDK's default auth - // chain resolves ONLY from a `~/.databrickscfg` profile — it reads no - // `DATABRICKS_*` env vars — so on the Databricks Apps runtime (which injects - // the app's SP credentials via env, with no config file) it would find no - // credentials and every request would fail. Resolve them here instead: - // M2M (client id + secret, what Apps injects) first, then a PAT, else fall - // through to the default chain for local dev with a config file. + // 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; @@ -89,59 +92,60 @@ function mapToClientOptions(opts: WorkspaceClientOptions): ClientOptions { // 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; } -// The modular SDK has no per-client User-Agent option; product/client-info is a -// process-global set once via `setProduct`/`addToDefault` before any client is -// built. The AppKit product/version/userAgentExtra arrive on `opts.clientOptions` -// (from `getClientOptions()`); build-time callers omit them and are left unstamped, -// preserving the legacy behavior where build-time clients carry no AppKit UA. The -// flag latches only once we actually stamp, so a first (unstamped) build-time -// client never blocks a later runtime client from stamping. -let clientInfoStamped = false; - /** - * Coerce an arbitrary string into a valid client-info segment. The modular SDK - * validates keys as simple tokens and throws `ClientInfoError` on anything else, - * so the legacy product name `@databricks/appkit` (with `@` and `/`) is rejected - * — collapse invalid runs to `-` and trim the ends (`@databricks/appkit` → - * `databricks-appkit`). + * 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 toClientInfoKey(value: string): string { - return value.replace(/[^A-Za-z0-9._-]+/g, "-").replace(/^-+|-+$/g, ""); -} - -function ensureClientInfo(opts: WorkspaceClientOptions): void { - if (clientInfoStamped) { - return; - } +function buildHttpClient(opts: WorkspaceClientOptions): HttpClient | undefined { const co = opts.clientOptions; if (!co?.product || !co?.productVersion) { - return; + return undefined; } - // User-Agent stamping is best-effort: a value the SDK's client-info validator - // rejects must NEVER break client construction (the legacy SDK stamped the UA - // without validating). On failure the outbound request just carries the SDK's - // default User-Agent. - try { - setProduct(toClientInfoKey(co.product), co.productVersion); - if (co.userAgentExtra) { - for (const [key, value] of Object.entries(co.userAgentExtra)) { - addToDefault(toClientInfoKey(key), String(value)); - } + const segments = [`${co.product}/${co.productVersion}`]; + if (co.userAgentExtra) { + for (const [key, value] of Object.entries(co.userAgentExtra)) { + segments.push(`${key}/${String(value)}`); } - clientInfoStamped = true; - } catch { - clientInfoStamped = true; } + 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 { - ensureClientInfo(opts); return new WarehousesClient(mapToClientOptions(opts)); } @@ -149,7 +153,6 @@ export function buildWarehousesClient( export function buildStatementExecutionClient( opts: WorkspaceClientOptions, ): StatementExecutionClient { - ensureClientInfo(opts); return new StatementExecutionClient(mapToClientOptions(opts)); } diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 07af9ef07..51814a293 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -3,11 +3,10 @@ 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, productCalls } = vi.hoisted(() => ({ +const { ctorOpts, patTokens, m2mOpts } = vi.hoisted(() => ({ ctorOpts: [] as Array>, patTokens: [] as string[], m2mOpts: [] as Array>, - productCalls: [] as Array<[string, string]>, })); vi.mock("@databricks/sdk-warehouses/v1", () => ({ @@ -29,19 +28,42 @@ vi.mock("@databricks/sdk-auth/credentials", () => ({ return { kind: "m2m", ...opts }; }), })); -vi.mock("@databricks/sdk-core/clientinfo", () => ({ - setProduct: vi.fn((name: string, version: string) => { - // Mirror the real SDK: reject client-info keys that aren't simple tokens. - if (/[^A-Za-z0-9._-]/.test(name)) { - throw new Error(`Invalid key: ${name}.`); - } - productCalls.push([name, version]); - }), - addToDefault: vi.fn(), +// 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. @@ -57,7 +79,6 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { ctorOpts.length = 0; patTokens.length = 0; m2mOpts.length = 0; - productCalls.length = 0; for (const key of AUTH_ENV) { originalEnv[key] = process.env[key]; delete process.env[key]; @@ -116,9 +137,10 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { }); test("service-principal by default: DATABRICKS_CLIENT_ID/SECRET + host env → M2M creds", () => { - // The Databricks Apps runtime injects the app's SP credentials this way - // (env only, no config file). The modular SDK's default chain reads no env, - // so we must map them to M2M credentials ourselves. + // 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"; @@ -170,18 +192,35 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(ctorOpts[0].credentials).toBeUndefined(); }); - test("client-info: sanitizes an invalid product name (e.g. @databricks/appkit) rather than crashing the client build", () => { + test("User-Agent: prepends the exact @databricks/appkit product segment (dashboards match on it)", async () => { // Regression: the modular SDK's `setProduct` rejects `@databricks/appkit` - // (INVALID_KEY), which the legacy SDK accepted. UA stamping must be - // best-effort — a bad product string must never break client construction. - const client = buildWarehousesClient({ + // (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); - expect(client).toBeDefined(); - expect(productCalls[0]).toEqual(["databricks-appkit", "0.64.0"]); + // 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(); }); }); From 04d0d935b1e2e905f6b5a5eeb77223211c1a6cd3 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 15 Sep 2026 18:35:22 +0200 Subject: [PATCH 11/12] fix(appkit): bundle patched statementexecution to reach deployed apps The Reyden `attachment` fix is a pnpm patch, which applies only at this monorepo's install and does NOT travel through a normal `npm install`. A deployed app resolved the unpatched registry copy and lost `attachment` on the INLINE+ARROW_STREAM path (masked by the EXTERNAL_LINKS fallback where supported, broken on INLINE-only Reyden warehouses). Ship the patched copy inside appkit's tarball via npm `bundledDependencies`: `dist-appkit.ts` copies the pnpm-patched package into the tarball's node_modules and declares it bundled, so the consumer (including the Databricks Apps runtime's npm) resolves appkit's imports to the patched copy. Verified: a consumer install resolves the nested patched copy with no unpatched top-level copy. Extensible via the BUNDLED_PATCHED_PACKAGES list. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- tools/dist-appkit.ts | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tools/dist-appkit.ts b/tools/dist-appkit.ts index fee75f96a..1d6aaf10b 100644 --- a/tools/dist-appkit.ts +++ b/tools/dist-appkit.ts @@ -23,6 +23,18 @@ 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"]; + if (prerelease) { pkg.version = `${pkg.version}-pr.${prerelease}`; } @@ -64,6 +76,30 @@ 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) { + 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 }); From ebd530306f910e7524b7f031e7f7489df899bc91 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 15 Sep 2026 19:39:56 +0200 Subject: [PATCH 12/12] fix(appkit): bundle patched deps only for packages that declare them dist-appkit.ts builds both the appkit and appkit-ui tarballs. The bundledDependencies step hard-failed on appkit-ui, which does not depend on @databricks/sdk-statementexecution. Gate bundling on the tarball package's own declared dependencies (captured before the CLI-dependency merge), so appkit-ui bundles nothing and appkit still bundles the patched statementexecution. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- tools/dist-appkit.ts | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tools/dist-appkit.ts b/tools/dist-appkit.ts index 1d6aaf10b..83bc975a0 100644 --- a/tools/dist-appkit.ts +++ b/tools/dist-appkit.ts @@ -35,6 +35,12 @@ const WORKSPACE_PACKAGE_REPLACEMENTS = ["@databricks/lakebase"]; // 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}`; } @@ -84,6 +90,10 @@ Object.assign(pkg.dependencies, CLI_DEPENDENCIES); // 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(