From b666a9446552db242c4d33628d6280926c993e33 Mon Sep 17 00:00:00 2001
From: Ben King <9087625+benfdking@users.noreply.github.com>
Date: Sun, 13 Jul 2025 22:09:19 +0100
Subject: [PATCH] chore(vscode): refactoring e2e tests

---
 vscode/extension/tests/bad_setup.spec.ts      | 31 ++-------
 vscode/extension/tests/broken_project.spec.ts | 26 ++-----
 vscode/extension/tests/completions.spec.ts    | 15 ++--
 vscode/extension/tests/diagnostics.spec.ts    | 37 +++-------
 .../extension/tests/find_references.spec.ts   | 68 ++++++-------------
 vscode/extension/tests/format.spec.ts         |  7 +-
 .../extension/tests/go_to_definition.spec.ts  | 12 +---
 vscode/extension/tests/hints.spec.ts          |  7 +-
 vscode/extension/tests/lineage.spec.ts        | 23 +++----
 .../extension/tests/lineage_settings.spec.ts  |  6 +-
 vscode/extension/tests/python_env.spec.ts     |  9 +--
 vscode/extension/tests/rename_cte.spec.ts     | 12 ++--
 vscode/extension/tests/render.spec.ts         | 24 +++----
 vscode/extension/tests/stop.spec.ts           |  9 +--
 vscode/extension/tests/tcloud.spec.ts         | 32 ++++-----
 vscode/extension/tests/utils.ts               | 18 ++++-
 vscode/extension/tests/venv_naming.spec.ts    |  5 +-
 17 files changed, 115 insertions(+), 226 deletions(-)

