Wait for request to be sent to sourcekitd before cancelling it in SwiftSourceKitPluginTests.testCancellation

Sources/SKTestSupport/Assertions.swift

@@ -119,6 +119,26 @@ package func assertNotNil<T>(
   XCTAssertNotNil(expression, message(), file: file, line: line)
 }
 
+/// Check that the string contains the given substring.
+package func assertContains(
+  _ string: some StringProtocol,
+  _ substring: some StringProtocol,
+  file: StaticString = #filePath,
+  line: UInt = #line
+) {
+  XCTAssert(string.contains(substring), "Expected to contain '\(substring)': \(string)", file: file, line: line)
+}
+
+/// Check that the sequence contains the given element.
+package func assertContains<Element: Equatable>(
+  _ sequence: some Sequence<Element>,
+  _ element: Element,
+  file: StaticString = #filePath,
+  line: UInt = #line
+) {
+  XCTAssert(sequence.contains(element), "Expected to contain '\(element)': \(sequence)", file: file, line: line)
+}
+
 /// Same as `XCTUnwrap` but doesn't take autoclosures and thus `expression`
 /// can contain `await`.
 package func unwrap<T>(

Sources/SourceKitD/DynamicallyLoadedSourceKitD.swift

@@ -108,6 +108,9 @@ package actor DynamicallyLoadedSourceKitD: SourceKitD {
   /// List of notification handlers that will be called for each notification.
   private var notificationHandlers: [WeakSKDNotificationHandler] = []
 
+  /// List of hooks that should be executed for every request sent to sourcekitd.
+  private var requestHandlingHooks: [UUID: (SKDRequestDictionary) -> Void] = [:]
+
   package static func getOrCreate(
     dylibPath: URL,
     pluginPaths: PluginPaths?
@@ -204,6 +207,29 @@ package actor DynamicallyLoadedSourceKitD: SourceKitD {
     notificationHandlers.removeAll(where: { $0.value == nil || $0.value === handler })
   }
 
+  /// Execute `body` and invoke `hook` for every sourcekitd request that is sent during the execution time of `body`.
+  ///
+  /// Note that `hook` will not only be executed for requests sent *by* body but this may also include sourcekitd
+  /// requests that were sent by other clients of the same `DynamicallyLoadedSourceKitD` instance that just happen to
+  /// send a request during that time.
+  ///
+  /// This is intended for testing only.
+  package func withRequestHandlingHook(
+    body: () async throws -> Void,
+    hook: @escaping (SKDRequestDictionary) -> Void
+  ) async rethrows {
+    let id = UUID()
+    requestHandlingHooks[id] = hook
+    defer { requestHandlingHooks[id] = nil }
+    try await body()
+  }
+
+  package func didSend(request: SKDRequestDictionary) {
+    for hook in requestHandlingHooks.values {
+      hook(request)
+    }
+  }
+
   package nonisolated func log(request: SKDRequestDictionary) {
     logger.info(
       """

Sources/SourceKitD/SourceKitD.swift

@@ -83,6 +83,10 @@ package protocol SourceKitD: AnyObject, Sendable {
   /// Removes a previously registered notification handler.
   func removeNotificationHandler(_ handler: SKDNotificationHandler) async
 
+  /// A function that gets called after the request has been sent to sourcekitd but but does not wait for results to be
+  /// received. This can be used by clients to implement hooks that should be executed for every sourcekitd request.
+  func didSend(request: SKDRequestDictionary) async
+
   /// Log the given request.
   ///
   /// This log call is issued during normal operation. It is acceptable for the logger to truncate the log message
@@ -146,6 +150,9 @@ extension SourceKitD {
         self.api.send_request(request.dict, &handle) { response in
           continuation.resume(returning: SKDResponse(response!, sourcekitd: self))
         }
+        Task {
+          await self.didSend(request: request)
+        }
         if let handle {
           return SourceKitDRequestHandle(handle: handle)
         }

Tests/BuildSystemIntegrationTests/BuildServerBuildSystemTests.swift

@@ -501,7 +501,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
         allowFallbackSettings: false
       )
     )
-    XCTAssert(try XCTUnwrap(firstOptions).compilerArguments.contains("-DFIRST"))
+    assertContains(try XCTUnwrap(firstOptions).compilerArguments, "-DFIRST")
 
     let secondOptions = try await project.testClient.send(
       SourceKitOptionsRequest(
@@ -511,7 +511,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
         allowFallbackSettings: false
       )
     )
-    XCTAssert(try XCTUnwrap(secondOptions).compilerArguments.contains("-DSECOND"))
+    assertContains(try XCTUnwrap(secondOptions).compilerArguments, "-DSECOND")
 
     let optionsWithoutTarget = try await project.testClient.send(
       SourceKitOptionsRequest(
@@ -521,7 +521,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
       )
     )
     // We currently pick the canonical target alphabetically, which means that `bsp://first` wins over `bsp://second`
-    XCTAssert(try XCTUnwrap(optionsWithoutTarget).compilerArguments.contains("-DFIRST"))
+    assertContains(try XCTUnwrap(optionsWithoutTarget).compilerArguments, "-DFIRST")
   }
 
   func testDontBlockBuildServerInitializationIfBuildSystemIsUnresponsive() async throws {

Tests/BuildSystemIntegrationTests/BuildSystemManagerTests.swift

@@ -116,7 +116,7 @@ final class BuildSystemManagerTests: XCTestCase {
     await manager.registerForChangeNotifications(for: d, language: .c)
     await assertEqual(manager.cachedMainFile(for: a), c)
     let bMain = await manager.cachedMainFile(for: b)
-    XCTAssert(Set([c, d]).contains(bMain))
+    assertContains([c, d], bMain)
     await assertEqual(manager.cachedMainFile(for: c), c)
     await assertEqual(manager.cachedMainFile(for: d), d)
 

Tests/BuildSystemIntegrationTests/SwiftPMBuildSystemTests.swift

@@ -1118,8 +1118,8 @@ final class SwiftPMBuildSystemTests: XCTestCase {
         fallbackAfterTimeout: false
       )
       let compilerArgs = try XCTUnwrap(settings?.compilerArguments)
-      XCTAssert(compilerArgs.contains("-package-description-version"))
-      XCTAssert(compilerArgs.contains(try versionSpecificManifestURL.filePath))
+      assertContains(compilerArgs, "-package-description-version")
+      assertContains(compilerArgs, try versionSpecificManifestURL.filePath)
     }
   }
 
@@ -1152,7 +1152,7 @@ final class SwiftPMBuildSystemTests: XCTestCase {
       )
       let compilerArgs = try XCTUnwrap(settings?.compilerArguments)
       assertArgumentsContain("-package-description-version", "5.1.0", arguments: compilerArgs)
-      XCTAssert(compilerArgs.contains(try manifestURL.filePath))
+      assertContains(compilerArgs, try manifestURL.filePath)
     }
   }
 }

Tests/SourceKitDTests/SourceKitDRegistryTests.swift

@@ -74,6 +74,7 @@ final class FakeSourceKitD: SourceKitD {
   var values: sourcekitd_api_values { fatalError() }
   func addNotificationHandler(_ handler: SKDNotificationHandler) { fatalError() }
   func removeNotificationHandler(_ handler: SKDNotificationHandler) { fatalError() }
+  func didSend(request: SKDRequestDictionary) {}
   private init() {
     token = nextToken.fetchAndIncrement()
   }

Tests/SourceKitLSPTests/CompilationDatabaseTests.swift

@@ -151,7 +151,7 @@ final class CompilationDatabaseTests: XCTestCase {
         XCTFail("Expected ResponseError, got \(error)")
         return
       }
-      XCTAssert(error.message.contains("No language service"))
+      assertContains(error.message, "No language service")
     }
   }
 
@@ -217,7 +217,7 @@ final class CompilationDatabaseTests: XCTestCase {
       )
     )
     let hoverContent = try XCTUnwrap(hover?.contents.markupContent?.value)
-    XCTAssert(hoverContent.contains("void main()"))
+    assertContains(hoverContent, "void main()")
 
     // But for `testFromDummyToolchain.swift`, we can't launch sourcekitd (because it doesn't exist, we just provided a
     // dummy), so we should receive an error. The exact error here is not super relevant, the important part is that we
@@ -235,7 +235,7 @@ final class CompilationDatabaseTests: XCTestCase {
         XCTFail("Expected ResponseError, got \(error)")
         return
       }
-      XCTAssert(error.message.contains("No language service"))
+      assertContains(error.message, "No language service")
     }
   }
 }

Tests/SourceKitLSPTests/LocalSwiftTests.swift

@@ -730,7 +730,7 @@ final class LocalSwiftTests: XCTestCase {
 
     XCTAssertEqual(fixit.title, "Use 'new(_:hotness:)' instead")
     XCTAssertEqual(fixit.diagnostics?.count, 1)
-    XCTAssert(fixit.diagnostics?.first?.message.contains("is deprecated") == true)
+    assertContains(try XCTUnwrap(fixit.diagnostics?.first?.message), "is deprecated")
     XCTAssertEqual(
       fixit.edit?.changes?[uri],
       [

Tests/SourceKitLSPTests/WorkspaceTests.swift

@@ -1181,7 +1181,7 @@ final class WorkspaceTests: XCTestCase {
       )
     )
     let options = try XCTUnwrap(optionsOptional)
-    XCTAssert(options.compilerArguments.contains("-module-name"))
+    assertContains(options.compilerArguments, "-module-name")
     XCTAssertEqual(options.kind, .normal)
     XCTAssertNil(options.didPrepareTarget)
   }
@@ -1366,7 +1366,7 @@ final class WorkspaceTests: XCTestCase {
         allowFallbackSettings: false
       )
     )
-    XCTAssert(options.compilerArguments.contains("-DHAVE_SETTINGS"))
+    assertContains(options.compilerArguments, "-DHAVE_SETTINGS")
   }
 
   func testOutputPaths() async throws {

Tests/SwiftSourceKitPluginTests/SwiftSourceKitPluginTests.swift

@@ -311,22 +311,27 @@ final class SwiftSourceKitPluginTests: XCTestCase {
       compilerArguments: [path]
     )
 
+    let slowCompletionRequestSent = self.expectation(description: "slow completion result sent")
     let slowCompletionResultReceived = self.expectation(description: "slow completion")
-    let slowCompletionTask = Task {
-      do {
-        _ = try await sourcekitd.completeOpen(
-          path: path,
-          position: positions["1️⃣"],
-          filter: ""
-        )
-        XCTFail("Expected completion to be cancelled")
-      } catch {
-        XCTAssert(error is CancellationError, "Expected completion to be cancelled, failed with \(error)")
+    let dynamicallyLoadedSourceKitd = try XCTUnwrap(sourcekitd as? DynamicallyLoadedSourceKitD)
+    try await dynamicallyLoadedSourceKitd.withRequestHandlingHook {
+      let slowCompletionTask = Task {
+        await assertThrowsError(try await sourcekitd.completeOpen(path: path, position: positions["1️⃣"], filter: "")) {
+          XCTAssert($0 is CancellationError, "Expected completion to be cancelled, failed with \($0)")
+        }
+        slowCompletionResultReceived.fulfill()
       }
-      slowCompletionResultReceived.fulfill()
+      // Wait for the slow completion request to actually be sent to sourcekitd. Otherwise, we might hit a cancellation
+      // check somewhere during request sending and we aren't actually sending the completion request to sourcekitd.
+      try await fulfillmentOfOrThrow(slowCompletionRequestSent)
+
+      slowCompletionTask.cancel()
+      try await fulfillmentOfOrThrow(slowCompletionResultReceived, timeout: 30)
+    } hook: { request in
+      // Check that we aren't matching against a request sent by something else that has handle to the same sourcekitd.
+      assertContains(request.description.replacing(#"\\"#, with: #"\"#), path)
+      slowCompletionRequestSent.fulfill()
     }
-    slowCompletionTask.cancel()
-    try await fulfillmentOfOrThrow(slowCompletionResultReceived, timeout: 30)
 
     let fastCompletionStarted = Date()
     let result = try await sourcekitd.completeOpen(