Add infrastructure to test which targets are being prepared

Author: ahoppen

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -84,6 +84,7 @@ public class MultiFileTestProject {
     serverOptions: SourceKitLSPServer.Options = .testDefault,
     usePullDiagnostics: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
+    cleanUp: (() -> Void)? = nil,
     testName: String = #function
   ) async throws {
     scratchDirectory = try testScratchDir(testName: testName)
@@ -122,6 +123,7 @@ public class MultiFileTestProject {
         if cleanScratchDirectories {
           try? FileManager.default.removeItem(at: scratchDirectory)
         }
+        cleanUp?()
       }
     )
   }

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -47,6 +47,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
     pollIndex: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     usePullDiagnostics: Bool = true,
+    cleanUp: (() -> Void)? = nil,
     testName: String = #function
   ) async throws {
     var filesByPath: [RelativeFileLocation: String] = [:]
@@ -72,6 +73,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
       serverOptions: serverOptions,
       usePullDiagnostics: usePullDiagnostics,
       preInitialization: preInitialization,
+      cleanUp: cleanUp,
       testName: testName
     )
 

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -30,7 +30,7 @@ public struct PreparationTaskDescription: IndexTaskDescription {
   public let id = preparationIDForLogging.fetchAndIncrement()
 
   /// The targets that should be prepared.
-  private let targetsToPrepare: [ConfiguredTarget]
+  public let targetsToPrepare: [ConfiguredTarget]
 
   /// The build system manager that is used to get the toolchain and build settings for the files to index.
   private let buildSystemManager: BuildSystemManager
@@ -84,7 +84,7 @@ public struct PreparationTaskDescription: IndexTaskDescription {
           "Preparation failed: \(error.forLogging)"
         )
       }
-      testHooks.preparationTaskDidFinish?(self)
+      await testHooks.preparationTaskDidFinish?(self)
       logger.log(
         "Finished preparation in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(targetsToPrepareDescription)"
       )

Sources/SemanticIndex/TestHooks.swift

@@ -12,14 +12,14 @@
 
 /// Callbacks that allow inspection of internal state modifications during testing.
 public struct IndexTestHooks: Sendable {
-  public var preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) -> Void)?
+  public var preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)?
 
   /// A callback that is called when an index task finishes.
-  public var updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) -> Void)?
+  public var updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)?
 
   public init(
-    preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) -> Void)? = nil,
-    updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) -> Void)? = nil
+    preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
+    updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil
   ) {
     self.preparationTaskDidFinish = preparationTaskDidFinish
     self.updateIndexStoreTaskDidFinish = updateIndexStoreTaskDidFinish

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -99,7 +99,7 @@ public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
       for file in filesToIndex {
         await updateIndexStoreForSingleFile(file)
       }
-      testHooks.updateIndexStoreTaskDidFinish?(self)
+      await testHooks.updateIndexStoreTaskDidFinish?(self)
       logger.log(
         "Finished updating index store in \(Date().timeIntervalSince(startDate) * 1000, privacy: .public)ms: \(filesToIndexDescription)"
       )

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -12,14 +12,59 @@
 
 import LSPTestSupport
 import LanguageServerProtocol
+import SKCore
 import SKTestSupport
+import SemanticIndex
 import SourceKitLSP
 import XCTest
 
 fileprivate let backgroundIndexingOptions = SourceKitLSPServer.Options(
   indexOptions: IndexOptions(enableBackgroundIndexing: true)
 )
 
+fileprivate actor ExpectedPreparationTracker {
+  /// The targets we expect to be prepared. For targets within the same set, we don't care about the exact order.
+  private var expectedPreparations: [Set<ConfiguredTarget>]
+
+  /// Implicitly-unwrapped optional so we can reference `self` when creating `IndexTestHooks`.
+  /// `nonisolated(unsafe)` is fine because this is not modified after `testHooks` is created.
+  nonisolated(unsafe) var testHooks: IndexTestHooks!
+
+  init(expectedPreparations: [Set<ConfiguredTarget>]) {
+    self.expectedPreparations = expectedPreparations
+    self.testHooks = IndexTestHooks(preparationTaskDidFinish: { [weak self] in
+      await self?.preparationTaskDidFinish(taskDescription: $0)
+    })
+  }
+
+  func preparationTaskDidFinish(taskDescription: PreparationTaskDescription) -> Void {
+    guard let expectedTargetsToPrepare = expectedPreparations.first else {
+      XCTFail("Didn't expect a preparation but received \(taskDescription)")
+      return
+    }
+    guard Set(taskDescription.targetsToPrepare).isSubset(of: expectedTargetsToPrepare) else {
+      XCTFail("Received unexpected preparation of \(taskDescription)")
+      return
+    }
+    let remainingExpectedTargetsToPrepare = expectedTargetsToPrepare.filter {
+      !taskDescription.targetsToPrepare.contains($0)
+    }
+    if remainingExpectedTargetsToPrepare.isEmpty {
+      expectedPreparations.remove(at: 0)
+    } else {
+      expectedPreparations[0] = remainingExpectedTargetsToPrepare
+    }
+  }
+
+  deinit {
+    let expectedPreparations = self.expectedPreparations
+    XCTAssert(
+      expectedPreparations.isEmpty,
+      "ExpectedPreparationTracker destroyed with unfulfilled expected preparations: \(expectedPreparations)."
+    )
+  }
+}
+
 final class BackgroundIndexingTests: XCTestCase {
   func testBackgroundIndexingOfSingleFile() async throws {
     let project = try await SwiftPMTestProject(
@@ -508,6 +553,19 @@ final class BackgroundIndexingTests: XCTestCase {
 
   func testPrepareTarget() async throws {
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
+    var serverOptions = backgroundIndexingOptions
+    let expectedPreparationTracker = ExpectedPreparationTracker(expectedPreparations: [
+      [
+        ConfiguredTarget(targetID: "LibA", runDestinationID: "dummy"),
+        ConfiguredTarget(targetID: "LibB", runDestinationID: "dummy"),
+      ],
+      [
+        ConfiguredTarget(targetID: "LibA", runDestinationID: "dummy"),
+        ConfiguredTarget(targetID: "LibB", runDestinationID: "dummy"),
+      ],
+    ])
+    serverOptions.indexTestHooks = expectedPreparationTracker.testHooks
+
     let project = try await SwiftPMTestProject(
       files: [
         "LibA/MyFile.swift": "",
@@ -531,7 +589,8 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      serverOptions: serverOptions,
+      cleanUp: { /* Keep expectedPreparationTracker alive */ _ = expectedPreparationTracker }
     )
 
     let (uri, _) = try project.openDocument("MyOtherFile.swift")