diff --git a/vscode/extension/tests/bad_setup.spec.ts b/vscode/extension/tests/bad_setup.spec.ts
index db3d481341..2284e2a95e 100644
--- a/vscode/extension/tests/bad_setup.spec.ts
+++ b/vscode/extension/tests/bad_setup.spec.ts
@@ -6,6 +6,7 @@ import {
   createVirtualEnvironment,
   openFile,
   openLineageView,
+  openServerPage,
   pipInstall,
   REPO_ROOT,
   SUSHI_SOURCE_PATH,
@@ -41,21 +42,13 @@ test('missing LSP dependencies shows install prompt', async ({
     spaces: 2,
   })
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-
-  // Open a SQL file to trigger SQLMesh activation
-  // Wait for the models folder to be visible
-  await page.waitForSelector('text=models')
+  await openServerPage(page, tempDir, sharedCodeServer)
 
-  // Click on the models folder
+  // Open a top_waiters model to trigger SQLMesh activation
   await page
     .getByRole('treeitem', { name: 'models', exact: true })
     .locator('a')
     .click()
-
-  // Open the top_waiters model
   await page
     .getByRole('treeitem', { name: 'customers.sql', exact: true })
     .locator('a')
@@ -66,12 +59,7 @@ test('missing LSP dependencies shows install prompt', async ({
   expect(await page.locator('text=Install').count()).toBeGreaterThanOrEqual(1)
 })
 
-test('lineage, no sqlmesh found', async ({
-  page,
-  sharedCodeServer,
-}, testInfo) => {
-  testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
-
+test('lineage, no sqlmesh found', async ({ page, sharedCodeServer }) => {
   const tempDir = await fs.mkdtemp(
     path.join(os.tmpdir(), 'vscode-test-tcloud-'),
   )
@@ -92,12 +80,8 @@ test('lineage, no sqlmesh found', async ({
   })
 
   // navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-  await page.waitForLoadState('networkidle')
+  await openServerPage(page, tempDir, sharedCodeServer)
 
-  // Open lineage view
   await openLineageView(page)
 
   // Assert shows that sqlmesh is not installed
@@ -143,10 +127,7 @@ test.skip('check that the LSP runs correctly by opening lineage when looking at
   await fs.ensureDir(path.dirname(sqlFile))
   await fs.writeFile(sqlFile, 'SELECT 1')
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-  await page.waitForLoadState('networkidle')
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Open the SQL file from the other directory
   await openFile(page, sqlFile)
diff --git a/vscode/extension/tests/broken_project.spec.ts b/vscode/extension/tests/broken_project.spec.ts
index 60c103fe7f..af7b022545 100644
--- a/vscode/extension/tests/broken_project.spec.ts
+++ b/vscode/extension/tests/broken_project.spec.ts
@@ -4,6 +4,7 @@ import os from 'os'
 import path from 'path'
 import {
   openLineageView,
+  openServerPage,
   runCommand,
   saveFile,
   SUSHI_SOURCE_PATH,
@@ -29,9 +30,7 @@ test('bad project, double model', async ({ page, sharedCodeServer }) => {
   )
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   await page.waitForSelector('text=models')
 
@@ -60,9 +59,7 @@ test('working project, then broken through adding double model, then refixed', a
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
   await page.waitForLoadState('networkidle')
 
   // Open the lineage view to confirm it loads properly
@@ -170,13 +167,7 @@ test('bad project, double model, then fixed', async ({
     customersSql,
   )
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-  await page.waitForLoadState('networkidle')
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   await page.waitForSelector('text=models')
 
@@ -245,10 +236,7 @@ test('bad project, double model, check lineage', async ({
   )
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-  await page.waitForLoadState('networkidle')
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Open the lineage view
   await openLineageView(page)
@@ -273,9 +261,7 @@ test('bad model block, then fixed', async ({ page, sharedCodeServer }) => {
     'MODEL ( name sushi.bad_block, test); SELECT * FROM sushi.customers'
   await fs.writeFile(badModelPath, contents)
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
   await page.waitForLoadState('networkidle')
 
   // Open the customers.sql model
diff --git a/vscode/extension/tests/completions.spec.ts b/vscode/extension/tests/completions.spec.ts
index 693934d4e4..da4eed6efc 100644
--- a/vscode/extension/tests/completions.spec.ts
+++ b/vscode/extension/tests/completions.spec.ts
@@ -2,7 +2,7 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { SUSHI_SOURCE_PATH } from './utils'
+import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Autocomplete for model names', async ({ page, sharedCodeServer }) => {
@@ -10,9 +10,7 @@ test('Autocomplete for model names', async ({ page, sharedCodeServer }) => {
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Wait for the models folder to be visible
   await page.waitForSelector('text=models')
@@ -67,9 +65,7 @@ test.describe('Macro Completions', () => {
     await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
     await createPythonInterpreterSettingsSpecifier(tempDir)
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Wait for the models folder to be visible
     await page.waitForSelector('text=models')
@@ -119,10 +115,7 @@ test.describe('Macro Completions', () => {
     await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
     await createPythonInterpreterSettingsSpecifier(tempDir)
-
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Wait for the models folder to be visible
     await page.waitForSelector('text=models')
diff --git a/vscode/extension/tests/diagnostics.spec.ts b/vscode/extension/tests/diagnostics.spec.ts
index f36cf81c59..a520cabcf7 100644
--- a/vscode/extension/tests/diagnostics.spec.ts
+++ b/vscode/extension/tests/diagnostics.spec.ts
@@ -2,7 +2,7 @@ import { expect, test } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { runCommand, SUSHI_SOURCE_PATH } from './utils'
+import { openServerPage, runCommand, SUSHI_SOURCE_PATH } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 import { execAsync } from '../src/utilities/exec'
 import yaml from 'yaml'
@@ -20,10 +20,7 @@ test('Workspace diagnostics show up in the diagnostics panel', async ({
   const updatedContent = configContent.replace('enabled=False', 'enabled=True')
   await fs.writeFile(configPath, updatedContent)
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-
+  await openServerPage(page, tempDir, sharedCodeServer)
   // Wait for the models folder to be visible
   await page.waitForSelector('text=models')
 
@@ -70,10 +67,7 @@ test.describe('Bad config.py/config.yaml file issues', () => {
     // Write an invalid YAML to config.yaml
     await fs.writeFile(configYamlPath, 'invalid_yaml; asdfasudfy')
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-    await page.waitForLoadState('networkidle')
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open full_model.sql model
     await page
@@ -116,10 +110,7 @@ test.describe('Bad config.py/config.yaml file issues', () => {
     // Write config to the yaml file
     await fs.writeFile(configYamlPath, yaml.stringify(config))
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-    await page.waitForLoadState('networkidle')
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open full_model.sql model
     await page
@@ -158,10 +149,7 @@ test.describe('Bad config.py/config.yaml file issues', () => {
     // Write an invalid Python to config.py
     await fs.writeFile(configPyPath, 'config = {}')
 
-    await page.goto(
-      `http://127.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-    await page.waitForLoadState('networkidle')
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open customers.sql model
     await page
@@ -197,10 +185,7 @@ test.describe('Bad config.py/config.yaml file issues', () => {
     // Write an invalid Python to config.py
     await fs.writeFile(configPyPath, 'invalid_python_code = [1, 2, 3')
 
-    await page.goto(
-      `http://127.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-    await page.waitForLoadState('networkidle')
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open customers.sql model
     await page
@@ -242,10 +227,7 @@ test.describe('Diagnostics for bad SQLMesh models', () => {
     )
     await fs.copy(customersSqlPath, duplicatedCustomersSqlPath)
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-    await page.waitForLoadState('networkidle')
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open full_model.sql model
     await page
@@ -334,10 +316,7 @@ test.describe('Diagnostics for bad audits', () => {
     await fs.writeFile(auditFilePath, updatedContent)
 
     // Navigate to the code-server instance
-    await page.goto(
-      `http://127.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-    await page.waitForLoadState('networkidle')
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open a the customers.sql model
     await page
diff --git a/vscode/extension/tests/find_references.spec.ts b/vscode/extension/tests/find_references.spec.ts
index 42dd0e6274..b952e30ef8 100644
--- a/vscode/extension/tests/find_references.spec.ts
+++ b/vscode/extension/tests/find_references.spec.ts
@@ -2,7 +2,12 @@ import { test, expect, Page } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { findAllReferences, goToReferences, SUSHI_SOURCE_PATH } from './utils'
+import {
+  findAllReferences,
+  goToReferences,
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 // Helper function to set up a test environment for model references
@@ -59,10 +64,7 @@ test.describe('Model References', () => {
     sharedCodeServer,
   }) => {
     const tempDir = await setupModelTestEnvironment()
-
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open customers.sql which contains references to other models
     await openCustomersFile(page)
@@ -130,9 +132,7 @@ test.describe('Model References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open customers.sql which contains multiple model references
     await openCustomersFile(page)
@@ -180,10 +180,7 @@ test.describe('Model References', () => {
     sharedCodeServer,
   }) => {
     const tempDir = await setupModelTestEnvironment()
-
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open assert_item_price_above_zero.sql audit file which references sushi.items model
     await navigateToAudits(page)
@@ -268,9 +265,7 @@ test.describe('Model References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open the audit file that validates item prices
     await navigateToAudits(page)
@@ -331,10 +326,7 @@ test.describe('CTE References', () => {
     sharedCodeServer,
   }) => {
     const tempDir = await setupModelTestEnvironment()
-
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     await openCustomersFile(page)
 
@@ -373,9 +365,7 @@ test.describe('CTE References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     await openCustomersFile(page)
 
@@ -414,10 +404,7 @@ test.describe('CTE References', () => {
     sharedCodeServer,
   }) => {
     const tempDir = await setupModelTestEnvironment()
-
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     await openCustomersFile(page)
 
@@ -455,10 +442,7 @@ test.describe('CTE References', () => {
   test('Find all references for CTE', async ({ page, sharedCodeServer }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-
+    await openServerPage(page, tempDir, sharedCodeServer)
     await openCustomersFile(page)
 
     // Click on the CTE definition "current_marketing_outer"
@@ -482,9 +466,7 @@ test.describe('CTE References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     await openCustomersFile(page)
 
@@ -509,10 +491,7 @@ test.describe('CTE References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-
+    await openServerPage(page, tempDir, sharedCodeServer)
     await openCustomersFile(page)
 
     // Click on the nested CTE "current_marketing" at line 33
@@ -541,10 +520,7 @@ test.describe('Macro References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-
+    await openServerPage(page, tempDir, sharedCodeServer)
     await openTopWaitersFile(page)
 
     // Click on the @ADD_ONE macro usage
@@ -578,10 +554,7 @@ test.describe('Macro References', () => {
     sharedCodeServer,
   }) => {
     const tempDir = await setupModelTestEnvironment()
-
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     await openTopWaitersFile(page)
 
@@ -635,10 +608,7 @@ test.describe('Macro References', () => {
   }) => {
     const tempDir = await setupModelTestEnvironment()
 
-    await page.goto(
-      `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-    )
-
+    await openServerPage(page, tempDir, sharedCodeServer)
     await openTopWaitersFile(page)
 
     // Click on the @SQL_LITERAL macro usage
diff --git a/vscode/extension/tests/format.spec.ts b/vscode/extension/tests/format.spec.ts
index 7525f88fdf..4e2b96bf94 100644
--- a/vscode/extension/tests/format.spec.ts
+++ b/vscode/extension/tests/format.spec.ts
@@ -2,7 +2,7 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { runCommand, SUSHI_SOURCE_PATH } from './utils'
+import { openServerPage, runCommand, SUSHI_SOURCE_PATH } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Format project works correctly', async ({ page, sharedCodeServer }) => {
@@ -10,10 +10,7 @@ test('Format project works correctly', async ({ page, sharedCodeServer }) => {
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   //   Wait for the models folder to be visible
   await page.waitForSelector('text=models')
diff --git a/vscode/extension/tests/go_to_definition.spec.ts b/vscode/extension/tests/go_to_definition.spec.ts
index 40d941669a..7d1749a1b7 100644
--- a/vscode/extension/tests/go_to_definition.spec.ts
+++ b/vscode/extension/tests/go_to_definition.spec.ts
@@ -2,18 +2,14 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { goToDefinition, SUSHI_SOURCE_PATH } from './utils'
+import { goToDefinition, openServerPage, SUSHI_SOURCE_PATH } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Stop server works', async ({ page, sharedCodeServer }) => {
   const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'))
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
   await createPythonInterpreterSettingsSpecifier(tempDir)
-
-  // Navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Wait for the models folder to be visible
   await page.waitForSelector('text=models')
@@ -47,9 +43,7 @@ test('Go to definition for model', async ({ page, sharedCodeServer }) => {
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
   // Navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Wait for the models folder to be visible
   await page.waitForSelector('text=models')
diff --git a/vscode/extension/tests/hints.spec.ts b/vscode/extension/tests/hints.spec.ts
index 13bf37fe8a..cb5dddb0eb 100644
--- a/vscode/extension/tests/hints.spec.ts
+++ b/vscode/extension/tests/hints.spec.ts
@@ -2,17 +2,14 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { SUSHI_SOURCE_PATH } from './utils'
+import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Model type hinting', async ({ page, sharedCodeServer }) => {
   const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'))
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  // Navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Wait for the models folder to be visible
   await page.waitForSelector('text=models')
diff --git a/vscode/extension/tests/lineage.spec.ts b/vscode/extension/tests/lineage.spec.ts
index 8f88c753f0..dfac5f5342 100644
--- a/vscode/extension/tests/lineage.spec.ts
+++ b/vscode/extension/tests/lineage.spec.ts
@@ -2,7 +2,7 @@ import { test, Page } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { openLineageView, SUSHI_SOURCE_PATH } from './utils'
+import { openLineageView, openServerPage, SUSHI_SOURCE_PATH } from './utils'
 import { writeFileSync } from 'fs'
 import {
   createPythonInterpreterSettingsSpecifier,
@@ -28,9 +28,7 @@ test('Lineage panel renders correctly - no project path config (default)', async
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
   await testLineageWithProjectPath(page)
 })
 
@@ -59,7 +57,7 @@ test.skip('Lineage panel renders correctly - relative project path', async ({
   )
 
   try {
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, workspaceDir, sharedCodeServer)
     await testLineageWithProjectPath(page)
   } finally {
     await fs.remove(workspaceDir)
@@ -91,7 +89,7 @@ test.skip('Lineage panel renders correctly - absolute project path', async ({
   )
 
   try {
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, tempDir, sharedCodeServer)
     await testLineageWithProjectPath(page)
   } finally {
     await stopCodeServer(context)
@@ -125,7 +123,7 @@ test.skip('Lineage panel renders correctly - relative project outside of workspa
   )
 
   try {
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, tempDir, sharedCodeServer)
     await testLineageWithProjectPath(page)
   } finally {
     await stopCodeServer(context)
@@ -156,9 +154,7 @@ test.skip('Lineage panel renders correctly - absolute path project outside of wo
     { spaces: 2 },
   )
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}?folder=${workspaceDir}`,
-  )
+  await openServerPage(page, workspaceDir, sharedCodeServer)
   await testLineageWithProjectPath(page)
 })
 
@@ -206,9 +202,7 @@ test.skip('Lineage panel renders correctly - multiworkspace setup', async ({
     { spaces: 2 },
   )
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}?folder=${workspaceDir}`,
-  )
+  await openServerPage(page, workspaceDir, sharedCodeServer)
   await page.waitForSelector('text=Open workspace')
   await page.click('text=Open workspace')
   await testLineageWithProjectPath(page)
@@ -216,6 +210,7 @@ test.skip('Lineage panel renders correctly - multiworkspace setup', async ({
 
 test.skip('Lineage panel renders correctly - multiworkspace setup reversed', async ({
   page,
+  sharedCodeServer,
 }) => {
   const workspaceDir = await fs.mkdtemp(
     path.join(os.tmpdir(), 'vscode-test-workspace-'),
@@ -267,7 +262,7 @@ test.skip('Lineage panel renders correctly - multiworkspace setup reversed', asy
   )
 
   try {
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, workspaceDir, sharedCodeServer)
     await page.waitForSelector('text=Open workspace')
     await page.click('text=Open workspace')
     await testLineageWithProjectPath(page)
diff --git a/vscode/extension/tests/lineage_settings.spec.ts b/vscode/extension/tests/lineage_settings.spec.ts
index f9d2e88781..21827e351f 100644
--- a/vscode/extension/tests/lineage_settings.spec.ts
+++ b/vscode/extension/tests/lineage_settings.spec.ts
@@ -2,7 +2,7 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { openLineageView, SUSHI_SOURCE_PATH } from './utils'
+import { openLineageView, openServerPage, SUSHI_SOURCE_PATH } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Settings button is visible in the lineage view', async ({
@@ -13,9 +13,7 @@ test('Settings button is visible in the lineage view', async ({
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   await page.waitForSelector('text=models')
 
diff --git a/vscode/extension/tests/python_env.spec.ts b/vscode/extension/tests/python_env.spec.ts
index bd5dfff3a3..04ccd2a7a0 100644
--- a/vscode/extension/tests/python_env.spec.ts
+++ b/vscode/extension/tests/python_env.spec.ts
@@ -3,6 +3,7 @@ import fs from 'fs-extra'
 import {
   createVirtualEnvironment,
   openLineageView,
+  openServerPage,
   pipInstall,
   PythonEnvironment,
   REPO_ROOT,
@@ -34,9 +35,7 @@ async function runTest(
   context: CodeServerContext,
   tempDir: string,
 ): Promise<void> {
-  await page.goto(
-    `http://127.0.0.1:${context.codeServerPort}` + `/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, context)
   await page.waitForSelector('text=models')
   await openLineageView(page)
 }
@@ -78,9 +77,7 @@ test.describe('python environment variable injection on sqlmesh_lsp', () => {
     await page.waitForSelector('text=Error creating context')
   })
 
-  test('normal setup - set', async ({ page, sharedCodeServer }, testInfo) => {
-    testInfo.setTimeout(120_000)
-
+  test('normal setup - set', async ({ page, sharedCodeServer }) => {
     const { tempDir } = await setupEnvironment()
     writeEnvironmentConfig(tempDir)
     const env_file = path.join(tempDir, '.env')
diff --git a/vscode/extension/tests/rename_cte.spec.ts b/vscode/extension/tests/rename_cte.spec.ts
index 0563f023b2..db7ab6bcb3 100644
--- a/vscode/extension/tests/rename_cte.spec.ts
+++ b/vscode/extension/tests/rename_cte.spec.ts
@@ -2,7 +2,12 @@ import { test, expect, Page } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { findAllReferences, renameSymbol, SUSHI_SOURCE_PATH } from './utils'
+import {
+  findAllReferences,
+  openServerPage,
+  renameSymbol,
+  SUSHI_SOURCE_PATH,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 async function setupTestEnvironment({
@@ -16,10 +21,7 @@ async function setupTestEnvironment({
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
-  // Navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Navigate to customers.sql which contains CTEs
   await page.waitForSelector('text=models')
diff --git a/vscode/extension/tests/render.spec.ts b/vscode/extension/tests/render.spec.ts
index 741d37ae14..8b5cec01ab 100644
--- a/vscode/extension/tests/render.spec.ts
+++ b/vscode/extension/tests/render.spec.ts
@@ -2,7 +2,12 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { openLineageView, runCommand, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openLineageView,
+  openServerPage,
+  runCommand,
+  SUSHI_SOURCE_PATH,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Render works correctly', async ({ page, sharedCodeServer }) => {
@@ -10,9 +15,7 @@ test('Render works correctly', async ({ page, sharedCodeServer }) => {
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   //   Wait for the models folder to be visible
   await page.waitForSelector('text=models')
@@ -48,9 +51,7 @@ test('Render works correctly with model without a description', async ({
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   //   Wait for the models folder to be visible
   await page.waitForSelector('text=models')
@@ -86,9 +87,7 @@ test('Render works correctly with every rendered model opening a new tab', async
   await fs.copy(SUSHI_SOURCE_PATH, tempDir)
 
   await createPythonInterpreterSettingsSpecifier(tempDir)
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Wait for the models folder to be visible
   await page.waitForSelector('text=models')
@@ -133,10 +132,7 @@ test('Render shows model picker when no active editor is open', async ({
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
   // Navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-  await page.waitForLoadState('networkidle')
+  await openServerPage(page, tempDir, sharedCodeServer)
 
   // Load the lineage view to initialize SQLMesh context (like lineage.spec.ts does)
   await openLineageView(page)
diff --git a/vscode/extension/tests/stop.spec.ts b/vscode/extension/tests/stop.spec.ts
index e9ba98fa03..12c3275a77 100644
--- a/vscode/extension/tests/stop.spec.ts
+++ b/vscode/extension/tests/stop.spec.ts
@@ -1,5 +1,5 @@
 import path from 'path'
-import { runCommand, SUSHI_SOURCE_PATH } from './utils'
+import { openServerPage, runCommand, SUSHI_SOURCE_PATH } from './utils'
 import os from 'os'
 import { test } from './fixtures'
 import fs from 'fs-extra'
@@ -12,12 +12,7 @@ test('Stop server works', async ({ page, sharedCodeServer }) => {
   await createPythonInterpreterSettingsSpecifier(tempDir)
 
   // Navigate to code-server instance
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
-
-  // Wait for code-server to load
-  await page.waitForLoadState('networkidle')
+  await openServerPage(page, tempDir, sharedCodeServer)
   await page.waitForSelector('[role="application"]', { timeout: 10000 })
 
   // Wait for the models folder to be visible in the file explorer
diff --git a/vscode/extension/tests/tcloud.spec.ts b/vscode/extension/tests/tcloud.spec.ts
index b677acaa27..ec334a3170 100644
--- a/vscode/extension/tests/tcloud.spec.ts
+++ b/vscode/extension/tests/tcloud.spec.ts
@@ -4,6 +4,7 @@ import fs from 'fs-extra'
 import os from 'os'
 import {
   createVirtualEnvironment,
+  openServerPage,
   pipInstall,
   REPO_ROOT,
   SUSHI_SOURCE_PATH,
@@ -38,7 +39,10 @@ async function setupPythonEnvironment(envDir: string): Promise<string> {
   return pythonDetails.pythonPath
 }
 
-test('not signed in, shows sign in window', async ({ page }) => {
+test('not signed in, shows sign in window', async ({
+  page,
+  sharedCodeServer,
+}) => {
   const tempDir = await fs.mkdtemp(
     path.join(os.tmpdir(), 'vscode-test-tcloud-'),
   )
@@ -79,8 +83,7 @@ test('not signed in, shows sign in window', async ({ page }) => {
       { spaces: 2 },
     )
 
-    // Start VS Code
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open a SQL file to trigger SQLMesh activation
     // Wait for the models folder to be visible
@@ -111,8 +114,8 @@ test('not signed in, shows sign in window', async ({ page }) => {
 
 test('signed in and not installed shows installation window', async ({
   page,
-}, testInfo) => {
-  testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
+  sharedCodeServer,
+}) => {
   const tempDir = await fs.mkdtemp(
     path.join(os.tmpdir(), 'vscode-test-tcloud-'),
   )
@@ -156,8 +159,7 @@ test('signed in and not installed shows installation window', async ({
       { spaces: 2 },
     )
 
-    // Start VS Code
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open a SQL file to trigger SQLMesh activation
     // Wait for the models folder to be visible
@@ -188,6 +190,7 @@ test('signed in and not installed shows installation window', async ({
 
 test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in old version when ready', async ({
   page,
+  sharedCodeServer,
 }) => {
   const tempDir = await fs.mkdtemp(
     path.join(os.tmpdir(), 'vscode-test-tcloud-'),
@@ -236,10 +239,7 @@ test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in old version when read
     )
 
     // Start VS Code
-    const context = await startCodeServer({
-      tempDir,
-    })
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open a SQL file to trigger SQLMesh activation
     // Wait for the models folder to be visible
@@ -267,8 +267,8 @@ test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in old version when read
 
 test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in new version when ready', async ({
   page,
-}, testInfo) => {
-  testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation
+  sharedCodeServer,
+}) => {
   const tempDir = await fs.mkdtemp(
     path.join(os.tmpdir(), 'vscode-test-tcloud-'),
   )
@@ -315,11 +315,7 @@ test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in new version when read
       { spaces: 2 },
     )
 
-    // Start VS Code
-    const context = await startCodeServer({
-      tempDir,
-    })
-    await page.goto(`http://127.0.0.1:${context.codeServerPort}`)
+    await openServerPage(page, tempDir, sharedCodeServer)
 
     // Open a SQL file to trigger SQLMesh activation
     // Wait for the models folder to be visible
diff --git a/vscode/extension/tests/utils.ts b/vscode/extension/tests/utils.ts
index d6a2ac7aaa..e783ae268e 100644
--- a/vscode/extension/tests/utils.ts
+++ b/vscode/extension/tests/utils.ts
@@ -1,6 +1,7 @@
 import path from 'path'
 import { Page } from '@playwright/test'
 import { execAsync } from '../src/utilities/exec'
+import { CodeServerContext } from './utils_code_server'
 
 // Where your extension lives on disk
 export const EXT_PATH = path.resolve(__dirname, '..')
@@ -163,8 +164,6 @@ export const goToReferences = async (page: Page): Promise<void> =>
 
 /**
  * Open the vscode code file picker and select the given file.
- * @param window The window to run the command in.
- * @param filePath The path to the file to select.
  */
 export const openFile = async (page: Page, file: string): Promise<void> => {
   const maxRetries = 3
@@ -196,3 +195,18 @@ export const openFile = async (page: Page, file: string): Promise<void> => {
     }
   }
 }
+
+/**
+ * Go to VSCode page
+ */
+export const openServerPage = async (
+  page: Page,
+  tempDir: string,
+  context: CodeServerContext,
+) => {
+  await page.goto(
+    `http://127.0.0.1:${context.codeServerPort}/?folder=${tempDir}`,
+  )
+  await page.waitForLoadState('networkidle')
+  await page.waitForSelector('[role="application"]', { timeout: 10000 })
+}
diff --git a/vscode/extension/tests/venv_naming.spec.ts b/vscode/extension/tests/venv_naming.spec.ts
index 4aeab2dcca..9e8006bff9 100644
--- a/vscode/extension/tests/venv_naming.spec.ts
+++ b/vscode/extension/tests/venv_naming.spec.ts
@@ -5,6 +5,7 @@ import path from 'path'
 import {
   createVirtualEnvironment,
   openLineageView,
+  openServerPage,
   pipInstall,
   REPO_ROOT,
   SUSHI_SOURCE_PATH,
@@ -34,9 +35,7 @@ test('venv being named .env', async ({ page, sharedCodeServer }) => {
     spaces: 2,
   })
 
-  await page.goto(
-    `http://127.0.0.1:${sharedCodeServer.codeServerPort}/?folder=${tempDir}`,
-  )
+  await openServerPage(page, tempDir, sharedCodeServer)
   await page.waitForSelector('text=models')
   await openLineageView(page)
   await page.waitForSelector('text=Loaded SQLMesh Context')