diff --git a/src/common/localize.ts b/src/common/localize.ts index e6f8db55..b528b7fc 100644 --- a/src/common/localize.ts +++ b/src/common/localize.ts @@ -26,6 +26,12 @@ export namespace WorkbenchStrings { export namespace InlineScriptStrings { export const updateExtension = l10n.t('Update Extension'); + export const setUpScriptEnvironment = l10n.t("Set up this script's Python environment"); + + export const saveFailedBeforeSetup = l10n.t( + 'Could not save this script, so its environment was not set up. Save the file and try again.', + ); + export const updatePythonExtension = l10n.t( 'The environment for this script was created. Update the Python extension for the full inline script experience.', ); diff --git a/src/features/inlineScript/setupCodeAction.ts b/src/features/inlineScript/setupCodeAction.ts new file mode 100644 index 00000000..37275ffa --- /dev/null +++ b/src/features/inlineScript/setupCodeAction.ts @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { + CancellationToken, + CodeAction, + CodeActionContext, + CodeActionKind, + CodeActionProvider, + Diagnostic, + Disposable, + languages, + Range, + TextDocument, +} from 'vscode'; +import { readInlineScriptMetadata, sliceHeaderBytes } from '../../common/inlineScript/metadata'; +import { getInlineScriptRoutingKey, InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry'; +import { InlineScriptStrings } from '../../common/localize'; +import { isInlineScriptsFeatureEnabled } from '../../helpers'; + +/** + * Diagnostic codes meaning "this import did not resolve", lowercased for comparison. + * + * `reportMissingModuleSource` is included deliberately, unlike in Pylance's own + * `isMissingImportDiagnostic`: a stub without source means the package is not installed, which + * setting the script's environment up fixes. + */ +const UNRESOLVED_IMPORT_DIAGNOSTIC_CODES: ReadonlySet = new Set([ + // Pyright / Pylance / basedpyright. + 'reportmissingimports', + 'reportmissingmodulesource', + // Ty. + 'unresolved-import', + 'possibly-missing-import', + // Pyrefly. + 'missing-import', + 'missing-source', + 'missing-source-for-stubs', + // mypy, via ms-python.mypy-type-checker. + 'import-not-found', + 'import-untyped', +]); + +function normalizeDiagnosticCode(code: Diagnostic['code']): string | undefined { + if (code === undefined || code === null) { + return undefined; + } + const value = typeof code === 'object' ? code.value : code; + return typeof value === 'string' || typeof value === 'number' ? String(value).toLowerCase() : undefined; +} + +/** + * Whether `diagnostic` reports an import that could not be resolved. Matches on `code`, never on + * `source`: Pyrefly-backed Pylance reports its source as the literal string `pylance + pyrefly`. + */ +export function isUnresolvedImportDiagnostic(diagnostic: Diagnostic): boolean { + const code = normalizeDiagnosticCode(diagnostic.code); + return code !== undefined && UNRESOLVED_IMPORT_DIAGNOSTIC_CODES.has(code); +} + +/** + * Offers "Set up this script's Python environment" as a quick fix on an unresolved import in a `.py` + * file that declares a PEP 723 `# /// script` block and has no inline-script environment yet. + * + * Complements the CodeLens, which is hidden while the document is dirty — the moment a user has just + * typed the import that does not resolve. This provider parses the in-memory buffer instead. + * + * `diagnostics` and `isPreferred` are both left unset: setup installs the block's declared + * dependencies verbatim and may not resolve the import at all, so the action must not claim to fix + * the diagnostic or pre-empt a real import fix. + */ +export class InlineScriptSetupCodeActionProvider implements CodeActionProvider { + constructor( + private readonly routing: InlineScriptRoutingRegistry, + private readonly setupCommand: string, + ) {} + + /** Gates run cheapest-first, and before any parsing: VS Code may call this on every cursor move. */ + public provideCodeActions( + document: TextDocument, + _range: Range, + context: CodeActionContext, + _token: CancellationToken, + ): CodeAction[] { + if (!isInlineScriptsFeatureEnabled()) { + return []; + } + if (!context.diagnostics.some(isUnresolvedImportDiagnostic)) { + return []; + } + const uri = document.uri; + if (!getInlineScriptRoutingKey(uri)) { + return []; + } + if (this.routing.shouldRoute(uri)) { + return []; + } + if (!readInlineScriptMetadata(sliceHeaderBytes(document.getText()), uri.fsPath)) { + return []; + } + const action = new CodeAction(InlineScriptStrings.setUpScriptEnvironment, CodeActionKind.QuickFix); + action.command = { + title: InlineScriptStrings.setUpScriptEnvironment, + command: this.setupCommand, + arguments: [uri], + }; + return [action]; + } +} + +/** Register the inline-script quick fix for local `.py` files. */ +export function registerInlineScriptSetupCodeAction( + routing: InlineScriptRoutingRegistry, + setupCommand: string, +): Disposable { + return languages.registerCodeActionsProvider( + { scheme: 'file', language: 'python' }, + new InlineScriptSetupCodeActionProvider(routing, setupCommand), + { providedCodeActionKinds: [CodeActionKind.QuickFix] }, + ); +} diff --git a/src/features/inlineScript/setupEnvironment.ts b/src/features/inlineScript/setupEnvironment.ts index f604fdde..488893fc 100644 --- a/src/features/inlineScript/setupEnvironment.ts +++ b/src/features/inlineScript/setupEnvironment.ts @@ -1,12 +1,13 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import { commands, Disposable, l10n, QuickPickItem, Uri, window } from 'vscode'; +import { commands, Disposable, l10n, QuickPickItem, TextDocument, Uri, window } from 'vscode'; import { PythonEnvironment } from '../../api'; import { INLINE_SCRIPT_MANAGER_ID } from '../../common/constants'; import { readInlineScriptMetadataFromFile } from '../../common/inlineScript/metadata'; import { InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry'; -import { traceError, traceInfo } from '../../common/logging'; +import { InlineScriptStrings } from '../../common/localize'; +import { traceError, traceInfo, traceVerbose } from '../../common/logging'; import { normalizePath } from '../../common/utils/pathUtils'; import { showErrorMessage, @@ -18,6 +19,7 @@ import { asRelativePath, findFiles, getOpenTextDocuments } from '../../common/wo import { EnvironmentManagers } from '../../internal.api'; import { registerInlineScriptCodeLens } from './codeLens'; import { promptUpdateExtensionsForInlineScripts } from './extensionVersionCheck'; +import { registerInlineScriptSetupCodeAction } from './setupCodeAction'; /** * Hidden command invoked by the inline-script CodeLens to set up the environment for one script. @@ -86,20 +88,50 @@ async function seedRoutingMetadataForClosedScript(scriptUri: Uri, routing: Inlin if (routing.getMetadata(scriptUri)) { return; } + if (findOpenDocument(scriptUri)) { + return; + } + const metadata = await readInlineScriptMetadataFromFile(scriptUri); + if (metadata && !routing.getMetadata(scriptUri)) { + routing.setMetadata(scriptUri, metadata); + } +} + +/** The open text document backing `scriptUri`, if the user has it open. */ +function findOpenDocument(scriptUri: Uri): TextDocument | undefined { const scriptPath = normalizePath(scriptUri.fsPath); - const isOpen = getOpenTextDocuments().some( + return getOpenTextDocuments().find( (document) => document.uri.scheme === 'file' && normalizePath(document.uri.fsPath) === scriptPath, ); - if (isOpen) { - return; +} + +/** + * Save `scriptUri` if it is open with unsaved changes, so setup reads what the user actually sees. + * + * Returns `false` when the document could not be saved; setup must not run in that case. + */ +async function saveScriptBeforeSetup(scriptUri: Uri, routing: InlineScriptRoutingRegistry): Promise { + const document = findOpenDocument(scriptUri); + if (!document?.isDirty) { + return true; + } + if (!(await document.save())) { + traceError(`Could not save ${scriptUri.fsPath} before setting up its inline-script environment.`); + return false; } + traceVerbose(`Saved ${scriptUri.fsPath} before setting up its inline-script environment.`); + // Seeding here is load-bearing: without it a just-typed block goes from no metadata to an + // identity while `create` runs, which `setUpInlineScriptEnvironment` reads as a concurrent edit + // and silently skips the association. const metadata = await readInlineScriptMetadataFromFile(scriptUri); - if (metadata && !routing.getMetadata(scriptUri)) { + if (metadata) { routing.setMetadata(scriptUri, metadata); } + return true; } -function setupInlineScriptEnvironmentHandler( +/** Handler for the single-file setup command, shared by the CodeLens and the quick fix. */ +export function setupInlineScriptEnvironmentHandler( em: EnvironmentManagers, routing: InlineScriptRoutingRegistry, ): (scriptUri?: Uri) => Promise { @@ -112,13 +144,13 @@ function setupInlineScriptEnvironmentHandler( showErrorMessage(l10n.t('The inline script environment manager is not available yet. Try again shortly.')); return; } + if (!(await saveScriptBeforeSetup(uri, routing))) { + showErrorMessage(InlineScriptStrings.saveFailedBeforeSetup); + return; + } + let environment: PythonEnvironment | undefined; try { - const environment = await setUpInlineScriptEnvironment(uri, em, routing); - if (!environment) { - notifyInlineScriptSetupOutcome(uri, routing); - return; - } - await promptUpdateExtensionsForInlineScripts(); + environment = await setUpInlineScriptEnvironment(uri, em, routing); } catch (error) { traceError(`Failed to set up the inline-script environment for ${uri.fsPath}:`, error); showErrorMessage( @@ -126,7 +158,17 @@ function setupInlineScriptEnvironmentHandler( 'Failed to set up the environment for this script. See the Python Environments output for details.', ), ); + return; + } + if (!environment) { + notifyInlineScriptSetupOutcome(uri, routing); + return; } + // Kept out of the try: the environment is already set up, so a failure in this follow-up + // must not be reported to the user as a setup failure. + await promptUpdateExtensionsForInlineScripts().catch((error) => + traceError('Failed to check companion extension versions for inline scripts:', error), + ); }; } @@ -314,13 +356,15 @@ async function filterInlineScriptFiles(files: readonly Uri[]): Promise { } /** - * Register the inline-script user-facing surfaces (the CodeLens and its setup commands). Only called - * when the PEP 723 inline-script feature flag is enabled. The single-file setup command is invoked by - * the CodeLens and stays out of `package.json`; the bulk command is palette-gated behind the flag. + * Register the inline-script user-facing surfaces (the CodeLens, the quick fix, and their setup + * commands). Only called when the PEP 723 inline-script feature flag is enabled. The single-file + * setup command is invoked by both surfaces and stays out of `package.json`; the bulk command is + * palette-gated behind the flag. */ export function registerInlineScriptUx(em: EnvironmentManagers, routing: InlineScriptRoutingRegistry): Disposable[] { return [ registerInlineScriptCodeLens(routing, SETUP_INLINE_SCRIPT_ENV_COMMAND), + registerInlineScriptSetupCodeAction(routing, SETUP_INLINE_SCRIPT_ENV_COMMAND), commands.registerCommand(SETUP_INLINE_SCRIPT_ENV_COMMAND, setupInlineScriptEnvironmentHandler(em, routing)), commands.registerCommand(SETUP_INLINE_SCRIPT_ENVS_COMMAND, () => setUpInlineScriptEnvironmentsInWorkspace(em, routing), diff --git a/src/test/features/inlineScript/setupCodeAction.unit.test.ts b/src/test/features/inlineScript/setupCodeAction.unit.test.ts new file mode 100644 index 00000000..37d2a129 --- /dev/null +++ b/src/test/features/inlineScript/setupCodeAction.unit.test.ts @@ -0,0 +1,253 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import assert from 'assert'; +import * as sinon from 'sinon'; +import { CodeActionContext, Diagnostic, Position, Range, TextDocument, Uri } from 'vscode'; +import { InlineScriptMetadata, MAX_HEADER_BYTES } from '../../../common/inlineScript/metadata'; +import { InlineScriptRoutingRegistry } from '../../../common/inlineScript/routingRegistry'; +import { InlineScriptStrings } from '../../../common/localize'; +import { + InlineScriptSetupCodeActionProvider, + isUnresolvedImportDiagnostic, +} from '../../../features/inlineScript/setupCodeAction'; +import * as helpers from '../../../helpers'; + +const SETUP_COMMAND = 'python-envs.setupInlineScriptEnv'; + +const SCRIPT_WITH_METADATA = [ + '# /// script', + '# dependencies = ["requests"]', + '# ///', + '', + 'import requests', + '', +].join('\n'); + +const SCRIPT_WITHOUT_METADATA = ['import requests', '', 'print(requests)', ''].join('\n'); + +function makeMetadata(): InlineScriptMetadata { + return { + dependencies: ['requests'], + range: { start: 0, end: 47 }, + sourceRange: { start: 0, end: 47 }, + }; +} + +function makeDocument(uri: Uri, text = SCRIPT_WITH_METADATA): TextDocument { + return { + uri, + languageId: 'python', + getText: () => text, + } as unknown as TextDocument; +} + +function makeDiagnostic(code: Diagnostic['code'], source = 'Pylance'): Diagnostic { + const range = new Range(new Position(4, 7), new Position(4, 15)); + const diagnostic = new Diagnostic(range, 'Import "requests" could not be resolved'); + diagnostic.code = code; + diagnostic.source = source; + return diagnostic; +} + +function makeContext(diagnostics: Diagnostic[]): CodeActionContext { + return { diagnostics, only: undefined, triggerKind: 1 } as unknown as CodeActionContext; +} + +const INVOCATION_RANGE = new Range(new Position(4, 7), new Position(4, 7)); + +suite('Inline script setup code action', () => { + const scriptUri = Uri.file('/workspace/app.py'); + let routing: InlineScriptRoutingRegistry; + let provider: InlineScriptSetupCodeActionProvider; + let featureEnabledStub: sinon.SinonStub; + + setup(() => { + featureEnabledStub = sinon.stub(helpers, 'isInlineScriptsFeatureEnabled').returns(true); + routing = new InlineScriptRoutingRegistry(); + provider = new InlineScriptSetupCodeActionProvider(routing, SETUP_COMMAND); + }); + + teardown(() => { + routing.dispose(); + sinon.restore(); + }); + + function provide(document: TextDocument, diagnostics: Diagnostic[]) { + return provider.provideCodeActions(document, INVOCATION_RANGE, makeContext(diagnostics), {} as never); + } + + suite('isUnresolvedImportDiagnostic', () => { + const matching: Array<[string, Diagnostic['code']]> = [ + ['Pyright/Pylance reportMissingImports', 'reportMissingImports'], + ['Pyright/Pylance reportMissingModuleSource', 'reportMissingModuleSource'], + ['Ty unresolved-import', 'unresolved-import'], + ['Ty possibly-missing-import', 'possibly-missing-import'], + ['Pyrefly missing-import', 'missing-import'], + ['Pyrefly missing-source', 'missing-source'], + ['Pyrefly missing-source-for-stubs', 'missing-source-for-stubs'], + ['mypy import-not-found', 'import-not-found'], + ['mypy import-untyped', 'import-untyped'], + ]; + + for (const [label, code] of matching) { + test(`matches ${label}`, () => { + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic(code)), true); + }); + } + + test('matches the {value, target} object form of Diagnostic.code', () => { + const code = { value: 'reportMissingImports', target: Uri.parse('https://example.invalid/rule') }; + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic(code)), true); + }); + + test('matches regardless of case', () => { + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic('ReportMissingImports')), true); + }); + + test('ignores the diagnostic source, which is "pylance + pyrefly" for Pyrefly', () => { + assert.strictEqual( + isUnresolvedImportDiagnostic(makeDiagnostic('missing-import', 'pylance + pyrefly')), + true, + ); + }); + + test('does not match unrelated diagnostic codes', () => { + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic('reportUndefinedVariable')), false); + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic('unresolved-reference')), false); + }); + + test('does not match a diagnostic with no code', () => { + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic(undefined)), false); + }); + + test('does not match a numeric code that is not an import rule', () => { + assert.strictEqual(isUnresolvedImportDiagnostic(makeDiagnostic(42)), false); + }); + }); + + suite('provideCodeActions', () => { + test('offers setup for an unresolved import in a not-yet-configured inline script', () => { + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 1); + assert.strictEqual(actions[0].title, InlineScriptStrings.setUpScriptEnvironment); + assert.strictEqual(actions[0].command?.command, SETUP_COMMAND); + assert.deepStrictEqual(actions[0].command?.arguments, [scriptUri]); + }); + + test('leaves diagnostics unset so VS Code is not told the action resolves them', () => { + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 1); + assert.strictEqual( + actions[0].diagnostics, + undefined, + 'setting diagnostics would claim the action fixes them and opt it into fix-all', + ); + }); + + test('leaves isPreferred unset so it never pre-empts a real import fix', () => { + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 1); + assert.strictEqual(actions[0].isPreferred, undefined); + }); + + test('offers nothing when no diagnostic reports an unresolved import', () => { + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportUndefinedVariable')]); + + assert.strictEqual(actions.length, 0); + }); + + test('offers nothing when there are no diagnostics at the invocation range', () => { + const actions = provide(makeDocument(scriptUri), []); + + assert.strictEqual(actions.length, 0); + }); + + test('offers nothing when the file has no PEP 723 block', () => { + const actions = provide(makeDocument(scriptUri, SCRIPT_WITHOUT_METADATA), [ + makeDiagnostic('reportMissingImports'), + ]); + + assert.strictEqual(actions.length, 0); + }); + + test('offers nothing when the PEP 723 block is malformed', () => { + const malformed = ['# /// script', '# dependencies = [', '# ///', '', 'import requests', ''].join('\n'); + + const actions = provide(makeDocument(scriptUri, malformed), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 0); + }); + + test('offers nothing once the script is already set up', () => { + routing.setMetadata(scriptUri, makeMetadata()); + routing.setValidatedAssociation(scriptUri, true); + assert.strictEqual(routing.shouldRoute(scriptUri), true); + + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 0); + }); + + test('offers nothing when the inline-scripts feature flag is off', () => { + featureEnabledStub.returns(false); + + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 0); + }); + + test('offers nothing for a file that cannot carry an inline-script environment', () => { + const notebookUri = Uri.parse('vscode-notebook-cell:/workspace/app.py#ch0'); + + const actions = provide(makeDocument(notebookUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 0); + }); + + test('offers setup again after the metadata block is edited to add the missing dependency', () => { + routing.setMetadata(scriptUri, makeMetadata()); + routing.setValidatedAssociation(scriptUri, true); + assert.strictEqual( + provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]).length, + 0, + 'precondition: hidden while the script is set up', + ); + + routing.setMetadata(scriptUri, { ...makeMetadata(), dependencies: ['requests', 'rich'] }); + assert.strictEqual(routing.shouldRoute(scriptUri), false); + + const edited = ['# /// script', '# dependencies = ["requests", "rich"]', '# ///', '', 'import rich', ''].join( + '\n', + ); + const actions = provide(makeDocument(scriptUri, edited), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 1); + assert.strictEqual(actions[0].command?.command, SETUP_COMMAND); + }); + + test('parses the in-memory buffer, so it works on an unsaved edit the CodeLens cannot see', () => { + assert.strictEqual(routing.getMetadata(scriptUri), undefined); + + const actions = provide(makeDocument(scriptUri), [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual(actions.length, 1); + }); + + test('ignores a metadata block that sits past the header byte budget setup reads', () => { + const padding = `${'# padding comment\n'.repeat(Math.ceil(MAX_HEADER_BYTES / 18) + 10)}`; + const document = makeDocument(scriptUri, padding + SCRIPT_WITH_METADATA); + + const actions = provide(document, [makeDiagnostic('reportMissingImports')]); + + assert.strictEqual( + actions.length, + 0, + 'setup reads only the first MAX_HEADER_BYTES from disk, so the block must be invisible here too', + ); + }); + }); +}); diff --git a/src/test/features/inlineScript/setupEnvironment.unit.test.ts b/src/test/features/inlineScript/setupEnvironment.unit.test.ts index 9f3ff1a8..6d184cbc 100644 --- a/src/test/features/inlineScript/setupEnvironment.unit.test.ts +++ b/src/test/features/inlineScript/setupEnvironment.unit.test.ts @@ -10,13 +10,16 @@ import { INLINE_SCRIPT_MANAGER_ID } from '../../../common/constants'; import { InlineScriptMetadata } from '../../../common/inlineScript/metadata'; import * as metadataApi from '../../../common/inlineScript/metadata'; import { InlineScriptRoutingRegistry } from '../../../common/inlineScript/routingRegistry'; +import { InlineScriptStrings } from '../../../common/localize'; import * as winapi from '../../../common/window.apis'; import * as wapi from '../../../common/workspace.apis'; import { notifyInlineScriptSetupOutcome, setUpInlineScriptEnvironment, setUpInlineScriptEnvironmentsInWorkspace, + setupInlineScriptEnvironmentHandler, } from '../../../features/inlineScript/setupEnvironment'; +import * as extensionVersionCheck from '../../../features/inlineScript/extensionVersionCheck'; import { EnvironmentManagers, InternalEnvironmentManager } from '../../../internal.api'; function makeEnv(): PythonEnvironment { @@ -313,3 +316,113 @@ suite('notifyInlineScriptSetupOutcome', () => { sinon.assert.notCalled(errorStub); }); }); + +suite('setupInlineScriptEnvironmentHandler', () => { + const scriptUri = Uri.file('/workspace/app.py'); + let em: typemoq.IMock; + let manager: typemoq.IMock; + let routing: InlineScriptRoutingRegistry; + let readMetadataStub: sinon.SinonStub; + let openDocumentsStub: sinon.SinonStub; + let errorStub: sinon.SinonStub; + let saveStub: sinon.SinonStub; + let promptStub: sinon.SinonStub; + + setup(() => { + em = typemoq.Mock.ofType(); + manager = typemoq.Mock.ofType(); + routing = new InlineScriptRoutingRegistry(); + em.setup((m) => m.getEnvironmentManager(INLINE_SCRIPT_MANAGER_ID)).returns(() => manager.object); + readMetadataStub = sinon.stub(metadataApi, 'readInlineScriptMetadataFromFile').resolves(undefined); + openDocumentsStub = sinon.stub(wapi, 'getOpenTextDocuments').returns([]); + errorStub = sinon.stub(winapi, 'showErrorMessage').resolves(undefined); + sinon.stub(winapi, 'showInformationMessage').resolves(undefined); + sinon.stub(winapi, 'showWarningMessage').resolves(undefined); + promptStub = sinon.stub(extensionVersionCheck, 'promptUpdateExtensionsForInlineScripts').resolves(); + saveStub = sinon.stub().resolves(true); + }); + + teardown(() => { + routing.dispose(); + sinon.restore(); + }); + + function openDirtyDocument(isDirty = true): void { + openDocumentsStub.returns([{ uri: scriptUri, isDirty, save: saveStub }]); + } + + function expectEnvironmentCreated(): PythonEnvironment { + const env = makeEnv(); + manager.setup((m) => m.create(scriptUri, undefined)).returns(() => Promise.resolve(env)); + em.setup((m) => m.setEnvironment(scriptUri, env)).returns(() => Promise.resolve()); + return env; + } + + test('saves a dirty document before setup, because setup reads the block from disk', async () => { + openDirtyDocument(); + expectEnvironmentCreated(); + + await setupInlineScriptEnvironmentHandler(em.object, routing)(scriptUri); + + sinon.assert.calledOnce(saveStub); + em.verify((m) => m.setEnvironment(scriptUri, typemoq.It.isAny()), typemoq.Times.once()); + }); + + test('seeds routing metadata from the saved file so setup does not misread it as a mid-setup edit', async () => { + openDirtyDocument(); + const metadata = makeMetadata(['requests']); + readMetadataStub.resolves(metadata); + const env = makeEnv(); + manager + .setup((m) => m.create(scriptUri, undefined)) + .returns(async () => { + routing.setMetadata(scriptUri, makeMetadata(['requests'])); + return env; + }); + em.setup((m) => m.setEnvironment(scriptUri, env)).returns(() => Promise.resolve()); + + await setupInlineScriptEnvironmentHandler(em.object, routing)(scriptUri); + + assert.deepStrictEqual(routing.getMetadata(scriptUri), metadata); + em.verify((m) => m.setEnvironment(scriptUri, env), typemoq.Times.once()); + }); + + test('does not save a document that has no unsaved changes', async () => { + openDirtyDocument(false); + expectEnvironmentCreated(); + + await setupInlineScriptEnvironmentHandler(em.object, routing)(scriptUri); + + sinon.assert.notCalled(saveStub); + }); + + test('does not run setup when the document could not be saved', async () => { + openDirtyDocument(); + saveStub.resolves(false); + const env = expectEnvironmentCreated(); + + await setupInlineScriptEnvironmentHandler(em.object, routing)(scriptUri); + + manager.verify((m) => m.create(scriptUri, undefined), typemoq.Times.never()); + em.verify((m) => m.setEnvironment(scriptUri, env), typemoq.Times.never()); + sinon.assert.calledOnce(errorStub); + assert.strictEqual(errorStub.firstCall.args[0], InlineScriptStrings.saveFailedBeforeSetup); + }); + + test('reports an error when setup throws', async () => { + manager.setup((m) => m.create(scriptUri, undefined)).returns(() => Promise.reject(new Error('boom'))); + + await setupInlineScriptEnvironmentHandler(em.object, routing)(scriptUri); + + sinon.assert.calledOnce(errorStub); + }); + + test('does not report a failing companion-extension prompt as a setup failure', async () => { + expectEnvironmentCreated(); + promptStub.rejects(new Error('boom')); + + await setupInlineScriptEnvironmentHandler(em.object, routing)(scriptUri); + + sinon.assert.notCalled(errorStub); + }); +});