Merge pull request #73997 from dibarbet/flaky_lsp_disconnect_test

Fix flaky test: LanguageServerCleansUpOnUnexpectedJsonRpcDisconnectAsync
dotnet · Jun 17, 2024 · 47298b8 · 47298b8
2 parents a1d0ddd + b13ed1b
commit 47298b8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs b/src/EditorFeatures/TestUtilities/LanguageServer/AbstractLanguageServerProtocolTests.cs
@@ -710,8 +710,13 @@ public async Task AssertServerShuttingDownAsync()
             {
                 var queueAccessor = GetQueueAccessor()!.Value;
                 await queueAccessor.WaitForProcessingToStopAsync().ConfigureAwait(false);
-                Assert.True(GetServerAccessor().HasShutdownStarted());
-                Assert.True(queueAccessor.IsComplete());
+
+                var shutdownTask = GetServerAccessor().GetShutdownTaskAsync();
+                AssertEx.NotNull(shutdownTask, "Unexpected shutdown not started");
+
+                // Shutdown task will close the queue, so we need to wait for it to complete.
+                await shutdownTask.ConfigureAwait(false);
+                Assert.True(queueAccessor.IsComplete(), "Unexpected queue not complete");
             }
 
             internal async Task WaitForDiagnosticsAsync()

diff --git a/...LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework/AbstractLanguageServer.cs b/...LanguageServer/Microsoft.CommonLanguageServerProtocol.Framework/AbstractLanguageServer.cs
@@ -402,10 +402,15 @@ internal Task<TResponse> ExecuteRequestAsync<TRequest, TResponse>(string methodN
         internal JsonRpc GetServerRpc() => _server._jsonRpc;
 
         internal bool HasShutdownStarted()
+        {
+            return GetShutdownTaskAsync() != null;
+        }
+
+        internal Task? GetShutdownTaskAsync()
         {
             lock (_server._lifeCycleLock)
             {
-                return _server._shutdownRequestTask != null;
+                return _server._shutdownRequestTask;
             }
         }
     }