ci: build the e2e virtualenv once per worker and re-enable test-vscode-e2e

.github/workflows/pr.yaml

@@ -17,6 +17,7 @@ jobs:
       python: ${{ steps.filter.outputs.python }}
       client: ${{ steps.filter.outputs.client }}
       vscode: ${{ steps.filter.outputs.vscode }}
+      lsp: ${{ steps.filter.outputs.lsp }}
       ci: ${{ steps.filter.outputs.ci }}
     steps:
       - uses: actions/checkout@v7
@@ -37,6 +38,8 @@ jobs:
               - 'web/client/**'
             vscode:
               - 'vscode/**'
+            lsp:
+              - 'sqlmesh/lsp/**'
             ci:
               - '.github/**'
               - 'Makefile'
@@ -427,10 +430,12 @@ jobs:
       - name: Run CI
         run: pnpm run ci
   test-vscode-e2e:
-    runs-on:
-      labels: [ubuntu-2204-8]
-    # As at 2026-01-12 this job flakes 100% of the time. It needs investigation
-    if: false
+    needs: changes
+    if:
+      needs.changes.outputs.vscode == 'true' || needs.changes.outputs.lsp ==
+      'true' || needs.changes.outputs.ci == 'true' || github.ref ==
+      'refs/heads/main'
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v7
       - uses: actions/setup-node@v7
@@ -453,7 +458,9 @@ jobs:
           source .venv/bin/activate
           make install-dev
       - name: Install code-server
-        run: curl -fsSL https://code-server.dev/install.sh | sh
+        # Pinned so that a code-server release cannot change the editor UI the
+        # tests drive without the change showing up here first.
+        run: curl -fsSL https://code-server.dev/install.sh | sh -s -- --version 4.137.0
       - name: Install Playwright browsers
         working-directory: ./vscode/extension
         run: pnpm exec playwright install

vscode/extension/playwright.config.ts

