Fix race condition that caused testMainFileChangesIfIncludeIsAdded to fail

In some situations, we could return the timeout error from the timeout task, but receive a notification from the `self.notifications` `AsyncStream`. That notification would then be dropped and never get delivered to the test case, which can cause test cases to fail.

Author: ahoppen

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -248,21 +248,26 @@ package final class TestSourceKitLSPClient: MessageHandler, Sendable {
   /// - Note: This also returns any notifications sent before the call to
   ///   `nextNotification`.
   package func nextNotification(timeout: Duration = .seconds(defaultTimeout)) async throws -> any NotificationType {
-    return try await withThrowingTaskGroup(of: (any NotificationType).self) { taskGroup in
-      taskGroup.addTask {
-        for await notification in self.notifications {
-          return notification
-        }
-        throw NotificationTimeoutError()
-      }
-      taskGroup.addTask {
-        try await Task.sleep(for: timeout)
-        throw NotificationTimeoutError()
+    // The task that gets the next notification from `self.notifications`.
+    let notificationYielder = Task {
+      for await notification in self.notifications {
+        return notification
       }
-      let result = try await taskGroup.next()!
-      taskGroup.cancelAll()
-      return result
+      throw NotificationTimeoutError()
+    }
+    // After `timeout`, we tell `notificationYielder` that we are no longer interested in its result by cancelling it.
+    // We wait for `notificationYielder` to accept this cancellation instead of returning immediately to avoid a
+    // situation where `notificationYielder` continues running, eats the first notification but it then never gets
+    // delivered to the test because we already delivered a timeout.
+    let cancellationTask = Task {
+      try await Task.sleep(for: timeout)
+      notificationYielder.cancel()
+    }
+    defer {
+      // We have received a value and don't need the cancellationTask anymore
+      cancellationTask.cancel()
     }
+    return try await notificationYielder.value
   }
 
   /// Await the next diagnostic notification sent to the client.

Tests/SourceKitLSPTests/MainFilesProviderTests.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import LanguageServerProtocol
+import SKLogging
 import SKTestSupport
 import SourceKitLSP
 import SwiftExtensions
@@ -202,10 +203,11 @@ final class MainFilesProviderTests: XCTestCase {
     // diagnostics.
     try await repeatUntilExpectedResult {
       let refreshedDiags = try? await project.testClient.nextDiagnosticsNotification(timeout: .seconds(1))
-      guard let diagnostic = refreshedDiags?.diagnostics.only else {
+      guard refreshedDiags?.diagnostics.only?.message == "Unused variable 'fromMyFancyLibrary'" else {
+        logger.debug("Received unexpected diagnostics: \(refreshedDiags?.forLogging)")
         return false
       }
-      return diagnostic.message == "Unused variable 'fromMyFancyLibrary'"
+      return true
     }
   }
 }