Skip to content

Commit edbe39a

Browse files
benfdkingbenfdking
committed
chore(vscode): refactor tests
1 parent 29833e7 commit edbe39a

19 files changed

+102
-52
lines changed

‎sqlmesh/lsp/main.py‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -254,10 +254,7 @@ def _reload_context_and_publish_diagnostics(
254254
self._ensure_context_for_document(uri)
255255
# If successful, context_state will be ContextLoaded
256256
if isinstance(self.context_state, ContextLoaded):
257-
ls.show_message(
258-
"Successfully loaded SQLMesh context",
259-
types.MessageType.Info,
260-
)
257+
loaded_sqlmesh_message(ls)
261258
except Exception as e:
262259
ls.log_trace(f"Still cannot load context: {e}")
263260
# 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:
342339
config_path = folder_path / f"config.{ext}"
343340
if config_path.exists():
344341
if self._create_lsp_context([folder_path]):
345-
loaded_sqlmesh_message(ls, folder_path)
342+
loaded_sqlmesh_message(ls)
346343
return # Exit after successfully loading any config
347344
except Exception as e:
348345
ls.log_trace(
@@ -863,15 +860,15 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
863860
try:
864861
if isinstance(self.context_state, NoContext):
865862
context = self.context_class(paths=paths)
866-
loaded_sqlmesh_message(self.server, paths[0])
863+
loaded_sqlmesh_message(self.server)
867864
elif isinstance(self.context_state, ContextFailed):
868865
if self.context_state.context:
869866
context = self.context_state.context
870867
context.load()
871868
else:
872869
# If there's no context (initial creation failed), try creating again
873870
context = self.context_class(paths=paths)
874-
loaded_sqlmesh_message(self.server, paths[0])
871+
loaded_sqlmesh_message(self.server)
875872
else:
876873
context = self.context_state.lsp_context.context
877874
context.load()
@@ -908,9 +905,9 @@ def start(self) -> None:
908905
self.server.start_io()
909906

910907

911-
def loaded_sqlmesh_message(ls: LanguageServer, folder: Path) -> None:
908+
def loaded_sqlmesh_message(ls: LanguageServer) -> None:
912909
ls.show_message(
913-
f"Loaded SQLMesh context from {folder}",
910+
f"Loaded SQLMesh Context",
914911
types.MessageType.Info,
915912
)
916913

‎vscode/extension/tests/bad_setup.spec.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
pipInstall,
1111
REPO_ROOT,
1212
SUSHI_SOURCE_PATH,
13+
waitForLoadedSQLMesh,
1314
} from './utils'
1415

1516
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
132133
// Open the SQL file from the other directory
133134
await openFile(page, sqlFile)
134135

135-
await page.waitForSelector('text=Loaded SQLMesh context')
136+
await waitForLoadedSQLMesh(page)
136137
})

‎vscode/extension/tests/broken_project.spec.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
openProblemsView,
99
saveFile,
1010
SUSHI_SOURCE_PATH,
11+
waitForLoadedSQLMesh,
1112
} from './utils'
1213
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
1314

