From 70fe74a52b09a1fea0b8a43d1ee772a8382e083f Mon Sep 17 00:00:00 2001
From: David Barbet <dabarbet@microsoft.com>
Date: Thu, 12 Oct 2023 11:09:48 -0700
Subject: [PATCH] Wait for operations on the server to improve integration test
 startup

---
 test/integrationTests/integrationHelpers.ts   | 28 ++++++-------------
 ...sp_tools_host_slnWithCsproj.code-workspace |  2 +-
 test/integrationTests/testHooks.ts            | 25 +++++++++++++++++
 3 files changed, 34 insertions(+), 21 deletions(-)
 create mode 100644 test/integrationTests/testHooks.ts

diff --git a/test/integrationTests/integrationHelpers.ts b/test/integrationTests/integrationHelpers.ts
index d34c6eca4..a69d9d986 100644
--- a/test/integrationTests/integrationHelpers.ts
+++ b/test/integrationTests/integrationHelpers.ts
@@ -7,8 +7,8 @@ import * as vscode from 'vscode';
 import * as path from 'path';
 import { CSharpExtensionExports } from '../../src/csharpExtensionExports';
 import { existsSync } from 'fs';
-import { ServerStateChange } from '../../src/lsptoolshost/serverStateChange';
 import testAssetWorkspace from './testAssets/testAssetWorkspace';
+import { WaitForAsyncOperationsRequest } from './testHooks';
 
 export async function activateCSharpExtension(): Promise<void> {
     // Ensure the dependent extension exists - when launching via F5 launch.json we can't install the extension prior to opening vscode.
@@ -28,31 +28,14 @@ export async function activateCSharpExtension(): Promise<void> {
     // Run a restore manually to make sure the project is up to date since we don't have automatic restore.
     await testAssetWorkspace.restoreLspToolsHostAsync();
 
-    // If the extension is already active, we need to restart it to ensure we start with a clean server state.
-    // For example, a previous test may have changed configs, deleted restored packages or made other changes that would put it in an invalid state.
-    let shouldRestart = false;
-    if (csharpExtension.isActive) {
-        shouldRestart = true;
-    }
-
     // Explicitly await the extension activation even if completed so that we capture any errors it threw during activation.
     await csharpExtension.activate();
     await csharpExtension.exports.initializationFinished();
     console.log('ms-dotnettools.csharp activated');
     console.log(`Extension Log Directory: ${csharpExtension.exports.logDirectory}`);
 
-    if (shouldRestart) {
-        // Register to wait for initialization events and restart the server.
-        const waitForInitialProjectLoad = new Promise<void>((resolve, _) => {
-            csharpExtension.exports.experimental.languageServerEvents.onServerStateChange(async (state) => {
-                if (state === ServerStateChange.ProjectInitializationComplete) {
-                    resolve();
-                }
-            });
-        });
-        await vscode.commands.executeCommand('dotnet.restartServer');
-        await waitForInitialProjectLoad;
-    }
+    await waitForAllAsyncOperationsAsync(csharpExtension.exports);
+    console.log('Async operations completed');
 }
 
 export async function openFileInWorkspaceAsync(relativeFilePath: string): Promise<void> {
@@ -73,6 +56,11 @@ export function isSlnWithGenerator(workspace: typeof vscode.workspace) {
     return isGivenSln(workspace, 'slnWithGenerator');
 }
 
+export async function waitForAllAsyncOperationsAsync(exports: CSharpExtensionExports): Promise<void> {
+    const source = new vscode.CancellationTokenSource();
+    await exports.experimental.sendServerRequest(WaitForAsyncOperationsRequest.type, { operations: [] }, source.token);
+}
+
 function isGivenSln(workspace: typeof vscode.workspace, expectedProjectFileName: string) {
     const primeWorkspace = workspace.workspaceFolders![0];
     const projectFileName = primeWorkspace.uri.fsPath.split(path.sep).pop();
diff --git a/test/integrationTests/testAssets/slnWithCsproj/.vscode/lsp_tools_host_slnWithCsproj.code-workspace b/test/integrationTests/testAssets/slnWithCsproj/.vscode/lsp_tools_host_slnWithCsproj.code-workspace
index cef8ba814..f447dcd76 100644
--- a/test/integrationTests/testAssets/slnWithCsproj/.vscode/lsp_tools_host_slnWithCsproj.code-workspace
+++ b/test/integrationTests/testAssets/slnWithCsproj/.vscode/lsp_tools_host_slnWithCsproj.code-workspace
@@ -8,6 +8,6 @@
         "dotnet.defaultSolution": "b_SecondInOrder_SlnFile.sln",	
         "dotnet.server.trace": "Trace",	
         "dotnet.server.useOmnisharp": false,	
-        "omnisharp.enableLspDriver": false,	
+        "omnisharp.enableLspDriver": false,
     }	
 }
\ No newline at end of file
diff --git a/test/integrationTests/testHooks.ts b/test/integrationTests/testHooks.ts
new file mode 100644
index 000000000..95af0a2fb
--- /dev/null
+++ b/test/integrationTests/testHooks.ts
@@ -0,0 +1,25 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as lsp from 'vscode-languageserver-protocol';
+
+export interface WaitForAsyncOperationsParams {
+    /**
+     * The operations to wait for.
+     */
+    operations: string[];
+}
+
+export namespace WaitForAsyncOperationsRequest {
+    export const method = 'workspace/waitForAsyncOperations';
+    export const messageDirection: lsp.MessageDirection = lsp.MessageDirection.clientToServer;
+    export const type = new lsp.RequestType<
+        WaitForAsyncOperationsParams,
+        {
+            /* empty */
+        },
+        void
+    >(method);
+}