Fix a non-deterministic test failure in testPrepareTargetAfterEditToDependency

We could get a `DiagnosticsRefreshRequest` before the modified target had been re-prepared.

Author: ahoppen

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -69,7 +69,11 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///
   /// Conceptually, this is an array of `RequestHandler<any RequestType>` but
   /// since we can't express this in the Swift type system, we use `[Any]`.
-  private nonisolated(unsafe) var requestHandlers = ThreadSafeBox<[Any]>(initialValue: [])
+  ///
+  /// `isOneShort` if the request handler should only serve a single request and should be removed from
+  /// `requestHandlers` after it has been called.
+  private nonisolated(unsafe) var requestHandlers: ThreadSafeBox<[(requestHandler: Any, isOneShot: Bool)]> =
+    ThreadSafeBox(initialValue: [])
 
   /// A closure that is called when the `TestSourceKitLSPClient` is destructed.
   ///
@@ -263,7 +267,12 @@ public final class TestSourceKitLSPClient: MessageHandler {
   /// The request handler will only handle a single request. If the request is called again, the request handler won't
   /// call again
   public func handleSingleRequest<R: RequestType>(_ requestHandler: @escaping RequestHandler<R>) {
-    requestHandlers.value.append(requestHandler)
+    requestHandlers.value.append((requestHandler: requestHandler, isOneShot: true))
+  }
+
+  /// Handle all requests of the given type that are sent to the client.
+  public func handleMultipleRequests<R: RequestType>(_ requestHandler: @escaping RequestHandler<R>) {
+    requestHandlers.value.append((requestHandler: requestHandler, isOneShot: false))
   }
 
   // MARK: - Conformance to MessageHandler
@@ -280,19 +289,21 @@ public final class TestSourceKitLSPClient: MessageHandler {
     reply: @escaping (LSPResult<Request.Response>) -> Void
   ) {
     requestHandlers.withLock { requestHandlers in
-      let requestHandlerAndIndex = requestHandlers.enumerated().compactMap {
-        (index, handler) -> (RequestHandler<Request>, Int)? in
-        guard let handler = handler as? RequestHandler<Request> else {
+      let requestHandlerIndexAndIsOneShot = requestHandlers.enumerated().compactMap {
+        (index, handlerAndIsOneShot) -> (RequestHandler<Request>, Int, Bool)? in
+        guard let handler = handlerAndIsOneShot.requestHandler as? RequestHandler<Request> else {
           return nil
         }
-        return (handler, index)
+        return (handler, index, handlerAndIsOneShot.isOneShot)
       }.first
-      guard let (requestHandler, index) = requestHandlerAndIndex else {
+      guard let (requestHandler, index, isOneShot) = requestHandlerIndexAndIsOneShot else {
         reply(.failure(.methodNotFound(Request.method)))
         return
       }
       reply(.success(requestHandler(params)))
-      requestHandlers.remove(at: index)
+      if isOneShot {
+        requestHandlers.remove(at: index)
+      }
     }
   }
 

Sources/sourcekit-lsp/CMakeLists.txt

@@ -5,6 +5,7 @@ target_link_libraries(sourcekit-lsp PRIVATE
   Diagnose
   LanguageServerProtocol
   LanguageServerProtocolJSONRPC
+  SemanticIndex
   SKCore
   SKSupport
   SourceKitLSP

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -569,16 +569,9 @@ final class BackgroundIndexingTests: XCTestCase {
       DidChangeWatchedFilesNotification(changes: [FileEvent(uri: try project.uri(for: "MyFile.swift"), type: .changed)])
     )
 
-    // Send a document request for `uri` to trigger re-preparation of its target. We don't actually care about the
-    // response for this request. Instead, we wait until SourceKit-LSP sends us a `DiagnosticsRefreshRequest`,
-    // indicating that the target of `uri` has been prepared.
-    _ = try await project.testClient.send(
-      DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(uri))
-    )
-
     let receivedEmptyDiagnostics = self.expectation(description: "Received diagnostic refresh request")
 
-    project.testClient.handleSingleRequest { (_: DiagnosticsRefreshRequest) in
+    project.testClient.handleMultipleRequests { (_: DiagnosticsRefreshRequest) in
       Task {
         let updatedDiagnostics = try await project.testClient.send(
           DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(uri))
@@ -594,6 +587,13 @@ final class BackgroundIndexingTests: XCTestCase {
       return VoidResponse()
     }
 
+    // Send a document request for `uri` to trigger re-preparation of its target. We don't actually care about the
+    // response for this request. Instead, we wait until SourceKit-LSP sends us a `DiagnosticsRefreshRequest`,
+    // indicating that the target of `uri` has been prepared.
+    _ = try await project.testClient.send(
+      DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
     try await fulfillmentOfOrThrow([receivedEmptyDiagnostics])
   }