Repository: TobikoData/sqlmesh
Pull request #4821 (2025-06-26 01:58:29)

Fix: display_name should not exclude the default catalog when environment_suffix_target is set to CATALOG
Prior to this, when using `environment_suffix_target: catalog` against a model in the default catalog, the CLI would give output like:

```
$ sqlmesh plan erin
...SNIP...
Models needing backfill:
└── none__erin.sqlmesh_example.incremental_model: [2020-01-01 - 2025-04-07]
Apply - Backfill Tables [y/n]: y
```

This was due to an output optimisation where we hid the catalog as unnecessary information if it matched the default catalog.

However, `environment_suffix_target: catalog` relies on the default catalog being present because it needs it to generate the catalog name for the virtual environment.

This PR fixes the output:

```
$ sqlmesh plan erin
...SNIP...
Models needing backfill:
└── tobiko-1__erin.sqlmesh_example.incremental_model: [2020-01-01 - 2025-04-07]
Apply - Backfill Tables [y/n]: y
```

Note that this was a display-only issue, the SQL queries being executed in the background still contained the correct catalog


Patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 37b03fbe..7a12d3c0 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -39,7 +39,7 @@ jobs:
           command: unset TWINE_USERNAME TWINE_PASSWORD && make publish-tests
   gh-release:
     docker:
-      - image: cimg/node:16.14
+      - image: cimg/node:20.19.0
     resource_class: small
     steps:
       - run:
@@ -54,7 +54,7 @@ jobs:
 
   ui-build:
     docker:
-      - image: cimg/node:19.8
+      - image: cimg/node:20.19.0
     resource_class: medium
     steps:
       - checkout
diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 765aea04..1d7f6ec8 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -17,7 +17,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: actions/setup-node@v4
         with:
-          node-version: '20'
+          node-version: '22'
       - uses: pnpm/action-setup@v4
         with:
           version: latest
@@ -25,3 +25,38 @@ jobs:
         run: pnpm install
       - name: Run CI
         run: pnpm run ci
+  test-vscode-e2e:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: '22'
+      - uses: pnpm/action-setup@v4
+        with:
+          version: latest
+      - name: Install dependencies
+        run: pnpm install
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+      - name: Install python dependencies
+        run: |
+          python -m venv .venv
+          source .venv/bin/activate
+          make install-dev
+      - name: Fetch VS Code
+        working-directory: ./vscode/extension
+        run: pnpm run fetch-vscode
+
+      - name: Install code-server
+        run: curl -fsSL https://code-server.dev/install.sh | sh
+      - name: Install Playwright browsers
+        working-directory: ./vscode/extension
+        run: pnpm exec playwright install
+      - name: Run e2e tests
+        working-directory: ./vscode/extension
+        run: |
+          source ../../.venv/bin/activate
+          pnpm run test:e2e tests/stop.spec.ts
diff --git a/.github/workflows/release_extension.yaml b/.github/workflows/release_extension.yaml
index 1807fee3..3807bc34 100644
--- a/.github/workflows/release_extension.yaml
+++ b/.github/workflows/release_extension.yaml
@@ -1,14 +1,32 @@
 name: Release VSCode Extension
 on:
-  push:
-    tags:
-      - 'vscode@v*'
+  workflow_dispatch:
+    inputs:
+      version:
+        description: 'Version to release (e.g., 1.0.0)'
+        required: true
+        type: string
 jobs:
   release:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
