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/python_env.spec.ts

@@ -1,4 +1,4 @@
-import { test, Page } from './fixtures'
+import { test, Page, SharedPythonEnvironment } from './fixtures'
 import fs from 'fs-extra'
 import {
   createVirtualEnvironment,
@@ -40,6 +40,42 @@ async function runTest(
   await openLineageView(page)
 }
 
+async function writeSettings(
+  tempDir: string,
+  pythonDetails: PythonEnvironment,
+  pythonEnvDir: string,
+): Promise<void> {
+  const settings = {
+    'python.defaultInterpreterPath': pythonDetails.pythonPath,
+    'sqlmesh.environmentPath': pythonEnvDir,
+  }
+  await fs.ensureDir(path.join(tempDir, '.vscode'))
+  await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, {
+    spaces: 2,
+  })
+}
+
+/**
+ * Copy the sushi project into the test directory and point it at the shared
+ * virtual environment.
+ */
+async function setupSharedEnvironment(
+  tempDir: string,
+  sharedPythonEnvironment: SharedPythonEnvironment,
+): Promise<void> {
+  await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+  await writeSettings(
+    tempDir,
+    sharedPythonEnvironment,
+    sharedPythonEnvironment.venvDir,
+  )
+}
+
+/**
+ * Copy the sushi project into the test directory and give it a virtual
+ * environment of its own. Used by the tcloud tests, which install a mock
+ * tcloud package that must not leak into other tests.
+ */
 async function setupEnvironment(tempDir: string): Promise<{
   pythonDetails: PythonEnvironment
 }> {
@@ -54,27 +90,30 @@ async function setupEnvironment(tempDir: string): Promise<{
   const sqlmeshWithExtras = `${REPO_ROOT}[bigquery,lsp]`
   await pipInstall(pythonDetails, [sqlmeshWithExtras, custom_materializations])
 
-  const settings = {
-    'python.defaultInterpreterPath': pythonDetails.pythonPath,
-    'sqlmesh.environmentPath': pythonEnvDir,
-  }
-  await fs.ensureDir(path.join(tempDir, '.vscode'))
-  await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, {
-    spaces: 2,
-  })
+  await writeSettings(tempDir, pythonDetails, pythonEnvDir)
   return { pythonDetails }
 }
 
 test.describe('python environment variable injection on sqlmesh_lsp', () => {
-  test('normal setup - error ', async ({ page, sharedCodeServer, tempDir }) => {
-    await setupEnvironment(tempDir)
+  test('normal setup - error ', async ({
+    page,
+    sharedCodeServer,
+    sharedPythonEnvironment,
+    tempDir,
+  }) => {
+    await setupSharedEnvironment(tempDir, sharedPythonEnvironment)
     writeEnvironmentConfig(tempDir)
     await runTest(page, sharedCodeServer, tempDir)
     await page.waitForSelector('text=Error creating context')
   })
 
-  test('normal setup - set', async ({ page, sharedCodeServer, tempDir }) => {
-    await setupEnvironment(tempDir)
+  test('normal setup - set', async ({
+    page,
+    sharedCodeServer,
+    sharedPythonEnvironment,
+    tempDir,
+  }) => {
+    await setupSharedEnvironment(tempDir, sharedPythonEnvironment)
     writeEnvironmentConfig(tempDir)
     const env_file = path.join(tempDir, '.env')
     fs.writeFileSync(env_file, 'TEST_VAR=test_value')

vscode/extension/tests/utils.ts

@@ -89,6 +89,26 @@ export const pipInstall = async (
   }
 }
 
+/**
+ * Import the language server entrypoint once so that Python writes its
+ * bytecode cache.
+ *
+ * The first import of sqlmesh in a fresh environment is several times slower
+ * than every import after it, because the whole dependency tree has to be
+ * compiled. Paying that here rather than inside a test keeps it out of the
+ * per-test timeout.
+ */
+export const warmUpVirtualEnvironment = async (
+  pythonDetails: PythonEnvironment,
+): Promise<void> => {
+  const { stderr, exitCode } = await execAsync(
+    `"${pythonDetails.pythonPath}" -c "import sqlmesh.lsp.main"`,
+  )
+  if (exitCode !== 0) {
+    throw new Error(`Failed to warm up the virtual environment: ${stderr}`)
+  }
+}
+
 /**
  * Open the lineage view in the given window.
  */