chore(vscode): refactor tests

sqlmesh/lsp/main.py

@@ -254,10 +254,7 @@ def _reload_context_and_publish_diagnostics(
                     self._ensure_context_for_document(uri)
                     # If successful, context_state will be ContextLoaded
                     if isinstance(self.context_state, ContextLoaded):
-                        ls.show_message(
-                            "Successfully loaded SQLMesh context",
-                            types.MessageType.Info,
-                        )
+                        loaded_sqlmesh_message(ls)
                 except Exception as e:
                     ls.log_trace(f"Still cannot load context: {e}")
                     # The error will be stored in context_state by _ensure_context_for_document
@@ -342,7 +339,7 @@ def initialize(ls: LanguageServer, params: types.InitializeParams) -> None:
                             config_path = folder_path / f"config.{ext}"
                             if config_path.exists():
                                 if self._create_lsp_context([folder_path]):
-                                    loaded_sqlmesh_message(ls, folder_path)
+                                    loaded_sqlmesh_message(ls)
                                     return  # Exit after successfully loading any config
             except Exception as e:
                 ls.log_trace(
@@ -863,15 +860,15 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
         try:
             if isinstance(self.context_state, NoContext):
                 context = self.context_class(paths=paths)
-                loaded_sqlmesh_message(self.server, paths[0])
+                loaded_sqlmesh_message(self.server)
             elif isinstance(self.context_state, ContextFailed):
                 if self.context_state.context:
                     context = self.context_state.context
                     context.load()
                 else:
                     # If there's no context (initial creation failed), try creating again
                     context = self.context_class(paths=paths)
-                    loaded_sqlmesh_message(self.server, paths[0])
+                    loaded_sqlmesh_message(self.server)
             else:
                 context = self.context_state.lsp_context.context
                 context.load()
@@ -908,9 +905,9 @@ def start(self) -> None:
         self.server.start_io()
 
 
-def loaded_sqlmesh_message(ls: LanguageServer, folder: Path) -> None:
+def loaded_sqlmesh_message(ls: LanguageServer) -> None:
     ls.show_message(
-        f"Loaded SQLMesh context from {folder}",
+        f"Loaded SQLMesh Context",
         types.MessageType.Info,
     )
 

vscode/extension/tests/bad_setup.spec.ts

@@ -10,6 +10,7 @@ import {
   pipInstall,
   REPO_ROOT,
   SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
 } from './utils'
 
 test('missing LSP dependencies shows install prompt', async ({
@@ -132,5 +133,5 @@ test.skip('check that the LSP runs correctly by opening lineage when looking at
   // Open the SQL file from the other directory
   await openFile(page, sqlFile)
 
-  await page.waitForSelector('text=Loaded SQLMesh context')
+  await waitForLoadedSQLMesh(page)
 })

vscode/extension/tests/broken_project.spec.ts

@@ -8,6 +8,7 @@ import {
   openProblemsView,
   saveFile,
   SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
 } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
@@ -64,7 +65,7 @@ test('working project, then broken through adding double model, then refixed', a
 
   // Open the lineage view to confirm it loads properly
   await openLineageView(page)
-  await page.waitForSelector('text=Loaded SQLMesh context')
+  await waitForLoadedSQLMesh(page)
 
   // Read the customers.sql file
   const customersSql = await fs.readFile(
@@ -292,6 +293,5 @@ test('bad model block, then fixed', async ({ page, sharedCodeServer }) => {
   await page.getByText('grain').click()
   await saveFile(page)
 
-  // Wait for successful context load
-  await page.waitForSelector('text=Loaded SQLMesh context')
+  await waitForLoadedSQLMesh(page)
 })

vscode/extension/tests/completions.spec.ts

@@ -2,7 +2,11 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Autocomplete for model names', async ({ page, sharedCodeServer }) => {
@@ -28,7 +32,7 @@ test('Autocomplete for model names', async ({ page, sharedCodeServer }) => {
     .click()
 
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 
   await page.locator('text=grain').first().click()
 
@@ -83,7 +87,7 @@ test.describe('Macro Completions', () => {
       .click()
 
     await page.waitForSelector('text=grain')
-    await page.waitForSelector('text=Loaded SQLMesh Context')
+    await waitForLoadedSQLMesh(page)
 
     await page.locator('text=grain').first().click()
 
@@ -133,7 +137,7 @@ test.describe('Macro Completions', () => {
       .click()
 
     await page.waitForSelector('text=grain')
-    await page.waitForSelector('text=Loaded SQLMesh Context')
+    await waitForLoadedSQLMesh(page)
 
     await page.locator('text=grain').first().click()
 

vscode/extension/tests/external_models.spec.ts

@@ -1,5 +1,9 @@
 import os from 'os'
-import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 import { test, expect } from './fixtures'
 import fs from 'fs-extra'
@@ -31,7 +35,7 @@ test.describe('External model files trigger lsp', () => {
       .click()
 
     await page.waitForSelector('text=raw.demographics')
-    await page.waitForSelector('text=Loaded SQLMesh Context')
+    await waitForLoadedSQLMesh(page)
   })
 
   test('external_models.yml', async ({ page, sharedCodeServer }) => {
@@ -63,6 +67,6 @@ test.describe('External model files trigger lsp', () => {
       .click()
 
     await page.waitForSelector('text=raw.demographics')
-    await page.waitForSelector('text=Loaded SQLMesh Context')
+    await waitForLoadedSQLMesh(page)
   })
 })

vscode/extension/tests/find_references.spec.ts

@@ -7,6 +7,7 @@ import {
   goToReferences,
   openServerPage,
   SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
 } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
@@ -44,7 +45,7 @@ async function openCustomersFile(page: Page) {
     .locator('a')
     .click()
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 }
 
 // Helper function to open top_waiters.sql and wait for SQLMesh context
@@ -55,7 +56,7 @@ async function openTopWaitersFile(page: Page) {
     .locator('a')
     .click()
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 }
 
 test.describe('Model References', () => {
@@ -194,7 +195,7 @@ test.describe('Model References', () => {
 
     // Wait for audit file to load and SQLMesh context to initialize
     await page.waitForSelector('text=standalone')
-    await page.waitForSelector('text=Loaded SQLMesh Context')
+    await waitForLoadedSQLMesh(page)
 
     // Step 4: Click on sushi.items model reference in the audit query
     await page.locator('text=sushi.items').first().click()
@@ -279,7 +280,7 @@ test.describe('Model References', () => {
 
     // Ensure audit file and SQLMesh context are fully loaded
     await page.waitForSelector('text=standalone')
-    await page.waitForSelector('text=Loaded SQLMesh Context')
+    await waitForLoadedSQLMesh(page)
 
     // Step 4: Position cursor on sushi.items model reference
     await page.locator('text=sushi.items').first().click()

vscode/extension/tests/format.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 { openServerPage, runCommand, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openServerPage,
+  runCommand,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Format project works correctly', async ({ page, sharedCodeServer }) => {
@@ -28,7 +33,7 @@ test('Format project works correctly', async ({ page, sharedCodeServer }) => {
     .click()
 
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 
   // Format the project
   await runCommand(page, 'SQLMesh: Format Project')

vscode/extension/tests/go_to_definition.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 { goToDefinition, openServerPage, SUSHI_SOURCE_PATH } from './utils'
+import {
+  goToDefinition,
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Stop server works', async ({ page, sharedCodeServer }) => {
@@ -27,7 +32,7 @@ test('Stop server works', async ({ page, sharedCodeServer }) => {
     .click()
 
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 
   // Render the model
   await page.locator('text=@MULTIPLY').click()
@@ -61,7 +66,7 @@ test('Go to definition for model', async ({ page, sharedCodeServer }) => {
     .click()
 
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 
   // Go to definition for the model
   await page.locator('text=sushi.waiter_revenue_by_day').first().click()

vscode/extension/tests/hints.spec.ts

@@ -2,7 +2,11 @@ import { test, expect } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Model type hinting', async ({ page, sharedCodeServer }) => {
@@ -30,7 +34,7 @@ test('Model type hinting', async ({ page, sharedCodeServer }) => {
     .click()
 
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 
   // Wait a moment for hints to appear
   await page.waitForTimeout(500)

vscode/extension/tests/lineage.spec.ts

@@ -2,7 +2,12 @@ import { test, Page } from './fixtures'
 import path from 'path'
 import fs from 'fs-extra'
 import os from 'os'
-import { openLineageView, openServerPage, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openLineageView,
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { writeFileSync } from 'fs'
 import {
   createPythonInterpreterSettingsSpecifier,
@@ -17,7 +22,7 @@ async function testLineageWithProjectPath(page: Page): Promise<void> {
   await page.waitForLoadState('networkidle')
   await page.waitForLoadState('domcontentloaded')
   await openLineageView(page)
-  await page.waitForSelector('text=Loaded SQLMesh context')
+  await waitForLoadedSQLMesh(page)
 }
 
 test('Lineage panel renders correctly - no project path config (default)', async ({

vscode/extension/tests/lineage_settings.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, openServerPage, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openLineageView,
+  openServerPage,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
 test('Settings button is visible in the lineage view', async ({
@@ -27,7 +32,7 @@ test('Settings button is visible in the lineage view', async ({
     .getByRole('treeitem', { name: 'waiters.py', exact: true })
     .locator('a')
     .click()
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 
   // Open lineage
   await openLineageView(page)

vscode/extension/tests/python_env.spec.ts

@@ -8,6 +8,7 @@ import {
   PythonEnvironment,
   REPO_ROOT,
   SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
 } from './utils'
 import os from 'os'
 import path from 'path'
@@ -83,7 +84,7 @@ test.describe('python environment variable injection on sqlmesh_lsp', () => {
     const env_file = path.join(tempDir, '.env')
     fs.writeFileSync(env_file, 'TEST_VAR=test_value')
     await runTest(page, sharedCodeServer, tempDir)
-    await page.waitForSelector('text=Loaded SQLMesh context')
+    await waitForLoadedSQLMesh(page)
   })
 })
 
@@ -127,6 +128,6 @@ test.describe('tcloud version', () => {
     const env_file = path.join(tempDir, '.env')
     fs.writeFileSync(env_file, 'TEST_VAR=test_value')
     await runTest(page, sharedCodeServer, tempDir)
-    await page.waitForSelector('text=Loaded SQLMesh context')
+    await waitForLoadedSQLMesh(page)
   })
 })

vscode/extension/tests/quickfix.spec.ts

@@ -1,7 +1,11 @@
 import fs from 'fs-extra'
 import path from 'path'
 import os from 'os'
-import { openProblemsView, SUSHI_SOURCE_PATH } from './utils'
+import {
+  openProblemsView,
+  SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
+} from './utils'
 import { test, expect } from './fixtures'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
@@ -66,7 +70,7 @@ test('noselectstar quickfix', async ({ page, sharedCodeServer }) => {
     .locator('a')
     .click()
 
-  await page.waitForSelector('text=Loaded SQLMesh context')
+  await waitForLoadedSQLMesh(page)
 
   await openProblemsView(page)
 

vscode/extension/tests/rename_cte.spec.ts

@@ -7,6 +7,7 @@ import {
   openServerPage,
   renameSymbol,
   SUSHI_SOURCE_PATH,
+  waitForLoadedSQLMesh,
 } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
@@ -34,7 +35,7 @@ async function setupTestEnvironment({
     .locator('a')
     .click()
   await page.waitForSelector('text=grain')
-  await page.waitForSelector('text=Loaded SQLMesh Context')
+  await waitForLoadedSQLMesh(page)
 }
 
 test.describe('CTE Rename', () => {