+      - name: Check branch is main
+        run: |
+          if [[ "${{ github.ref }}" != "refs/heads/main" ]]; then
+            echo "Error: This workflow can only be run from the main branch"
+            exit 1
+          fi
+          echo "Branch check passed: running from main branch"
+      - name: Validate version format
+        run: |
+          version="${{ github.event.inputs.version }}"
+          if ! [[ $version =~ ^[0-9]+\.[0-9]+\.[0-9]+(-[0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*)?(\+[0-9A-Za-z-]+(\.[0-9A-Za-z-]+)*)?$ ]]; then
+            echo "Error: Version must be a valid semantic version (e.g., 1.0.0, 1.0.0-beta.1, 1.0.0+build.1)"
+            exit 1
+          fi
+          echo "Version format is valid: $version"
       - name: Setup Node.js
         uses: actions/setup-node@v4
         with:
@@ -19,26 +37,20 @@ jobs:
           version: 10
       - name: Install dependencies
         run: pnpm install --frozen-lockfile
-      - name: Extract version from tag
-        id: extract_version
-        run: |
-          VERSION=${GITHUB_REF#refs/tags/vscode@v}
-          echo "VERSION=$VERSION" >> $GITHUB_OUTPUT
       - name: Update package.json version
         working-directory: vscode/extension
         run: |
-          npm version ${{ steps.extract_version.outputs.VERSION }} --no-git-tag-version
-      - name: Install dependencies
-        working-directory: vscode/extension
-        run: pnpm install
-      - name: Run CI
-        run: pnpm run ci
+          npm version ${{ github.event.inputs.version }} --no-git-tag-version
       - name: Build extension
         working-directory: vscode/extension
         run: pnpm run vscode:package
       - name: Upload extension to Marketplace
         working-directory: vscode/extension
         run: |
-          pnpx vsce publish --packagePath sqlmesh-${{ steps.extract_version.outputs.VERSION }}.vsix
+          pnpx vsce publish --packagePath sqlmesh-${{ github.event.inputs.version }}.vsix
         env:
           VSCE_PAT: ${{ secrets.VSCE_PAT }}
+      - name: Upload extension to OpenVSX
+        working-directory: vscode/extension
+        run: |
+          pnpx ovsx publish -p ${{ secrets.OPEN_VSX_TOKEN }} sqlmesh-${{ github.event.inputs.version }}.vsix
diff --git a/.nvmrc b/.nvmrc
new file mode 100644
index 00000000..209e3ef4
--- /dev/null
+++ b/.nvmrc
@@ -0,0 +1,1 @@
+20
diff --git a/package.json b/package.json
index d6383cfc..82be6472 100644
--- a/package.json
+++ b/package.json
@@ -1,4 +1,8 @@
 {
+  "engines": {
+    "node": ">=20.0.0",
+    "pnpm": ">=10.0.0"
+  },
   "scripts": {
     "ci": "pnpm run lint && pnpm run -r ci",
     "fmt": "prettier --write .",
diff --git a/pyproject.toml b/pyproject.toml
index acee956e..235bbb81 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -132,7 +132,7 @@ lsp = [
     "sse-starlette>=0.2.2",
     "pyarrow",
     # For lsp
-    "pygls",
+    "pygls>=1.2.0,<2.0.0",
     "lsprotocol",
 ]
 risingwave = ["psycopg2"]
diff --git a/sqlmesh/core/engine_adapter/mssql.py b/sqlmesh/core/engine_adapter/mssql.py
index a00e11e0..200640f3 100644
--- a/sqlmesh/core/engine_adapter/mssql.py
+++ b/sqlmesh/core/engine_adapter/mssql.py
@@ -6,11 +6,13 @@ import typing as t
 
 from sqlglot import exp
 
-from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.dialect import to_schema, add_table
 from sqlmesh.core.engine_adapter.base import (
     EngineAdapterWithIndexSupport,
     EngineAdapter,
     InsertOverwriteStrategy,
+    MERGE_SOURCE_ALIAS,
+    MERGE_TARGET_ALIAS,
 )
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
@@ -32,7 +34,7 @@ from sqlmesh.core.schema_diff import SchemaDiffer
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import SchemaName, TableName
-    from sqlmesh.core.engine_adapter._typing import DF, Query
+    from sqlmesh.core.engine_adapter._typing import DF, Query, QueryOrDF
 
 
 @set_catalog()
@@ -188,6 +190,87 @@ class MSSQLEngineAdapter(
                     )
         super().drop_schema(schema_name, ignore_if_not_exists=ignore_if_not_exists, cascade=False)
 
+    def merge(
+        self,
+        target_table: TableName,
+        source_table: QueryOrDF,
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
+        unique_key: t.Sequence[exp.Expression],
+        when_matched: t.Optional[exp.Whens] = None,
+        merge_filter: t.Optional[exp.Expression] = None,
+    ) -> None:
+        source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
+            source_table, columns_to_types, target_table=target_table
+        )
+        columns_to_types = columns_to_types or self.columns(target_table)
+        on = exp.and_(
+            *(
+                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
+                for part in unique_key
+            )
+        )
+        if merge_filter:
+            on = exp.and_(merge_filter, on)
+
+        if not when_matched:
+            match_condition = None
+            unique_key_names = [y.name for y in unique_key]
+            columns_to_types_no_keys = [c for c in columns_to_types if c not in unique_key_names]
+
+            target_columns_no_keys = [
+                exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types_no_keys
+            ]
+            source_columns_no_keys = [
+                exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types_no_keys
+            ]
+
+            match_condition = exp.Exists(
+                this=exp.select(*target_columns_no_keys).except_(
+                    exp.select(*source_columns_no_keys)
+                )
+            )
+
+            match_expressions = [
+                exp.When(
+                    matched=True,
+                    source=False,
+                    condition=match_condition,
+                    then=exp.Update(
+                        expressions=[
+                            exp.column(col, MERGE_TARGET_ALIAS).eq(
+                                exp.column(col, MERGE_SOURCE_ALIAS)
+                            )
+                            for col in columns_to_types_no_keys
+                        ],
+                    ),
+                )
+            ]
+        else:
+            match_expressions = when_matched.copy().expressions
+
+        match_expressions.append(
+            exp.When(
+                matched=False,
+                source=False,
+                then=exp.Insert(
+                    this=exp.Tuple(expressions=[exp.column(col) for col in columns_to_types]),
+                    expression=exp.Tuple(
+                        expressions=[
+                            exp.column(col, MERGE_SOURCE_ALIAS) for col in columns_to_types
+                        ]
+                    ),
+                ),
+            )
+        )
+        for source_query in source_queries:
+            with source_query as query:
+                self._merge(
+                    target_table=target_table,
+                    query=query,
+                    on=on,
+                    whens=exp.Whens(expressions=match_expressions),
+                )
+
     def _convert_df_datetime(self, df: DF, columns_to_types: t.Dict[str, exp.DataType]) -> None:
         import pandas as pd
         from pandas.api.types import is_datetime64_any_dtype  # type: ignore
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 573d3bc7..6e000764 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -15,6 +15,7 @@ from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import StandaloneAudit
+from sqlmesh.core.environment import EnvironmentSuffixTarget
 from sqlmesh.core.macros import call_macro
 from sqlmesh.core.model import Model, ModelKindMixin, ModelKindName, ViewKind, CustomKind
 from sqlmesh.core.model.definition import _Model
@@ -1589,12 +1590,18 @@ def display_name(
     if snapshot_info_like.is_audit:
         return snapshot_info_like.name
     view_name = exp.to_table(snapshot_info_like.name)
+
+    catalog = (
+        None
+        if (
+            environment_naming_info.suffix_target != EnvironmentSuffixTarget.CATALOG
+            and view_name.catalog == default_catalog
+        )
+        else view_name.catalog
+    )
+
     qvn = QualifiedViewName(
-        catalog=(
-            view_name.catalog
-            if view_name.catalog and view_name.catalog != default_catalog
-            else None
-        ),
+        catalog=catalog,
         schema_name=view_name.db or None,
         table=view_name.name,
     )
diff --git a/vscode/extension/package.json b/vscode/extension/package.json
index cc5722c6..e8b7d5f0 100644
--- a/vscode/extension/package.json
+++ b/vscode/extension/package.json
@@ -120,9 +120,10 @@
     "lint": "eslint src",
     "lint:fix": "eslint src --fix",
     "test:unit": "vitest run",
-    "test:e2e": "playwright test",
-    "test:e2e:ui": "playwright test --ui",
-    "test:e2e:headed": "playwright test --headed",
+    "code-server": "code-server",
+    "test:e2e": "pnpm run vscode:package && playwright test",
+    "test:e2e:ui": "pnpm run vscode:package && playwright test --ui",
+    "test:e2e:headed": "pnpm run vscode:package && playwright test --headed",
     "fetch-vscode": "tsx scripts/fetch-vscode.ts",
     "compile": "pnpm run check-types && node esbuild.js",
     "check-types": "tsc --noEmit -p ./tsconfig.build.json",
diff --git a/vscode/extension/playwright.config.ts b/vscode/extension/playwright.config.ts
index 6820f2c8..599e5e37 100644
--- a/vscode/extension/playwright.config.ts
+++ b/vscode/extension/playwright.config.ts
@@ -11,7 +11,7 @@ export default defineConfig({
       use: {
         // ⭢ we'll launch Electron ourselves – no browser needed
         browserName: 'chromium',
-        headless: false, // headed makes screenshots deterministic
+        headless: true, // headless mode for tests
         launchOptions: {
           slowMo: process.env.CI ? 0 : 100,
         },
diff --git a/vscode/react/src/components/graph/ModelLineage.tsx b/vscode/react/src/components/graph/ModelLineage.tsx
index b936dd02..652c0631 100644
--- a/vscode/react/src/components/graph/ModelLineage.tsx
+++ b/vscode/react/src/components/graph/ModelLineage.tsx
@@ -383,7 +383,10 @@ function ModelColumnLineage(): JSX.Element {
             <Divider />
           </Panel>
         )}
-        <Controls className="bg-light p-1 rounded-md !border-none !shadow-lg" />
+        <Controls
+          className="bg-light p-1 rounded-md !border-none !shadow-lg"
+          showInteractive={false}
+        />
         <Background
           variant={BackgroundVariant.Cross}
           gap={32}
diff --git a/vscode/react/src/pages/lineage.tsx b/vscode/react/src/pages/lineage.tsx
index d2d4e858..546e1ad4 100644
--- a/vscode/react/src/pages/lineage.tsx
+++ b/vscode/react/src/pages/lineage.tsx
@@ -110,6 +110,8 @@ function Lineage() {
       fetchFirstTimeModelIfNotSet(models).then(modelName => {
         if (modelName && selectedModel === undefined) {
           setSelectedModel(modelName)
+        } else {
+          setSelectedModel(models[0].name)
         }
       })
     }


Test patch:
diff --git a/tests/core/engine_adapter/test_mssql.py b/tests/core/engine_adapter/test_mssql.py
index 22f37174..a420a2d8 100644
--- a/tests/core/engine_adapter/test_mssql.py
+++ b/tests/core/engine_adapter/test_mssql.py
@@ -472,7 +472,7 @@ def test_merge_pandas(
 
     assert to_sql_calls(adapter) == [
         f"""IF NOT EXISTS (SELECT * FROM information_schema.tables WHERE table_name = '__temp_target_{temp_table_id}') EXEC('CREATE TABLE [__temp_target_{temp_table_id}] ([id] INTEGER, [ts] DATETIME2, [val] INTEGER)');""",
-        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] WHEN MATCHED THEN UPDATE SET [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id], [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
+        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] WHEN MATCHED AND EXISTS(SELECT [__MERGE_TARGET__].[ts], [__MERGE_TARGET__].[val] EXCEPT SELECT [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]) THEN UPDATE SET [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
         f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
     ]
 
@@ -495,7 +495,7 @@ def test_merge_pandas(
 
     assert to_sql_calls(adapter) == [
         f"""IF NOT EXISTS (SELECT * FROM information_schema.tables WHERE table_name = '__temp_target_{temp_table_id}') EXEC('CREATE TABLE [__temp_target_{temp_table_id}] ([id] INTEGER, [ts] DATETIME2, [val] INTEGER)');""",
-        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] AND [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts] WHEN MATCHED THEN UPDATE SET [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id], [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts], [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
+        f"MERGE INTO [target] AS [__MERGE_TARGET__] USING (SELECT CAST([id] AS INTEGER) AS [id], CAST([ts] AS DATETIME2) AS [ts], CAST([val] AS INTEGER) AS [val] FROM [__temp_target_{temp_table_id}]) AS [__MERGE_SOURCE__] ON [__MERGE_TARGET__].[id] = [__MERGE_SOURCE__].[id] AND [__MERGE_TARGET__].[ts] = [__MERGE_SOURCE__].[ts] WHEN MATCHED AND EXISTS(SELECT [__MERGE_TARGET__].[val] EXCEPT SELECT [__MERGE_SOURCE__].[val]) THEN UPDATE SET [__MERGE_TARGET__].[val] = [__MERGE_SOURCE__].[val] WHEN NOT MATCHED THEN INSERT ([id], [ts], [val]) VALUES ([__MERGE_SOURCE__].[id], [__MERGE_SOURCE__].[ts], [__MERGE_SOURCE__].[val]);",
         f"DROP TABLE IF EXISTS [__temp_target_{temp_table_id}];",
     ]
 
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index e4eb12c5..f09083f5 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -2178,6 +2178,21 @@ def test_deployability_index_missing_parent(make_snapshot):
             "snowflake",
             "CATALOG_OVERRIDE.test_db.test_model__DEV",
         ),
+        # EnvironmentSuffixTarget.CATALOG
+        (
+            "test_db.test_model",
+            EnvironmentNamingInfo(name="dev", suffix_target=EnvironmentSuffixTarget.CATALOG),
+            "default_catalog",
+            "duckdb",
+            "default_catalog__dev.test_db.test_model",
+        ),
+        (
+            "test_db.test_model",
+            EnvironmentNamingInfo(name="dev", suffix_target=EnvironmentSuffixTarget.CATALOG),
+            "default_catalog",
+            "snowflake",
+            "DEFAULT_CATALOG__DEV.test_db.test_model",
+        ),
     ),
 )
 def test_display_name(
diff --git a/vscode/extension/tests/stop.spec.ts b/vscode/extension/tests/stop.spec.ts
index 61422991..d7b2d08c 100644
--- a/vscode/extension/tests/stop.spec.ts
+++ b/vscode/extension/tests/stop.spec.ts
@@ -1,57 +1,67 @@
 import path from 'path'
-import { startVSCode, SUSHI_SOURCE_PATH } from './utils'
+import { SUSHI_SOURCE_PATH } from './utils'
 import os from 'os'
 import { test } from '@playwright/test'
 import fs from 'fs-extra'
+import { startCodeServer, stopCodeServer } from './utils_code_server'
 
-test('Stop server works', async () => {
+test('Stop server works', async ({ page }) => {
+  test.setTimeout(120000) // Increase timeout to 2 minutes
+
+  console.log('Starting test: Stop server works')
   const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'))
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
+  const context = await startCodeServer(tempDir, true)
+
   try {
-    const { window, close } = await startVSCode(tempDir)
+    // Navigate to code-server instance
+    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+
+    // Wait for code-server to load
+    await page.waitForLoadState('networkidle')
+    await page.waitForSelector('[role="application"]', { timeout: 10000 })
 
-    //   Wait for the models folder to be visible
-    await window.waitForSelector('text=models')
+    // Wait for the models folder to be visible in the file explorer
+    await page.waitForSelector('text=models')
 
     // Click on the models folder, excluding external_models
-    await window
+    await page
       .getByRole('treeitem', { name: 'models', exact: true })
       .locator('a')
       .click()
 
-    // Open the customer_revenue_lifetime model
-    await window
+    // Open the customers.sql model
+    await page
       .getByRole('treeitem', { name: 'customers.sql', exact: true })
       .locator('a')
       .click()
 
-    await window.waitForSelector('text=grain')
-    await window.waitForSelector('text=Loaded SQLMesh Context')
+    await page.waitForSelector('text=grain')
+    await page.waitForSelector('text=Loaded SQLMesh Context')
 
     // Stop the server
-    await window.keyboard.press(
+    await page.keyboard.press(
       process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P',
     )
-    await window.keyboard.type('SQLMesh: Stop Server')
-    await window.keyboard.press('Enter')
+    await page.keyboard.type('SQLMesh: Stop Server')
+    await page.keyboard.press('Enter')
 
     // Await LSP server stopped message
-    await window.waitForSelector('text=LSP server stopped')
+    await page.waitForSelector('text=LSP server stopped')
 
     // Render the model
-    await window.keyboard.press(
+    await page.keyboard.press(
       process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P',
     )
-    await window.keyboard.type('Render Model')
-    await window.keyboard.press('Enter')
+    await page.keyboard.type('Render Model')
+    await page.keyboard.press('Enter')
 
     // Await error message
-    await window.waitForSelector(
+    await page.waitForSelector(
       'text="Failed to render model: LSP client not ready."',
     )
-    await close()
   } finally {
-    await fs.remove(tempDir)
+    await stopCodeServer(context)
   }
 })
diff --git a/vscode/extension/tests/utils_code_server.ts b/vscode/extension/tests/utils_code_server.ts
new file mode 100644
index 00000000..c38fee61
--- /dev/null
+++ b/vscode/extension/tests/utils_code_server.ts
@@ -0,0 +1,165 @@
+import { spawn, ChildProcess, execSync } from 'child_process'
+import path from 'path'
+import fs from 'fs-extra'
+
+export interface CodeServerContext {
+  codeServerProcess: ChildProcess
+  codeServerPort: number
+  tempDir: string
+}
+
+/**
+ * @param tempDir - The temporary directory to use for the code-server instance
+ * @param placeFileWithPythonInterpreter - Whether to place a vscode/settings.json file in the temp directory that points to the python interpreter of the environmen the test is running in.
+ * @returns The code-server context
+ */
+export async function startCodeServer(
+  tempDir: string,
+  placeFileWithPythonInterpreter: boolean = false,
+): Promise<CodeServerContext> {
+  // Find an available port
+  const codeServerPort = Math.floor(Math.random() * 10000) + 50000
+
+  // Create .vscode/settings.json with Python interpreter if requested
+  if (placeFileWithPythonInterpreter) {
+    const vscodeDir = path.join(tempDir, '.vscode')
+    await fs.ensureDir(vscodeDir)
+
+    // Get the current Python interpreter path
+    const pythonPath = execSync('which python', {
+      encoding: 'utf-8',
+    }).trim()
+
+    const settings = {
+      'python.defaultInterpreterPath': path.join(
+        __dirname,
+        '..',
+        '..',
+        '..',
+        '.venv',
+        'bin',
+        'python',
+      ),
+    }
+
+    await fs.writeJson(path.join(vscodeDir, 'settings.json'), settings, {
+      spaces: 2,
+    })
+    console.log(
+      `Created .vscode/settings.json with Python interpreter: ${pythonPath}`,
+    )
+  }
+
+  // Get the extension version from package.json
+  const extensionDir = path.join(__dirname, '..')
+  const packageJson = JSON.parse(
+    fs.readFileSync(path.join(extensionDir, 'package.json'), 'utf-8'),
+  )
+  const version = packageJson.version
+  const extensionName = packageJson.name || 'sqlmesh'
+
+  // Look for the specific version .vsix file
+  const vsixFileName = `${extensionName}-${version}.vsix`
+  const vsixPath = path.join(extensionDir, vsixFileName)
+
+  if (!fs.existsSync(vsixPath)) {
+    throw new Error(
+      `Extension file ${vsixFileName} not found. Run "pnpm run vscode:package" first.`,
+    )
+  }
+
+  console.log(`Using extension: ${vsixFileName}`)
+
+  // Install the extension first
+  const extensionsDir = path.join(tempDir, 'extensions')
+  console.log('Installing extension...')
+  execSync(
+    `pnpm run code-server --user-data-dir "${tempDir}" --extensions-dir "${extensionsDir}" --install-extension "${vsixPath}"`,
+    { stdio: 'inherit' },
+  )
+
+  // Start code-server instance
+  const codeServerProcess = spawn(
+    'pnpm',
+    [
+      'run',
+      'code-server',
+      '--bind-addr',
+      `127.0.0.1:${codeServerPort}`,
+      '--auth',
+      'none',
+      '--disable-telemetry',
+      '--disable-update-check',
+      '--disable-workspace-trust',
+      '--user-data-dir',
+      tempDir,
+      '--extensions-dir',
+      extensionsDir,
+      tempDir,
+    ],
+    {
+      stdio: 'pipe',
+      cwd: path.join(__dirname, '..'),
+    },
+  )
+
+  // Wait for code-server to be ready
+  await new Promise<void>((resolve, reject) => {
+    let output = ''
+    const timeout = setTimeout(() => {
+      reject(new Error('Code-server failed to start within timeout'))
+    }, 30000)
+
+    codeServerProcess.stdout?.on('data', data => {
+      output += data.toString()
+      if (output.includes('HTTP server listening on')) {
+        clearTimeout(timeout)
+        resolve()
+      }
+    })
+
+    codeServerProcess.stderr?.on('data', data => {
+      console.error('Code-server stderr:', data.toString())
+    })
+
+    codeServerProcess.on('error', error => {
+      clearTimeout(timeout)
+      reject(error)
+    })
+
+    codeServerProcess.on('exit', code => {
+      if (code !== 0) {
+        clearTimeout(timeout)
+        reject(new Error(`Code-server exited with code ${code}`))
+      }
+    })
+  })
+
+  return { codeServerProcess, codeServerPort, tempDir }
+}
+
+export async function stopCodeServer(
+  context: CodeServerContext,
+): Promise<void> {
+  const { codeServerProcess, tempDir } = context
+
+  // Clean up code-server process
+  codeServerProcess.kill('SIGTERM')
+
+  // Wait for process to exit
+  await new Promise<void>(resolve => {
+    codeServerProcess.on('exit', () => {
+      resolve()
+    })
+    // Force kill after 5 seconds
+    setTimeout(() => {
+      if (!codeServerProcess.killed) {
+        codeServerProcess.kill('SIGKILL')
+      }
+      resolve()
+    }, 5000)
+  })
+
+  // Clean up temporary directory
+  await fs.remove(tempDir)
+}