From 0f029045fa21e04da237158963bd084b4d840a95 Mon Sep 17 00:00:00 2001 From: "@rugpanov" Date: Mon, 17 Aug 2026 16:56:50 +0200 Subject: [PATCH] test(e2e): cover uv setup-local on a serverless target *Why*: * The uv-native "Setup python environment" flow (which shells out to `databricks environments setup-local`) has no end-to-end coverage; the bundled CLI now ships the subcommand, so the full flow is testable. * Guards the main user journey against a real serverless target: initialization, that databricks-connect imports and runs, and that a re-run is idempotent (stays ready, no drift). *What:* * Add `setup_local.ucws.e2e.ts` (serverless/UC shard): opt into the uv feature at Workspace scope, attach serverless + pick the recommended version, run setup-local, verify the `.venv` interpreter and the "Python environment ready" config-view row, run a file with Databricks Connect against the provisioned env, and re-run to assert idempotency. * Assertions gate on filesystem/tree ground truth (not transient toasts) with serverless-cold-start budgets, following the serverless DBConnect spec's approach; the re-run is invoked by command id since it is palette-hidden. *Verification:* * `eslint` and `prettier -c` pass on the new spec. * The spec is auto-discovered by `list_integration_tests.ts` and tagged `ucws` (routed to the serverless-capable workspace); it runs in the eng-dev-ecosystem integration workflow. Co-authored-by: Isaac --- .../src/test/e2e/setup_local.ucws.e2e.ts | 310 ++++++++++++++++++ 1 file changed, 310 insertions(+) create mode 100644 packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts diff --git a/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts b/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts new file mode 100644 index 000000000..25dd21c9a --- /dev/null +++ b/packages/databricks-vscode/src/test/e2e/setup_local.ucws.e2e.ts @@ -0,0 +1,310 @@ +import path from "node:path"; +import * as fs from "fs/promises"; +import assert from "node:assert"; +import {CustomTreeSection, InputBox} from "wdio-vscode-service"; +import { + dismissNotifications, + executeCommandWhenAvailable, + getViewSection, + openFile, + waitForLogin, + waitForNotification, +} from "./utils/commonUtils.ts"; +import { + getBasicBundleConfig, + writeRootBundleConfig, +} from "./utils/dabsFixtures.ts"; + +// The uv-native setup entry is opt-in: it unlocks only when this feature id is +// present in `databricks.experiments.optInto` (PYTHON_SETUP_FEATURE_ID). We set +// it at Workspace scope from the test so it stays contained to this project's +// folder and never reroutes another spec's "Setup python environment" command. +const PYTHON_SETUP_FEATURE_ID = "environment.pythonSetup"; + +// Ground-truth text the flow surfaces on success — the config-view row label +// (persistent) and the completion toast (transient). We assert on the row. +const READY_LABEL = "Python environment ready"; +const DRIFTED_LABEL = "out of sync"; + +// uv writes the interpreter into `.venv` at the project root. Resolve it the +// same way the extension does (platform-specific), so the fs ground-truth check +// works on the Windows shard too. +function venvPython(projectDir: string): string { + return process.platform === "win32" + ? path.join(projectDir, ".venv", "Scripts", "python.exe") + : path.join(projectDir, ".venv", "bin", "python"); +} + +async function fileExists(p: string): Promise { + try { + await fs.access(p); + return true; + } catch { + return false; + } +} + +/** + * Wait for the quick-input widget to appear, WITHOUT gating on its progress bar + * clearing (which `waitForInput` does). The compute picker keeps a progress bar + * up while it lazily loads the cluster list — slow enough on the Windows shard to + * blow `waitForInput`'s 10s budget — yet the static "Serverless" item is + * selectable the whole time. The version picker likewise only appears after an + * I/O-bound version resolution, so allow it a generous window to show. + */ +async function waitForQuickInput(timeoutMs = 60_000): Promise { + const workbench = await browser.getWorkbench(); + return new InputBox(workbench.locatorMap).wait(timeoutMs); +} + +/** + * Wait for uv to finish provisioning: the ground truth is the `.venv` + * interpreter existing on disk. Provisioning first auto-installs uv (CI runners + * lack it), then shells out to uv for a Python download + databricks-connect + * install against a cold serverless target, so this is minutes, not seconds — + * hence the generous budget, mirroring the serverless DBConnect specs. We gate on + * the file rather than the completion toast, which can expire before a poll sees + * it on the slow Windows shard. + */ +async function waitForVenvInterpreter(projectDir: string, timeoutMs = 420_000) { + const python = venvPython(projectDir); + await browser.waitUntil(async () => fileExists(python), { + timeout: timeoutMs, + interval: 2000, + timeoutMsg: `uv did not create the .venv interpreter at "${python}"`, + }); +} + +/** + * Assert the config-view Python-environment row reports the ready end-state and + * is not showing the out-of-sync (drift) state — the persistent observable of a + * successful setup, more reliable than the transient toast. + * + * In the uv-native flow the entry is a TOP-LEVEL row (label "Python environment + * ready" / "…is out of sync"), not a child of a "Python Environment" group — the + * group wrapper exists only for the legacy checklist (see + * `EnvironmentComponent.getRoot`). So scan the CONFIGURATION section's visible + * items directly rather than opening a group that isn't there. + */ +async function waitForPythonEnvReady(timeoutMs = 60_000) { + await browser.waitUntil( + async () => { + const section = (await getViewSection("CONFIGURATION")) as + | CustomTreeSection + | undefined; + if (!section) { + return false; + } + for (const item of await section.getVisibleItems()) { + const label = await item.getLabel(); + console.log("CONFIGURATION item label:", label); + if (label.includes(DRIFTED_LABEL)) { + return false; + } + if (label.includes(READY_LABEL)) { + return true; + } + } + return false; + }, + { + timeout: timeoutMs, + interval: 2000, + timeoutMsg: `Config view never reported "${READY_LABEL}"`, + } + ); +} + +/** + * Poll for a DBConnect run's output file to contain the expected text, then + * remove it so a later run can't pass on a stale file. + */ +async function checkOutputFile( + filePath: string, + expected: string, + timeoutMs = 180_000 +) { + await browser.waitUntil( + async () => { + try { + const content = await fs.readFile(filePath, "utf-8"); + return content.includes(expected); + } catch { + return false; + } + }, + { + timeout: timeoutMs, + interval: 2000, + timeoutMsg: `Output file "${filePath}" did not contain "${expected}" within ${timeoutMs}ms`, + } + ); + await fs.rm(filePath); +} + +describe("Set up local Python environment (uv) on serverless", async function () { + let projectDir: string; + this.timeout(15 * 60 * 1000); + + before(async () => { + assert(process.env.WORKSPACE_PATH, "WORKSPACE_PATH doesn't exist"); + projectDir = process.env.WORKSPACE_PATH; + + // A DBConnect entrypoint that runs a query and writes the result, so the + // uv-provisioned databricks-connect can be verified end-to-end. Mirrors + // the layout the serverless DBConnect spec uses (lib at the root, the run + // file under nested/, output written to the run's cwd = projectDir). + await fs.writeFile( + path.join(projectDir, "lib.py"), + `def func(spark):\treturn spark.sql('SELECT "hello world"')` + ); + const nestedDir = path.join(projectDir, "nested"); + await fs.mkdir(nestedDir, {recursive: true}); + await fs.writeFile( + path.join(nestedDir, "hello.py"), + [ + `from lib import func`, + `import os`, + `df = func(spark).toPandas()`, + `df.to_json(os.path.join(os.getcwd(), "file-output.json"))`, + ].join("\n") + ); + + // Serverless project: no top-level cluster_id, so compute resolves to the + // serverless target the ticket calls for. + await writeRootBundleConfig( + getBasicBundleConfig({}, false), + projectDir + ); + }); + + it("should wait for connection", async () => { + await waitForLogin("DEFAULT"); + await dismissNotifications(); + const workbench = await driver.getWorkbench(); + await workbench.getEditorView().closeAllEditors(); + }); + + it("should opt into uv setup and select a serverless version", async () => { + // Opt in at Workspace scope. `isPythonSetupEnabled()` reads this config + // live, so the setup command routes to the uv flow on the next click + // without a window reload. + await browser.executeWorkbench(async (vscode, featureId) => { + await vscode.workspace + .getConfiguration("databricks.experiments") + .update( + "optInto", + [featureId], + vscode.ConfigurationTarget.Workspace + ); + }, PYTHON_SETUP_FEATURE_ID); + + // setup-local requires uv, and the CI runners do not ship it, so preflight + // would fail fast with E_UV_MISSING. Set the CLI's documented auto-install + // opt-in in the extension host's env — the extension spawns setup-local + // with `{...process.env}`, so the CLI installs uv itself instead of aborting. + await browser.executeWorkbench(async () => { + // eslint-disable-next-line @typescript-eslint/naming-convention + process.env.DATABRICKS_LOCALENV_AUTO_INSTALL_UV = "1"; + }); + + // Attach serverless. With the feature opted in, selecting serverless + // prompts for (and persists) the environment version up front, so the + // subsequent setup run resolves compute without re-prompting. + await executeCommandWhenAvailable("Databricks: Configure compute"); + const computeInput = await waitForQuickInput(); + await computeInput.selectQuickPick("Serverless"); + + // The version picker lists candidates best-first (the recommended one is + // starred at index 0). Take the recommendation rather than pinning a + // version string, which drifts as new serverless versions ship. + const versionInput = await waitForQuickInput(); + await versionInput.selectQuickPick(0); + }); + + it("should set up the environment with uv (setup-local)", async () => { + // The router command; opted in + a clean (uv-suitable) project routes it + // to the uv flow, which shells out to `databricks environments + // setup-local`. Compute + version are already resolved, so no prompt. + await executeCommandWhenAvailable( + "Databricks: Setup python environment" + ); + + await waitForVenvInterpreter(projectDir); + + // The completion toast is informational and can expire before a poll + // catches it; the persistent config-view row is the real end-state. + try { + await waitForNotification(READY_LABEL, undefined, 15_000); + } catch (e) { + console.log( + `"${READY_LABEL}" toast not observed; relying on the config-view row.`, + e + ); + } + + await browser.executeWorkbench(async (vscode) => { + await vscode.commands.executeCommand( + "databricks.environment.refresh" + ); + }); + await waitForPythonEnvReady(); + }); + + it("should run a python file with Databricks Connect from the uv env", async () => { + // Verifies the uv-provisioned databricks-connect imports and runs against + // serverless — no manual install fallback, so a broken provision fails + // the test rather than being papered over. + await openFile("hello.py"); + await executeCommandWhenAvailable( + "Databricks: Run current file with Databricks Connect" + ); + await checkOutputFile( + path.join(projectDir, "file-output.json"), + "hello world" + ); + }); + + it("should be idempotent on re-run (stays ready, no drift)", async () => { + // The re-run affordance is palette-hidden (when:false), so invoke it by + // id. A no-change re-run must leave the environment ready and not flip the + // row to the out-of-sync (drift) state; the CLI's own diskMutated=false is + // the underlying idempotency guarantee, covered at the CLI/unit level. + const python = venvPython(projectDir); + assert(await fileExists(python), "expected .venv from the setup step"); + + // Clear the setup step's success toast (it is persistent — it carries a + // "View Details" button) so the one awaited below is the RE-RUN's, not a + // stale match. + await dismissNotifications(); + + // executeCommand resolves when the (re-entrancy-guarded) run settles; a + // re-run over an already-provisioned env is a warm uv sync, so it returns + // quickly rather than paying another cold provision. + await browser.executeWorkbench(async (vscode) => { + await vscode.commands.executeCommand( + "databricks.environment.rerunPythonEnv" + ); + }); + + // A fresh "Python environment ready" toast is the re-run-SPECIFIC success + // signal. A failed run never clears readiness (readyRoots/state persist) + // and never removes the .venv, so the interpreter and the ready row alone + // can't tell a successful re-run from a failed one — but a failure shows + // an error toast instead of this one, so requiring it fails the test on a + // broken re-run. + await waitForNotification(READY_LABEL); + + // End-state is intact: interpreter present, row ready, not drifted. + assert( + await fileExists(python), + ".venv interpreter vanished after re-run" + ); + await browser.executeWorkbench(async (vscode) => { + await vscode.commands.executeCommand( + "databricks.environment.refresh" + ); + }); + await waitForPythonEnvReady(); + }); +});