@@ -2,7 +2,10 @@ import { defineConfig } from '@playwright/test'
 
 export default defineConfig({
   testDir: 'tests',
-  timeout: 60_000,
+  // The first test a worker runs pays for the editor and extension host starting
+  // up, which measured at 2.5-4 minutes locally against 21-39 seconds for every
+  // later test in the same worker. 60 seconds only ever fit the latter.
+  timeout: 180_000,
   // TODO: When stable, allow retries in CI
   retries: process.env.CI ? 2 : 0,
   workers: process.env.CI ? 2 : 4,

vscode/extension/tests/configuration.spec.ts

@@ -1,33 +1,12 @@
-import { test, expect } from './fixtures'
+import { test, expect, SharedPythonEnvironment } from './fixtures'
 import {
-  createVirtualEnvironment,
   openServerPage,
-  pipInstall,
-  REPO_ROOT,
   SUSHI_SOURCE_PATH,
   waitForLoadedSQLMesh,
 } from './utils'
 import path from 'path'
 import fs from 'fs-extra'
 
-async function setupPythonEnvironment(tempDir: string): Promise<void> {
-  // Create a temporary directory for the virtual environment
-  const venvDir = path.join(tempDir, '.venv')
-  fs.mkdirSync(venvDir, { recursive: true })
-
-  // Create virtual environment
-  const pythonDetails = await createVirtualEnvironment(venvDir)
-
-  // Install sqlmesh from the local repository with LSP support
-  const customMaterializations = path.join(
-    REPO_ROOT,
-    'examples',
-    'custom_materializations',
-  )
-  const sqlmeshWithExtras = `${REPO_ROOT}[lsp,bigquery]`
-  await pipInstall(pythonDetails, [sqlmeshWithExtras, customMaterializations])
-}
-
 /**
  * Creates an entrypoint file used to test the LSP configuration.
  *
@@ -36,14 +15,18 @@ async function setupPythonEnvironment(tempDir: string): Promise<void> {
 const createEntrypointFile = (
   tempDir: string,
   entrypointFileName: string,
+  pythonEnvironment: SharedPythonEnvironment,
   bitToStripFromArgs = '',
 ): {
   entrypointFile: string
   fileWhereStoredInputs: string
 } => {
   const entrypointFile = path.join(tempDir, entrypointFileName)
   const fileWhereStoredInputs = path.join(tempDir, 'inputs.txt')
-  const sqlmeshLSPFile = path.join(tempDir, '.venv/bin/sqlmesh_lsp')
+  const sqlmeshLSPFile = path.join(
+    path.dirname(pythonEnvironment.pythonPath),
+    'sqlmesh_lsp',
+  )
 
   // Create the entrypoint file
   fs.writeFileSync(
@@ -69,15 +52,15 @@ test.describe('Test LSP Entrypoint configuration', () => {
   test('specify single entrypoint relative path', async ({
     page,
     sharedCodeServer,
+    sharedPythonEnvironment,
     tempDir,
   }) => {
     await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
-    await setupPythonEnvironment(tempDir)
-
     const { fileWhereStoredInputs } = createEntrypointFile(
       tempDir,
       'entrypoint.sh',
+      sharedPythonEnvironment,
     )
 
     const settings = {
@@ -116,15 +99,15 @@ test.describe('Test LSP Entrypoint configuration', () => {
   test('specify one entrypoint absolute path', async ({
     page,
     sharedCodeServer,
+    sharedPythonEnvironment,
     tempDir,
   }) => {
     await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
-    await setupPythonEnvironment(tempDir)
-
     const { entrypointFile, fileWhereStoredInputs } = createEntrypointFile(
       tempDir,
       'entrypoint.sh',
+      sharedPythonEnvironment,
     )
     // Assert that the entrypoint file is an absolute path
     expect(path.isAbsolute(entrypointFile)).toBe(true)
@@ -165,15 +148,15 @@ test.describe('Test LSP Entrypoint configuration', () => {
   test('specify entrypoint with arguments', async ({
     page,
     sharedCodeServer,
+    sharedPythonEnvironment,
     tempDir,
   }) => {
     await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
-    await setupPythonEnvironment(tempDir)
-
     const { fileWhereStoredInputs } = createEntrypointFile(
       tempDir,
       'entrypoint.sh',
+      sharedPythonEnvironment,
       '--argToIgnore',
     )
 

vscode/extension/tests/fixtures.ts

@@ -7,6 +7,26 @@ import {
   stopCodeServer,
   CodeServerContext,
 } from './utils_code_server'
+import {
+  createVirtualEnvironment,
+  pipInstall,
+  PythonEnvironment,
+  REPO_ROOT,
+  warmUpVirtualEnvironment,
+} from './utils'
+
+/**
+ * A virtual environment with sqlmesh installed, shared by every test in a
+ * worker that does not need an environment of its own.
+ */
+export interface SharedPythonEnvironment extends PythonEnvironment {
+  venvDir: string
+}
+
+// Creating the environment installs sqlmesh from source, which is far slower
+// than a test. Give it a budget of its own so it is not charged to the first
+// test that happens to ask for it.
+const SHARED_PYTHON_ENVIRONMENT_TIMEOUT_MS = 300_000
 
 const REMOVE_RETRY_DELAYS_MS = [100, 250, 500, 1000, 2000]
 
@@ -39,7 +59,11 @@ const removeWithRetry = async (dir: string): Promise<void> => {
 export const test = base.extend<
   // eslint-disable-next-line @typescript-eslint/no-empty-object-type
   {},
-  { sharedCodeServer: CodeServerContext; tempDir: string }
+  {
+    sharedCodeServer: CodeServerContext
+    sharedPythonEnvironment: SharedPythonEnvironment
+    tempDir: string
+  }
 >({
   sharedCodeServer: [
     // eslint-disable-next-line no-empty-pattern
@@ -67,6 +91,33 @@ export const test = base.extend<
     },
     { scope: 'worker', auto: true },
   ],
+  sharedPythonEnvironment: [
+    // eslint-disable-next-line no-empty-pattern
+    async ({}, use) => {
+      // The environment lives outside the per-test temporary directory, which
+      // is removed after every test.
+      const envDir = await fs.mkdtemp(
+        path.join(os.tmpdir(), 'vscode-test-shared-python-env-'),
+      )
+      const venvDir = path.join(envDir, '.venv')
+      const pythonDetails = await createVirtualEnvironment(venvDir)
+      await pipInstall(pythonDetails, [
+        `${REPO_ROOT}[lsp,bigquery]`,
+        path.join(REPO_ROOT, 'examples', 'custom_materializations'),
+      ])
+      await warmUpVirtualEnvironment(pythonDetails)
+
+      console.log(
+        `Created shared Python environment for worker ${test.info().workerIndex} at ${venvDir}`,
+      )
+
+      await use({ ...pythonDetails, venvDir })
+
+      console.log(`Removing shared Python environment: ${envDir}`)
+      await fs.remove(envDir)
+    },
+    { scope: 'worker', timeout: SHARED_PYTHON_ENVIRONMENT_TIMEOUT_MS },
+  ],
   tempDir: [
     // eslint-disable-next-line no-empty-pattern
     async ({}, use) => {

vscode/extension/tests/lineage_settings.spec.ts

@@ -12,24 +12,39 @@ import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 /**
  * Find the iframe that hosts the lineage UI (the one containing the
  * Settings cog button). Returns null if it can't be located.
+ *
+ * The webview iframes attach asynchronously, so a single sweep can run before
+ * the lineage one exists and find nothing. The sweep is repeated until the
+ * Settings button answers from one of them or `timeout` elapses.
  */
-async function findLineageFrame(page: Page): Promise<FrameLocator | null> {
-  const iframes = page.locator('iframe')
-  const iframeCount = await iframes.count()
-
-  for (let i = 0; i < iframeCount; i++) {
-    const contentFrame = iframes.nth(i).contentFrame()
-    if (!contentFrame) continue
-    const activeFrame = contentFrame.locator('#active-frame').contentFrame()
-    if (!activeFrame) continue
-    try {
-      await activeFrame
-        .getByRole('button', { name: 'Settings' })
-        .waitFor({ timeout: 1000 })
-      return activeFrame
-    } catch {
-      continue
+async function findLineageFrame(
+  page: Page,
+  timeout = 60_000,
+): Promise<FrameLocator | null> {
+  const deadline = Date.now() + timeout
+
+  while (Date.now() < deadline) {
+    const iframes = page.locator('iframe')
+    const iframeCount = await iframes.count()
+
+    for (let i = 0; i < iframeCount; i++) {
+      const contentFrame = iframes.nth(i).contentFrame()
+      if (!contentFrame) continue
+      const activeFrame = contentFrame.locator('#active-frame').contentFrame()
+      if (!activeFrame) continue
+      try {
+        await activeFrame
+          .getByRole('button', { name: 'Settings' })
+          .waitFor({ timeout: 1000 })
+        return activeFrame
+      } catch {
+        continue
+      }
     }
+
+    // Nothing answered this time round; give the next webview a chance to
+    // attach instead of spinning when the page has no iframes yet.
+    await page.waitForTimeout(250)
   }
   return null
 }
@@ -97,9 +112,11 @@ test('Only Direct Neighbors toggle filters the lineage graph', async ({
   const nodesBefore = await lineageFrame.locator('.react-flow__node').count()
   expect(nodesBefore).toBeGreaterThan(0)
 
-  // Open the settings menu and toggle "Only Direct Neighbors"
+  // Open the settings menu and toggle "Only Direct Neighbors". Headless UI
+  // renders the entry as `<button role="menuitem">`, and the explicit role is
+  // the one the accessibility tree reports.
   await lineageFrame.getByRole('button', { name: 'Settings' }).click()
-  const toggle = lineageFrame.getByRole('button', {
+  const toggle = lineageFrame.getByRole('menuitem', {
     name: 'Only Direct Neighbors',
   })
   await toggle.waitFor()
@@ -115,10 +132,12 @@ test('Only Direct Neighbors toggle filters the lineage graph', async ({
 
   // Toggle off → graph returns to the full size
   await lineageFrame.getByRole('button', { name: 'Settings' }).click()
-  await lineageFrame
-    .getByRole('button', { name: 'Only Direct Neighbors' })
-    .click()
-  await page.waitForTimeout(250)
-  const nodesRestored = await lineageFrame.locator('.react-flow__node').count()
-  expect(nodesRestored).toBe(nodesBefore)
+  const toggleOff = lineageFrame.getByRole('menuitem', {
+    name: 'Only Direct Neighbors',
+  })
+  await toggleOff.waitFor()
+  await toggleOff.click()
+  await expect(lineageFrame.locator('.react-flow__node')).toHaveCount(
+    nodesBefore,
+  )
 })