@@ -64,7 +65,7 @@ test('working project, then broken through adding double model, then refixed', a
6465

6566
// Open the lineage view to confirm it loads properly
6667
await openLineageView(page)
67-
await page.waitForSelector('text=Loaded SQLMesh context')
68+
await waitForLoadedSQLMesh(page)
6869

6970
// Read the customers.sql file
7071
const customersSql = await fs.readFile(
@@ -292,6 +293,5 @@ test('bad model block, then fixed', async ({ page, sharedCodeServer }) => {
292293
await page.getByText('grain').click()
293294
await saveFile(page)
294295

295-
// Wait for successful context load
296-
await page.waitForSelector('text=Loaded SQLMesh context')
296+
await waitForLoadedSQLMesh(page)
297297
})

‎vscode/extension/tests/completions.spec.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@ import { test, expect } from './fixtures'
22
import path from 'path'
33
import fs from 'fs-extra'
44
import os from 'os'
5-
import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
5+
import {
6+
openServerPage,
7+
SUSHI_SOURCE_PATH,
8+
waitForLoadedSQLMesh,
9+
} from './utils'
610
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
711

812
test('Autocomplete for model names', async ({ page, sharedCodeServer }) => {
@@ -28,7 +32,7 @@ test('Autocomplete for model names', async ({ page, sharedCodeServer }) => {
2832
.click()
2933

3034
await page.waitForSelector('text=grain')
31-
await page.waitForSelector('text=Loaded SQLMesh Context')
35+
await waitForLoadedSQLMesh(page)
3236

3337
await page.locator('text=grain').first().click()
3438

@@ -83,7 +87,7 @@ test.describe('Macro Completions', () => {
8387
.click()
8488

8589
await page.waitForSelector('text=grain')
86-
await page.waitForSelector('text=Loaded SQLMesh Context')
90+
await waitForLoadedSQLMesh(page)
8791

8892
await page.locator('text=grain').first().click()
8993

@@ -133,7 +137,7 @@ test.describe('Macro Completions', () => {
133137
.click()
134138

135139
await page.waitForSelector('text=grain')
136-
await page.waitForSelector('text=Loaded SQLMesh Context')
140+
await waitForLoadedSQLMesh(page)
137141

138142
await page.locator('text=grain').first().click()
139143

‎vscode/extension/tests/external_models.spec.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import os from 'os'
2-
import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
2+
import {
3+
openServerPage,
4+
SUSHI_SOURCE_PATH,
5+
waitForLoadedSQLMesh,
6+
} from './utils'
37
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
48
import { test, expect } from './fixtures'
59
import fs from 'fs-extra'
@@ -31,7 +35,7 @@ test.describe('External model files trigger lsp', () => {
3135
.click()
3236

3337
await page.waitForSelector('text=raw.demographics')
34-
await page.waitForSelector('text=Loaded SQLMesh Context')
38+
await waitForLoadedSQLMesh(page)
3539
})
3640

3741
test('external_models.yml', async ({ page, sharedCodeServer }) => {
@@ -63,6 +67,6 @@ test.describe('External model files trigger lsp', () => {
6367
.click()
6468

6569
await page.waitForSelector('text=raw.demographics')
66-
await page.waitForSelector('text=Loaded SQLMesh Context')
70+
await waitForLoadedSQLMesh(page)
6771
})
6872
})

‎vscode/extension/tests/find_references.spec.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
goToReferences,
88
openServerPage,
99
SUSHI_SOURCE_PATH,
10+
waitForLoadedSQLMesh,
1011
} from './utils'
1112
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
1213

@@ -44,7 +45,7 @@ async function openCustomersFile(page: Page) {
4445
.locator('a')
4546
.click()
4647
await page.waitForSelector('text=grain')
47-
await page.waitForSelector('text=Loaded SQLMesh Context')
48+
await waitForLoadedSQLMesh(page)
4849
}
4950

5051
// Helper function to open top_waiters.sql and wait for SQLMesh context
@@ -55,7 +56,7 @@ async function openTopWaitersFile(page: Page) {
5556
.locator('a')
5657
.click()
5758
await page.waitForSelector('text=grain')
58-
await page.waitForSelector('text=Loaded SQLMesh Context')
59+
await waitForLoadedSQLMesh(page)
5960
}
6061

6162
test.describe('Model References', () => {
@@ -194,7 +195,7 @@ test.describe('Model References', () => {
194195

195196
// Wait for audit file to load and SQLMesh context to initialize
196197
await page.waitForSelector('text=standalone')
197-
await page.waitForSelector('text=Loaded SQLMesh Context')
198+
await waitForLoadedSQLMesh(page)
198199

199200
// Step 4: Click on sushi.items model reference in the audit query
200201
await page.locator('text=sushi.items').first().click()
@@ -279,7 +280,7 @@ test.describe('Model References', () => {
279280

280281
// Ensure audit file and SQLMesh context are fully loaded
281282
await page.waitForSelector('text=standalone')
282-
await page.waitForSelector('text=Loaded SQLMesh Context')
283+
await waitForLoadedSQLMesh(page)
283284

284285
// Step 4: Position cursor on sushi.items model reference
285286
await page.locator('text=sushi.items').first().click()

‎vscode/extension/tests/format.spec.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import { test, expect } from './fixtures'
22
import path from 'path'
33
import fs from 'fs-extra'
44
import os from 'os'
5-
import { openServerPage, runCommand, SUSHI_SOURCE_PATH } from './utils'
5+
import {
6+
openServerPage,
7+
runCommand,
8+
SUSHI_SOURCE_PATH,
9+
waitForLoadedSQLMesh,
10+
} from './utils'
611
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
712

813
test('Format project works correctly', async ({ page, sharedCodeServer }) => {
@@ -28,7 +33,7 @@ test('Format project works correctly', async ({ page, sharedCodeServer }) => {
2833
.click()
2934

3035
await page.waitForSelector('text=grain')
31-
await page.waitForSelector('text=Loaded SQLMesh Context')
36+
await waitForLoadedSQLMesh(page)
3237

3338
// Format the project
3439
await runCommand(page, 'SQLMesh: Format Project')

‎vscode/extension/tests/go_to_definition.spec.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import { test, expect } from './fixtures'
22
import path from 'path'
33
import fs from 'fs-extra'
44
import os from 'os'
5-
import { goToDefinition, openServerPage, SUSHI_SOURCE_PATH } from './utils'
5+
import {
6+
goToDefinition,
7+
openServerPage,
8+
SUSHI_SOURCE_PATH,
9+
waitForLoadedSQLMesh,
10+
} from './utils'
611
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
712

813
test('Stop server works', async ({ page, sharedCodeServer }) => {
@@ -27,7 +32,7 @@ test('Stop server works', async ({ page, sharedCodeServer }) => {
2732
.click()
2833

2934
await page.waitForSelector('text=grain')
30-
await page.waitForSelector('text=Loaded SQLMesh Context')
35+
await waitForLoadedSQLMesh(page)
3136

3237
// Render the model
3338
await page.locator('text=@MULTIPLY').click()
@@ -61,7 +66,7 @@ test('Go to definition for model', async ({ page, sharedCodeServer }) => {
6166
.click()
6267

6368
await page.waitForSelector('text=grain')
64-
await page.waitForSelector('text=Loaded SQLMesh Context')
69+
await waitForLoadedSQLMesh(page)
6570

6671
// Go to definition for the model
6772
await page.locator('text=sushi.waiter_revenue_by_day').first().click()

‎vscode/extension/tests/hints.spec.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@ import { test, expect } from './fixtures'
22
import path from 'path'
33
import fs from 'fs-extra'
44
import os from 'os'
5-
import { openServerPage, SUSHI_SOURCE_PATH } from './utils'
5+
import {
6+
openServerPage,
7+
SUSHI_SOURCE_PATH,
8+
waitForLoadedSQLMesh,
9+
} from './utils'
610
import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
711

812
test('Model type hinting', async ({ page, sharedCodeServer }) => {
@@ -30,7 +34,7 @@ test('Model type hinting', async ({ page, sharedCodeServer }) => {
3034
.click()
3135

3236
await page.waitForSelector('text=grain')
33-
await page.waitForSelector('text=Loaded SQLMesh Context')
37+
await waitForLoadedSQLMesh(page)
3438

3539
// Wait a moment for hints to appear
3640
await page.waitForTimeout(500)

‎vscode/extension/tests/lineage.spec.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import { test, Page } from './fixtures'
22
import path from 'path'
33
import fs from 'fs-extra'
44
import os from 'os'
5-
import { openLineageView, openServerPage, SUSHI_SOURCE_PATH } from './utils'
5+
import {
6+
openLineageView,
7+
openServerPage,
8+
SUSHI_SOURCE_PATH,
9+
waitForLoadedSQLMesh,
10+
} from './utils'
611
import { writeFileSync } from 'fs'
712
import {
813
createPythonInterpreterSettingsSpecifier,
@@ -17,7 +22,7 @@ async function testLineageWithProjectPath(page: Page): Promise<void> {
1722
await page.waitForLoadState('networkidle')
1823
await page.waitForLoadState('domcontentloaded')
1924
await openLineageView(page)
20-
await page.waitForSelector('text=Loaded SQLMesh context')
25+
await waitForLoadedSQLMesh(page)
2126
}
2227

2328
test('Lineage panel renders correctly - no project path config (default)', async ({

0 commit comments

Comments
 (0)