diff --git a/vscode/extension/src/commands/format.ts b/vscode/extension/src/commands/format.ts index 5e3465921a..bedeaa16be 100644 --- a/vscode/extension/src/commands/format.ts +++ b/vscode/extension/src/commands/format.ts @@ -11,12 +11,18 @@ export const format = ( authProvider: AuthenticationProviderTobikoCloud, lsp: LSPClient | undefined, + restartLSP: () => Promise, ) => async (): Promise => { traceLog('Calling format') const out = await internalFormat(lsp) if (isErr(out)) { - return handleError(authProvider, out.error, 'Project format failed') + return handleError( + authProvider, + restartLSP, + out.error, + 'Project format failed', + ) } vscode.window.showInformationMessage('Project formatted successfully') } diff --git a/vscode/extension/src/commands/signin.ts b/vscode/extension/src/commands/signin.ts index e59c2b2161..77131a8253 100644 --- a/vscode/extension/src/commands/signin.ts +++ b/vscode/extension/src/commands/signin.ts @@ -1,13 +1,30 @@ import { AuthenticationProviderTobikoCloud } from '../auth/auth' import * as vscode from 'vscode' import { isCodespaces } from '../utilities/isCodespaces' +import { traceInfo } from '../utilities/common/log' export const signIn = - (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => { + ( + authenticationProvider: AuthenticationProviderTobikoCloud, + onSignInSuccess: () => Promise, + ) => + async () => { if (isCodespaces()) { await authenticationProvider.sign_in_device_flow() } else { await authenticationProvider.createSession() } - await vscode.window.showInformationMessage('Signed in successfully') + + // Do not await this, as this will block the thread, you just need to show the message, but not block + vscode.window.showInformationMessage('Signed in successfully') + + // Execute callback after successful sign-in + if (onSignInSuccess) { + traceInfo('Executing post sign-in callback') + try { + await onSignInSuccess() + } catch (error) { + traceInfo(`Error in post sign-in callback: ${error}`) + } + } } diff --git a/vscode/extension/src/commands/signinSpecifyFlow.ts b/vscode/extension/src/commands/signinSpecifyFlow.ts index 8e277b28b0..2e0c0cfe15 100644 --- a/vscode/extension/src/commands/signinSpecifyFlow.ts +++ b/vscode/extension/src/commands/signinSpecifyFlow.ts @@ -3,7 +3,11 @@ import { traceInfo } from '../utilities/common/log' import { window } from 'vscode' export const signInSpecifyFlow = - (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => { + ( + authenticationProvider: AuthenticationProviderTobikoCloud, + onSignInSuccess?: () => Promise, + ) => + async () => { traceInfo('Sign in specify flow') const flowOptions = [ { @@ -24,11 +28,31 @@ export const signInSpecifyFlow = await authenticationProvider.sign_in_oauth_flow() await authenticationProvider.getSessions() await window.showInformationMessage('Sign in success') + + // Execute callback after successful sign-in + if (onSignInSuccess) { + traceInfo('Executing post sign-in callback') + try { + await onSignInSuccess() + } catch (error) { + traceInfo(`Error in post sign-in callback: ${error}`) + } + } return } else if (selectedFlow.label === 'Device Flow') { await authenticationProvider.sign_in_device_flow() await authenticationProvider.getSessions() await window.showInformationMessage('Sign in success') + + // Execute callback after successful sign-in + if (onSignInSuccess) { + traceInfo('Executing post sign-in callback') + try { + await onSignInSuccess() + } catch (error) { + traceInfo(`Error in post sign-in callback: ${error}`) + } + } return } else { traceInfo('Invalid flow selected') diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts index 8749c61fb2..8e4b57c907 100644 --- a/vscode/extension/src/extension.ts +++ b/vscode/extension/src/extension.ts @@ -45,12 +45,21 @@ export async function activate(context: vscode.ExtensionContext) { traceInfo('Authentication provider registered') context.subscriptions.push( - vscode.commands.registerCommand('sqlmesh.signin', signIn(authProvider)), + vscode.commands.registerCommand( + 'sqlmesh.signin', + signIn(authProvider, async () => { + traceInfo('Restarting LSP after sign-in') + await restart() + }), + ), ) context.subscriptions.push( vscode.commands.registerCommand( 'sqlmesh.signinSpecifyFlow', - signInSpecifyFlow(authProvider), + signInSpecifyFlow(authProvider, async () => { + traceInfo('Restarting LSP after sign-in') + await restart() + }), ), ) context.subscriptions.push( @@ -76,13 +85,6 @@ export async function activate(context: vscode.ExtensionContext) { ), ) - context.subscriptions.push( - vscode.commands.registerCommand( - 'sqlmesh.format', - format(authProvider, lspClient), - ), - ) - // Register the webview const lineagePanel = new LineagePanel(context.extensionUri, lspClient) context.subscriptions.push( @@ -118,14 +120,35 @@ export async function activate(context: vscode.ExtensionContext) { if (isErr(restartResult)) { return handleError( authProvider, + restart, restartResult.error, 'LSP restart failed', ) } context.subscriptions.push(lspClient) + } else { + lspClient = new LSPClient() + const result = await lspClient.start() + if (isErr(result)) { + return handleError( + authProvider, + restart, + result.error, + 'Failed to start LSP', + ) + } else { + context.subscriptions.push(lspClient) + } } } + context.subscriptions.push( + vscode.commands.registerCommand( + 'sqlmesh.format', + format(authProvider, lspClient, restart), + ), + ) + context.subscriptions.push( onDidChangePythonInterpreter(async () => { await restart() @@ -140,7 +163,12 @@ export async function activate(context: vscode.ExtensionContext) { const result = await lspClient.start() if (isErr(result)) { - return handleError(authProvider, result.error, 'Failed to start LSP') + return handleError( + authProvider, + restart, + result.error, + 'Failed to start LSP', + ) } else { context.subscriptions.push(lspClient) } diff --git a/vscode/extension/src/utilities/errors.ts b/vscode/extension/src/utilities/errors.ts index 62d0d016b9..f062d89369 100644 --- a/vscode/extension/src/utilities/errors.ts +++ b/vscode/extension/src/utilities/errors.ts @@ -69,6 +69,7 @@ interface SqlmeshLspDependenciesMissingError { export async function handleError( authProvider: AuthenticationProviderTobikoCloud, + restartLsp: () => Promise, error: ErrorType, genericErrorPrefix?: string, ): Promise { @@ -85,7 +86,7 @@ export async function handleError( ) return case 'not_signed_in': - return handleNotSignedInError(authProvider) + return handleNotSignedInError(authProvider, restartLsp) case 'sqlmesh_not_found': return handleSqlmeshNotFoundError() case 'sqlmesh_lsp_not_found': @@ -110,6 +111,7 @@ export async function handleError( */ const handleNotSignedInError = async ( authProvider: AuthenticationProviderTobikoCloud, + restartLsp: () => Promise, ): Promise => { traceInfo('handleNotSginedInError') const result = await window.showInformationMessage( @@ -117,7 +119,7 @@ const handleNotSignedInError = async ( 'Sign In', ) if (result === 'Sign In') { - await signIn(authProvider)() + await signIn(authProvider, restartLsp)() } } diff --git a/vscode/extension/tests/tcloud.spec.ts b/vscode/extension/tests/tcloud.spec.ts index f01dfc1a33..20629f5b79 100644 --- a/vscode/extension/tests/tcloud.spec.ts +++ b/vscode/extension/tests/tcloud.spec.ts @@ -34,311 +34,390 @@ async function setupPythonEnvironment(envDir: string): Promise { return pythonDetails.pythonPath } -test.describe('Tcloud', () => { - test('not signed in, shows sign in window', async ({}, testInfo) => { - testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation - const tempDir = await fs.mkdtemp( - path.join(os.tmpdir(), 'vscode-test-tcloud-'), +test('not signed in, shows sign in window', async ({}, testInfo) => { + testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation + const tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'vscode-test-tcloud-'), + ) + const pythonEnvDir = path.join(tempDir, '.venv') + + try { + // Copy sushi project + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + + // Create a tcloud.yaml to mark this as a tcloud project + const tcloudConfig = { + url: 'https://mock.tobikodata.com', + org: 'test-org', + project: 'test-project', + } + await fs.writeFile( + path.join(tempDir, 'tcloud.yaml'), + `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, ) - const pythonEnvDir = path.join(tempDir, '.venv') - - try { - // Copy sushi project - await fs.copy(SUSHI_SOURCE_PATH, tempDir) - - // Create a tcloud.yaml to mark this as a tcloud project - const tcloudConfig = { - url: 'https://mock.tobikodata.com', - org: 'test-org', - project: 'test-project', - } - await fs.writeFile( - path.join(tempDir, 'tcloud.yaml'), - `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, - ) - - // Set tcloud version to 2.10.0 - await setTcloudVersion(tempDir, '2.10.0') - - // Set up Python environment with mock tcloud and sqlmesh - const pythonPath = await setupPythonEnvironment(pythonEnvDir) - - // Configure VS Code settings to use our Python environment - const settings = { - 'python.defaultInterpreterPath': pythonPath, - 'sqlmesh.environmentPath': pythonEnvDir, - } - await fs.ensureDir(path.join(tempDir, '.vscode')) - await fs.writeJson( - path.join(tempDir, '.vscode', 'settings.json'), - settings, - { spaces: 2 }, - ) - - // Start VS Code - const { window, close } = await startVSCode(tempDir) - - // Open a SQL file to trigger SQLMesh activation - // Wait for the models folder to be visible - await window.waitForSelector('text=models') - - // Click on the models folder - await window - .getByRole('treeitem', { name: 'models', exact: true }) - .locator('a') - .click() - - // Open the top_waiters model - await window - .getByRole('treeitem', { name: 'customers.sql', exact: true }) - .locator('a') - .click() - - // Wait for the file to open - await window.waitForTimeout(2000) - - await window.waitForSelector( - 'text=Please sign in to Tobiko Cloud to use SQLMesh', - ) - - // Close VS Code - await close() - } finally { - // Clean up - await fs.remove(tempDir) + + // Set tcloud version to 2.10.0 + await setTcloudVersion(tempDir, '2.10.0') + + // Set up Python environment with mock tcloud and sqlmesh + const pythonPath = await setupPythonEnvironment(pythonEnvDir) + + // Configure VS Code settings to use our Python environment + const settings = { + 'python.defaultInterpreterPath': pythonPath, + 'sqlmesh.environmentPath': pythonEnvDir, } - }) + await fs.ensureDir(path.join(tempDir, '.vscode')) + await fs.writeJson( + path.join(tempDir, '.vscode', 'settings.json'), + settings, + { spaces: 2 }, + ) - test('signed in and not installed shows installation window', async ({}, testInfo) => { - testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation - const tempDir = await fs.mkdtemp( - path.join(os.tmpdir(), 'vscode-test-tcloud-'), + // Start VS Code + const { window, close } = await startVSCode(tempDir) + + // Open a SQL file to trigger SQLMesh activation + // Wait for the models folder to be visible + await window.waitForSelector('text=models') + + // Click on the models folder + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + + // Open the top_waiters model + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + + // Wait for the file to open + await window.waitForTimeout(2000) + + await window.waitForSelector( + 'text=Please sign in to Tobiko Cloud to use SQLMesh', ) - const pythonEnvDir = path.join(tempDir, '.venv') - - try { - // Copy sushi project - await fs.copy(SUSHI_SOURCE_PATH, tempDir) - - // Create a tcloud.yaml to mark this as a tcloud project - const tcloudConfig = { - url: 'https://mock.tobikodata.com', - org: 'test-org', - project: 'test-project', - } - await fs.writeFile( - path.join(tempDir, 'tcloud.yaml'), - `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, - ) - - // Write mock ".tcloud_auth_state.json" file - await setupAuthenticatedState(tempDir) - - // Set tcloud version to 2.10.0 - await setTcloudVersion(tempDir, '2.10.0') - - // Set up Python environment with mock tcloud and sqlmesh - const pythonPath = await setupPythonEnvironment(pythonEnvDir) - - // Configure VS Code settings to use our Python environment - const settings = { - 'python.defaultInterpreterPath': pythonPath, - 'sqlmesh.environmentPath': pythonEnvDir, - } - await fs.ensureDir(path.join(tempDir, '.vscode')) - await fs.writeJson( - path.join(tempDir, '.vscode', 'settings.json'), - settings, - { spaces: 2 }, - ) - - // Start VS Code - const { window, close } = await startVSCode(tempDir) - - // Open a SQL file to trigger SQLMesh activation - // Wait for the models folder to be visible - await window.waitForSelector('text=models') - - // Click on the models folder - await window - .getByRole('treeitem', { name: 'models', exact: true }) - .locator('a') - .click() - - // Open the top_waiters model - await window - .getByRole('treeitem', { name: 'customers.sql', exact: true }) - .locator('a') - .click() - - await window.waitForSelector('text=Installing enterprise python package') - expect( - await window.locator('text=Installing enterprise python package'), - ).toHaveCount(2) - - await window.waitForSelector('text=Loaded SQLMesh context') - - // Close VS Code - await close() - } finally { - // Clean up - await fs.remove(tempDir) + + // Close VS Code + await close() + } finally { + // Clean up + await fs.remove(tempDir) + } +}) + +test('signed in and not installed shows installation window', async ({}, testInfo) => { + testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation + const tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'vscode-test-tcloud-'), + ) + const pythonEnvDir = path.join(tempDir, '.venv') + + try { + // Copy sushi project + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + + // Create a tcloud.yaml to mark this as a tcloud project + const tcloudConfig = { + url: 'https://mock.tobikodata.com', + org: 'test-org', + project: 'test-project', } - }) + await fs.writeFile( + path.join(tempDir, 'tcloud.yaml'), + `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, + ) - test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in old version when ready', async ({}, testInfo) => { - testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation - const tempDir = await fs.mkdtemp( - path.join(os.tmpdir(), 'vscode-test-tcloud-'), + // Write mock ".tcloud_auth_state.json" file + await setupAuthenticatedState(tempDir) + + // Set tcloud version to 2.10.0 + await setTcloudVersion(tempDir, '2.10.0') + + // Set up Python environment with mock tcloud and sqlmesh + const pythonPath = await setupPythonEnvironment(pythonEnvDir) + + // Configure VS Code settings to use our Python environment + const settings = { + 'python.defaultInterpreterPath': pythonPath, + 'sqlmesh.environmentPath': pythonEnvDir, + } + await fs.ensureDir(path.join(tempDir, '.vscode')) + await fs.writeJson( + path.join(tempDir, '.vscode', 'settings.json'), + settings, + { spaces: 2 }, ) - const pythonEnvDir = path.join(tempDir, '.venv') - - try { - // Copy sushi project - await fs.copy(SUSHI_SOURCE_PATH, tempDir) - - // Create a tcloud.yaml to mark this as a tcloud project - const tcloudConfig = { - url: 'https://mock.tobikodata.com', - org: 'test-org', - project: 'test-project', - } - await fs.writeFile( - path.join(tempDir, 'tcloud.yaml'), - `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, - ) - - // Write mock ".tcloud_auth_state.json" file - await setupAuthenticatedState(tempDir) - - // Set tcloud version to 2.10.0 - await setTcloudVersion(tempDir, '2.10.0') - - // Set up Python environment with mock tcloud and sqlmesh - const pythonPath = await setupPythonEnvironment(pythonEnvDir) - - // Mark sqlmesh as installed - const binDir = path.dirname(pythonPath) - const installStateFile = path.join(binDir, '.sqlmesh_installed') - await fs.writeFile(installStateFile, '') - - // Configure VS Code settings to use our Python environment - const settings = { - 'python.defaultInterpreterPath': pythonPath, - 'sqlmesh.environmentPath': pythonEnvDir, - } - await fs.ensureDir(path.join(tempDir, '.vscode')) - await fs.writeJson( - path.join(tempDir, '.vscode', 'settings.json'), - settings, - { spaces: 2 }, - ) - - // Start VS Code - const { window, close } = await startVSCode(tempDir) - - // Open a SQL file to trigger SQLMesh activation - // Wait for the models folder to be visible - await window.waitForSelector('text=models') - - // Click on the models folder - await window - .getByRole('treeitem', { name: 'models', exact: true }) - .locator('a') - .click() - - // Open the top_waiters model - await window - .getByRole('treeitem', { name: 'customers.sql', exact: true }) - .locator('a') - .click() - - // Verify the context loads successfully - await window.waitForSelector('text=Loaded SQLMesh context') - - // Close VS Code - await close() - } finally { - // Clean up - await fs.remove(tempDir) + + // Start VS Code + const { window, close } = await startVSCode(tempDir) + + // Open a SQL file to trigger SQLMesh activation + // Wait for the models folder to be visible + await window.waitForSelector('text=models') + + // Click on the models folder + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + + // Open the top_waiters model + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + + await window.waitForSelector('text=Installing enterprise python package') + expect( + await window.locator('text=Installing enterprise python package'), + ).toHaveCount(2) + + await window.waitForSelector('text=Loaded SQLMesh context') + + // Close VS Code + await close() + } finally { + // Clean up + await fs.remove(tempDir) + } +}) + +test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in old version when ready', async ({}, testInfo) => { + testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation + const tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'vscode-test-tcloud-'), + ) + const pythonEnvDir = path.join(tempDir, '.venv') + + try { + // Copy sushi project + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + + // Create a tcloud.yaml to mark this as a tcloud project + const tcloudConfig = { + url: 'https://mock.tobikodata.com', + org: 'test-org', + project: 'test-project', } - }) + await fs.writeFile( + path.join(tempDir, 'tcloud.yaml'), + `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, + ) - test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in new version when ready', async ({}, testInfo) => { - testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation - const tempDir = await fs.mkdtemp( - path.join(os.tmpdir(), 'vscode-test-tcloud-'), + // Write mock ".tcloud_auth_state.json" file + await setupAuthenticatedState(tempDir) + + // Set tcloud version to 2.10.0 + await setTcloudVersion(tempDir, '2.10.0') + + // Set up Python environment with mock tcloud and sqlmesh + const pythonPath = await setupPythonEnvironment(pythonEnvDir) + + // Mark sqlmesh as installed + const binDir = path.dirname(pythonPath) + const installStateFile = path.join(binDir, '.sqlmesh_installed') + await fs.writeFile(installStateFile, '') + + // Configure VS Code settings to use our Python environment + const settings = { + 'python.defaultInterpreterPath': pythonPath, + 'sqlmesh.environmentPath': pythonEnvDir, + } + await fs.ensureDir(path.join(tempDir, '.vscode')) + await fs.writeJson( + path.join(tempDir, '.vscode', 'settings.json'), + settings, + { spaces: 2 }, ) - const pythonEnvDir = path.join(tempDir, '.venv') - - try { - // Copy sushi project - await fs.copy(SUSHI_SOURCE_PATH, tempDir) - - // Create a tcloud.yaml to mark this as a tcloud project - const tcloudConfig = { - url: 'https://mock.tobikodata.com', - org: 'test-org', - project: 'test-project', - } - await fs.writeFile( - path.join(tempDir, 'tcloud.yaml'), - `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, - ) - - // Write mock ".tcloud_auth_state.json" file - await setupAuthenticatedState(tempDir) - - // Set tcloud version to 2.10.0 - await setTcloudVersion(tempDir, '2.10.1') - - // Set up Python environment with mock tcloud and sqlmesh - const pythonPath = await setupPythonEnvironment(pythonEnvDir) - - // Mark sqlmesh as installed - const binDir = path.dirname(pythonPath) - const installStateFile = path.join(binDir, '.sqlmesh_installed') - await fs.writeFile(installStateFile, '') - - // Configure VS Code settings to use our Python environment - const settings = { - 'python.defaultInterpreterPath': pythonPath, - 'sqlmesh.environmentPath': pythonEnvDir, - } - await fs.ensureDir(path.join(tempDir, '.vscode')) - await fs.writeJson( - path.join(tempDir, '.vscode', 'settings.json'), - settings, - { spaces: 2 }, - ) - - // Start VS Code - const { window, close } = await startVSCode(tempDir) - - // Open a SQL file to trigger SQLMesh activation - // Wait for the models folder to be visible - await window.waitForSelector('text=models') - - // Click on the models folder - await window - .getByRole('treeitem', { name: 'models', exact: true }) - .locator('a') - .click() - - // Open the top_waiters model - await window - .getByRole('treeitem', { name: 'customers.sql', exact: true }) - .locator('a') - .click() - - // Verify the context loads successfully - await window.waitForSelector('text=Loaded SQLMesh context') - - // Close VS Code - await close() - } finally { - // Clean up - await fs.remove(tempDir) + + // Start VS Code + const { window, close } = await startVSCode(tempDir) + + // Open a SQL file to trigger SQLMesh activation + // Wait for the models folder to be visible + await window.waitForSelector('text=models') + + // Click on the models folder + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + + // Open the top_waiters model + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + + // Verify the context loads successfully + await window.waitForSelector('text=Loaded SQLMesh context') + + // Close VS Code + await close() + } finally { + // Clean up + await fs.remove(tempDir) + } +}) + +test('tcloud sqlmesh_lsp command starts the sqlmesh_lsp in new version when ready', async ({}, testInfo) => { + testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation + const tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'vscode-test-tcloud-'), + ) + const pythonEnvDir = path.join(tempDir, '.venv') + + try { + // Copy sushi project + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + + // Create a tcloud.yaml to mark this as a tcloud project + const tcloudConfig = { + url: 'https://mock.tobikodata.com', + org: 'test-org', + project: 'test-project', + } + await fs.writeFile( + path.join(tempDir, 'tcloud.yaml'), + `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, + ) + + // Write mock ".tcloud_auth_state.json" file + await setupAuthenticatedState(tempDir) + + // Set tcloud version to 2.10.0 + await setTcloudVersion(tempDir, '2.10.1') + + // Set up Python environment with mock tcloud and sqlmesh + const pythonPath = await setupPythonEnvironment(pythonEnvDir) + + // Mark sqlmesh as installed + const binDir = path.dirname(pythonPath) + const installStateFile = path.join(binDir, '.sqlmesh_installed') + await fs.writeFile(installStateFile, '') + + // Configure VS Code settings to use our Python environment + const settings = { + 'python.defaultInterpreterPath': pythonPath, + 'sqlmesh.environmentPath': pythonEnvDir, } + await fs.ensureDir(path.join(tempDir, '.vscode')) + await fs.writeJson( + path.join(tempDir, '.vscode', 'settings.json'), + settings, + { spaces: 2 }, + ) + + // Start VS Code + const { window, close } = await startVSCode(tempDir) + + // Open a SQL file to trigger SQLMesh activation + // Wait for the models folder to be visible + await window.waitForSelector('text=models') + + // Click on the models folder + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + + // Open the top_waiters model + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + + // Verify the context loads successfully + await window.waitForSelector('text=Loaded SQLMesh context') + + // Close VS Code + await close() + } finally { + // Clean up + await fs.remove(tempDir) + } +}) + +// This test is skipped becuase of the way the sign in window is shown is not useable by playwright. It's not solvable +// but the test is still useful when running it manually. +test.skip('tcloud not signed in and not installed, shows sign in window and then fact that loaded', async ({}, testInfo) => { + testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation + const tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'vscode-test-tcloud-'), + ) + const pythonEnvDir = path.join(tempDir, '.venv') + + // Create a tcloud.yaml to mark this as a tcloud project + const tcloudConfig = { + url: 'https://mock.tobikodata.com', + org: 'test-org', + project: 'test-project', + } + await fs.writeFile( + path.join(tempDir, 'tcloud.yaml'), + `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, + ) + + // Set up Python environment with mock tcloud and sqlmesh + const pythonPath = await setupPythonEnvironment(pythonEnvDir) + + // Configure VS Code settings to use our Python environment + const settings = { + 'python.defaultInterpreterPath': pythonPath, + 'sqlmesh.environmentPath': pythonEnvDir, + } + await fs.ensureDir(path.join(tempDir, '.vscode')) + await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, { + spaces: 2, }) + + // Set tcloud version to 2.10.0 + await setTcloudVersion(tempDir, '2.10.1') + + // Start VS Code + const { window, close } = await startVSCode(tempDir) + + try { + // Copy sushi project + await fs.copy(SUSHI_SOURCE_PATH, tempDir) + + // Open a SQL file to trigger SQLMesh activation + // Wait for the models folder to be visible + await window.waitForSelector('text=models') + + // Click on the models folder + await window + .getByRole('treeitem', { name: 'models', exact: true }) + .locator('a') + .click() + + // Open the top_waiters model + await window + .getByRole('treeitem', { name: 'customers.sql', exact: true }) + .locator('a') + .click() + + // Verify the sign in window is shown + await window.waitForSelector( + 'text=Please sign in to Tobiko Cloud to use SQLMesh', + ) + + // Click on the sign in button + await window + .getByRole('button', { name: 'Sign in' }) + .filter({ hasText: 'Sign in' }) + .click() + await window.waitForSelector('text="Signed in successfully"') + + await window.waitForSelector('text=Installing enterprise python package') + + await window.waitForSelector('text=Loaded SQLMesh context') + } finally { + // Clean up + await close() + await fs.remove(tempDir) + } })