Make OnAutoInsert more reliable and add integration tests

dibarbet · dibarbet · commit 75e63a192e3c · 2024-06-12T19:43:36.000-07:00
diff --git a/src/lsptoolshost/onAutoInsert.ts b/src/lsptoolshost/onAutoInsert.ts
@@ -4,27 +4,45 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
-import { UriConverter } from './uriConverter';
 
-import { FormattingOptions, TextDocumentIdentifier } from 'vscode-languageclient/node';
+import { FormattingOptions, LanguageClient, TextDocumentIdentifier } from 'vscode-languageclient/node';
 import * as RoslynProtocol from './roslynProtocol';
 import { RoslynLanguageServer } from './roslynLanguageServer';
 
-export function registerOnAutoInsert(languageServer: RoslynLanguageServer) {
+export function registerOnAutoInsert(languageServer: RoslynLanguageServer, languageClient: LanguageClient) {
     let source = new vscode.CancellationTokenSource();
-    vscode.workspace.onDidChangeTextDocument(async (e) => {
+
+    // We explicitly register against the server's didChange notification (instead of the VSCode workspace API)
+    // as we want to ensure that the server has processed the change before we request edits from auto insert.
+    // The VSCode workspace API will sometimes call this before the LSP client can send the change, leading to auto insert not working.
+    languageClient.getFeature('textDocument/didChange').onNotificationSent(async (event) => {
+        const e = event.params;
         if (e.contentChanges.length > 1 || e.contentChanges.length === 0) {
             return;
         }
 
         const change = e.contentChanges[0];
+        // TextDocumentContentChangeEvent is a union type that does not return a range if the change event is for a full document.
+        // Full document changes are not supported for onautoinsert.
+        if (!('range' in change)) {
+            return;
+        }
+
+        // Convert to a VSCode range for ease of handling.
+        const vscodeRange = languageClient.protocol2CodeConverter.asRange(change.range);
 
-        if (!change.range.isEmpty) {
+        // Empty or multiline changes are not supported for onautoinsert.
+        if (!vscodeRange.isEmpty || !vscodeRange.isSingleLine) {
             return;
         }
 
+        // We need to convert to a vscode TextDocument to apply the correct capabilities.
+        const uri = languageClient.protocol2CodeConverter.asUri(e.textDocument.uri);
+        // This is a no-op because the document is already open (in order to be edited).
+        const document = await vscode.workspace.openTextDocument(uri);
+
         const onAutoInsertFeature = languageServer.getOnAutoInsertFeature();
-        const onAutoInsertOptions = onAutoInsertFeature?.getOptions(e.document);
+        const onAutoInsertOptions = onAutoInsertFeature?.getOptions(document);
         const vsTriggerCharacters = onAutoInsertOptions?._vs_triggerCharacters;
 
         if (vsTriggerCharacters === undefined) {
@@ -38,10 +56,17 @@ export function registerOnAutoInsert(languageServer: RoslynLanguageServer) {
             return;
         }
 
+        // We have a single line range so we can compute the length by comparing the start and end character positions.
+        const rangeLength = vscodeRange.end.character - vscodeRange.start.character;
+
+        // The server expects the request position to represent the caret position in the text after the change has already been applied.
+        // We need to calculate what that position would be after the change is applied and send that to the server.
+        const position = vscodeRange.start.translate(0, change.text.length - rangeLength);
+
         source.cancel();
         source = new vscode.CancellationTokenSource();
         try {
-            await applyAutoInsertEdit(e, changeTrimmed, languageServer, source.token);
+            await applyAutoInsertEdit(position, changeTrimmed, e.textDocument, uri, languageServer, source.token);
         } catch (e) {
             if (e instanceof vscode.CancellationError) {
                 return;
@@ -53,25 +78,18 @@ export function registerOnAutoInsert(languageServer: RoslynLanguageServer) {
 }
 
 async function applyAutoInsertEdit(
-    e: vscode.TextDocumentChangeEvent,
-    changeTrimmed: string,
+    position: vscode.Position,
+    changeTextTrimmed: string,
+    textDocumentIdentifier: TextDocumentIdentifier,
+    uri: vscode.Uri,
     languageServer: RoslynLanguageServer,
     token: vscode.CancellationToken
 ) {
-    const change = e.contentChanges[0];
-    // The server expects the request position to represent the caret position in the text after the change has already been applied.
-    // We need to calculate what that position would be after the change is applied and send that to the server.
-    const position = new vscode.Position(
-        change.range.start.line,
-        change.range.start.character + (change.text.length - change.rangeLength)
-    );
-    const uri = UriConverter.serialize(e.document.uri);
-    const textDocument = TextDocumentIdentifier.create(uri);
     const formattingOptions = getFormattingOptions();
     const request: RoslynProtocol.OnAutoInsertParams = {
-        _vs_textDocument: textDocument,
+        _vs_textDocument: textDocumentIdentifier,
         _vs_position: position,
-        _vs_ch: changeTrimmed,
+        _vs_ch: changeTextTrimmed,
         _vs_options: formattingOptions,
     };
 
@@ -84,7 +102,7 @@ async function applyAutoInsertEdit(
         const code: any = vscode;
         const textEdits = [new code.SnippetTextEdit(new vscode.Range(startPosition, endPosition), docComment)];
         const edit = new vscode.WorkspaceEdit();
-        edit.set(e.document.uri, textEdits);
+        edit.set(uri, textEdits);
 
         const applied = vscode.workspace.applyEdit(edit);
         if (!applied) {
diff --git a/src/lsptoolshost/roslynLanguageServer.ts b/src/lsptoolshost/roslynLanguageServer.ts
@@ -131,6 +131,8 @@ export class RoslynLanguageServer {
 
         registerShowToastNotification(this._languageClient);
 
+        registerOnAutoInsert(this, this._languageClient);
+
         this._onAutoInsertFeature = new OnAutoInsertFeature(this._languageClient);
     }
 
@@ -990,8 +992,6 @@ export async function activateRoslynLanguageServer(
 
     registerRestoreCommands(context, languageServer, dotnetChannel);
 
-    registerOnAutoInsert(languageServer);
-
     context.subscriptions.push(registerLanguageServerOptionChanges(optionObservable));
 
     return languageServer;
diff --git a/test/integrationTests/diagnosticsHelpers.ts b/test/integrationTests/diagnosticsHelpers.ts
@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import { expect } from '@jest/globals';
-import { restartLanguageServer } from './integrationHelpers';
+import { restartLanguageServer, waitForExpectedResult } from './integrationHelpers';
 
 function sortDiagnostics(diagnostics: vscode.Diagnostic[]): vscode.Diagnostic[] {
     return diagnostics.sort((a, b) => {
@@ -21,32 +21,26 @@ export async function waitForExpectedDiagnostics(
     assertExpectedDiagnostics: (input: [vscode.Uri, vscode.Diagnostic[]][]) => void,
     file?: vscode.Uri
 ): Promise<void> {
-    let duration = 30 * 1000;
+    const duration = 30 * 1000;
     const step = 500;
-    let error: any = undefined;
-    while (duration > 0) {
-        const diagnostics: [vscode.Uri, vscode.Diagnostic[]][] = [];
-        if (file) {
-            diagnostics.push([file, sortDiagnostics(vscode.languages.getDiagnostics(file))]);
-        } else {
-            const allDiagnostics = vscode.languages.getDiagnostics();
-            for (const [uri, uriDiagnostics] of allDiagnostics) {
-                diagnostics.push([uri, sortDiagnostics(uriDiagnostics)]);
+    await waitForExpectedResult<[vscode.Uri, vscode.Diagnostic[]][]>(
+        () => {
+            const diagnostics: [vscode.Uri, vscode.Diagnostic[]][] = [];
+            if (file) {
+                diagnostics.push([file, sortDiagnostics(vscode.languages.getDiagnostics(file))]);
+            } else {
+                const allDiagnostics = vscode.languages.getDiagnostics();
+                for (const [uri, uriDiagnostics] of allDiagnostics) {
+                    diagnostics.push([uri, sortDiagnostics(uriDiagnostics)]);
+                }
             }
-        }
-
-        try {
-            assertExpectedDiagnostics(diagnostics);
-            return;
-        } catch (e) {
-            error = e;
-            // Wait for a bit and try again.
-            await new Promise((r) => setTimeout(r, step));
-            duration -= step;
-        }
-    }
 
-    throw new Error(`Polling did not succeed within the alotted duration: ${error}`);
+            return diagnostics;
+        },
+        duration,
+        step,
+        assertExpectedDiagnostics
+    );
 }
 
 export async function setBackgroundAnalysisScopes(scopes: { compiler: string; analyzer: string }): Promise<void> {
diff --git a/test/integrationTests/integrationHelpers.ts b/test/integrationTests/integrationHelpers.ts
@@ -58,6 +58,14 @@ export async function openFileInWorkspaceAsync(relativeFilePath: string): Promis
     return uri;
 }
 
+/**
+ * Reverts any unsaved changes to the active file.
+ * Useful to reset state between tests without fully reloading everything.
+ */
+export async function revertActiveFile(): Promise<void> {
+    await vscode.commands.executeCommand('workbench.action.files.revert');
+}
+
 export async function restartLanguageServer(): Promise<void> {
     const csharpExtension = vscode.extensions.getExtension<CSharpExtensionExports>('ms-dotnettools.csharp');
     // Register to wait for initialization events and restart the server.
@@ -112,3 +120,33 @@ function isGivenSln(workspace: typeof vscode.workspace, expectedProjectFileName:
 
     return projectFileName === expectedProjectFileName;
 }
+
+export async function waitForExpectedResult<T>(
+    getValue: () => Promise<T> | T,
+    duration: number,
+    step: number,
+    expression: (input: T) => void
+): Promise<void> {
+    let value: T;
+    let error: any = undefined;
+
+    while (duration > 0) {
+        value = await getValue();
+
+        try {
+            expression(value);
+            return;
+        } catch (e) {
+            error = e;
+            // Wait for a bit and try again.
+            await new Promise((r) => setTimeout(r, step));
+            duration -= step;
+        }
+    }
+
+    throw new Error(`Polling did not succeed within the alotted duration: ${error}`);
+}
+
+export async function sleep(ms = 0) {
+    return new Promise((r) => setTimeout(r, ms));
+}
diff --git a/test/integrationTests/onAutoInsert.integration.test.ts b/test/integrationTests/onAutoInsert.integration.test.ts
@@ -0,0 +1,77 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as path from 'path';
+import testAssetWorkspace from './testAssets/testAssetWorkspace';
+import {
+    activateCSharpExtension,
+    openFileInWorkspaceAsync,
+    revertActiveFile,
+    sleep,
+    waitForExpectedResult,
+} from './integrationHelpers';
+import { describe, beforeAll, beforeEach, afterAll, test, expect, afterEach } from '@jest/globals';
+
+describe(`[${testAssetWorkspace.description}] Test OnAutoInsert`, function () {
+    beforeAll(async function () {
+        await activateCSharpExtension();
+    });
+
+    beforeEach(async function () {
+        await openFileInWorkspaceAsync(path.join('src', 'app', 'DocComments.cs'));
+    });
+
+    afterEach(async () => {
+        await revertActiveFile();
+    });
+
+    afterAll(async () => {
+        await testAssetWorkspace.cleanupWorkspace();
+    });
+
+    test('Triple slash inserts doc comment snippet', async () => {
+        await sleep(1);
+        await vscode.window.activeTextEditor!.edit((editBuilder) => {
+            editBuilder.insert(new vscode.Position(2, 6), '/');
+        });
+
+        // OnAutoInsert is triggered by the change event but completes asynchronously, so wait for the buffer to be updated.
+
+        await waitForExpectedResult<string | undefined>(
+            async () => vscode.window.activeTextEditor?.document.getText(),
+            10000,
+            100,
+            (input) => {
+                expect(normalizeNewlines(input)).toContain(
+                    '/// <summary>\n    /// \n    /// </summary>\n    /// <param name="param1"></param>\n    /// <param name="param2"></param>\n    /// <returns></returns>'
+                );
+            }
+        );
+    });
+
+    test('Enter in comment inserts triple-slashes preceding', async () => {
+        await vscode.window.activeTextEditor!.edit((editBuilder) => {
+            editBuilder.insert(new vscode.Position(8, 17), '\n');
+        });
+
+        // OnAutoInsert is triggered by the change event but completes asynchronously, so wait for the buffer to be updated.
+
+        await waitForExpectedResult<string | undefined>(
+            async () => vscode.window.activeTextEditor?.document.getText(),
+            10000,
+            100,
+            (input) => {
+                expect(normalizeNewlines(input)).toContain(
+                    '/// <summary>\n    /// \n\n    /// </summary>\n    void M2() {}'
+                );
+            }
+        );
+    });
+});
+
+function normalizeNewlines(text: string | undefined): string | undefined {
+    return text?.replaceAll('\r\n', '\n');
+}
diff --git a/test/integrationTests/testAssets/slnWithCsproj/src/app/DocComments.cs b/test/integrationTests/testAssets/slnWithCsproj/src/app/DocComments.cs
@@ -1,6 +1,6 @@
 class DocComments
 {
-    ///
+    //
     string M(int param1, string param2)
     {
         return null;

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`class DocComments`
`2`	`2`	`{`
`3`		`- ///`
	`3`	`+ //`
`4`	`4`	`string M(int param1, string param2)`
`5`	`5`	`{`
`6`	`6`	`return null;`