Merge pull request #2035 from ahoppen/build-server-tests

Add a new test project type that uses a custom build server

Sources/BuildSystemIntegration/BuildSystemManager.swift

@@ -179,9 +179,9 @@ private extension BuildSystemSpec {
     _ createBuildSystem: @Sendable (_ connectionToSourceKitLSP: any Connection) async throws -> BuiltInBuildSystem?
   ) async -> BuildSystemAdapter? {
     let connectionToSourceKitLSP = LocalConnection(
-      receiverName: "BuildSystemManager for \(projectRoot.lastPathComponent)"
+      receiverName: "BuildSystemManager for \(projectRoot.lastPathComponent)",
+      handler: messagesToSourceKitLSPHandler
     )
-    connectionToSourceKitLSP.start(handler: messagesToSourceKitLSPHandler)
 
     let buildSystem = await orLog("Creating build system") {
       try await createBuildSystem(connectionToSourceKitLSP)
@@ -197,9 +197,9 @@ private extension BuildSystemSpec {
       buildSystemHooks: buildSystemHooks
     )
     let connectionToBuildSystem = LocalConnection(
-      receiverName: "\(type(of: buildSystem)) for \(projectRoot.lastPathComponent)"
+      receiverName: "\(type(of: buildSystem)) for \(projectRoot.lastPathComponent)",
+      handler: buildSystemAdapter
     )
-    connectionToBuildSystem.start(handler: buildSystemAdapter)
     return .builtIn(buildSystemAdapter, connectionToBuildSystem: connectionToBuildSystem)
   }
 
@@ -266,9 +266,9 @@ private extension BuildSystemSpec {
       #endif
     case .injected(let injector):
       let connectionToSourceKitLSP = LocalConnection(
-        receiverName: "BuildSystemManager for \(projectRoot.lastPathComponent)"
+        receiverName: "BuildSystemManager for \(projectRoot.lastPathComponent)",
+        handler: messagesToSourceKitLSPHandler
       )
-      connectionToSourceKitLSP.start(handler: messagesToSourceKitLSPHandler)
       return .injected(
         await injector(projectRoot, connectionToSourceKitLSP)
       )
@@ -510,8 +510,7 @@ package actor BuildSystemManager: QueueBasedMessageHandler {
           connectionToSourceKitLSP: legacyBuildServer.connectionToSourceKitLSP,
           buildSystemHooks: buildSystemHooks
         )
-        let connectionToBuildSystem = LocalConnection(receiverName: "Legacy BSP server")
-        connectionToBuildSystem.start(handler: adapter)
+        let connectionToBuildSystem = LocalConnection(receiverName: "Legacy BSP server", handler: adapter)
         self.buildSystemAdapter = .builtIn(adapter, connectionToBuildSystem: connectionToBuildSystem)
       }
       Task {

Sources/BuildSystemIntegration/CMakeLists.txt

@@ -19,8 +19,7 @@ add_library(BuildSystemIntegration STATIC
   LegacyBuildServerBuildSystem.swift
   MainFilesProvider.swift
   SplitShellCommand.swift
-  SwiftPMBuildSystem.swift
-  TestBuildSystem.swift)
+  SwiftPMBuildSystem.swift)
 set_target_properties(BuildSystemIntegration PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})
 target_link_libraries(BuildSystemIntegration PUBLIC

Sources/BuildSystemIntegration/LegacyBuildServerBuildSystem.swift

@@ -73,8 +73,10 @@ actor LegacyBuildServerBuildSystem: MessageHandler, BuiltInBuildSystem {
     self.configPath = configPath
     self.indexDatabasePath = nil
     self.indexStorePath = nil
-    self.connectionToSourceKitLSP = LocalConnection(receiverName: "BuildSystemManager")
-    await self.connectionToSourceKitLSP.start(handler: externalBuildSystemAdapter.messagesToSourceKitLSPHandler)
+    self.connectionToSourceKitLSP = LocalConnection(
+      receiverName: "BuildSystemManager",
+      handler: await externalBuildSystemAdapter.messagesToSourceKitLSPHandler
+    )
     await externalBuildSystemAdapter.changeMessageToSourceKitLSPHandler(to: self)
     self.buildServer = await externalBuildSystemAdapter.connectionToBuildServer
   }

Sources/LanguageServerProtocolExtensions/LocalConnection.swift

@@ -56,6 +56,11 @@ package final class LocalConnection: Connection, Sendable {
     self.name = receiverName
   }
 
+  package convenience init(receiverName: String, handler: MessageHandler) {
+    self.init(receiverName: receiverName)
+    self.start(handler: handler)
+  }
+
   deinit {
     queue.sync {
       if state != .closed {

Sources/SKTestSupport/CustomBuildServerTestProject.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -10,45 +10,53 @@
 //
 //===----------------------------------------------------------------------===//
 
-import Foundation
+import BuildSystemIntegration
 import LanguageServerProtocolExtensions
 import SKLogging
-import SKOptions
+import SourceKitLSP
+import SwiftExtensions
 import ToolchainRegistry
+import XCTest
 
 #if compiler(>=6)
 package import BuildServerProtocol
+package import Foundation
 package import LanguageServerProtocol
+package import SKOptions
 #else
 import BuildServerProtocol
+import Foundation
 import LanguageServerProtocol
+import SKOptions
 #endif
 
-/// Build system to be used for testing BuildSystem and BuildSystemDelegate functionality with SourceKitLSPServer
-/// and other components.
-package actor TestBuildSystem: MessageHandler {
-  private let connectionToSourceKitLSP: any Connection
+// MARK: - CustomBuildServer
 
-  /// Build settings by file.
-  private var buildSettingsByFile: [DocumentURI: TextDocumentSourceKitOptionsResponse] = [:]
+/// A build server that can be injected into `CustomBuildServerTestProject`.
+package protocol CustomBuildServer: MessageHandler {
+  init(projectRoot: URL, connectionToSourceKitLSP: any Connection)
 
-  package func setBuildSettings(for uri: DocumentURI, to buildSettings: TextDocumentSourceKitOptionsResponse?) {
-    buildSettingsByFile[uri] = buildSettings
-    connectionToSourceKitLSP.send(OnBuildTargetDidChangeNotification(changes: nil))
-  }
-
-  private let initializeData: SourceKitInitializeBuildResponseData
-
-  package init(
-    initializeData: SourceKitInitializeBuildResponseData = SourceKitInitializeBuildResponseData(
-      sourceKitOptionsProvider: true
-    ),
-    connectionToSourceKitLSP: any Connection
-  ) {
-    self.initializeData = initializeData
-    self.connectionToSourceKitLSP = connectionToSourceKitLSP
-  }
+  func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse
+  func onBuildInitialized(_ notification: OnBuildInitializedNotification) throws
+  func buildShutdown(_ request: BuildShutdownRequest) async throws -> VoidResponse
+  func onBuildExit(_ notification: OnBuildExitNotification) throws
+  func workspaceBuildTargetsRequest(
+    _ request: WorkspaceBuildTargetsRequest
+  ) async throws -> WorkspaceBuildTargetsResponse
+  func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse
+  func textDocumentSourceKitOptionsRequest(
+    _ request: TextDocumentSourceKitOptionsRequest
+  ) async throws -> TextDocumentSourceKitOptionsResponse?
+  func prepareTarget(_ request: BuildTargetPrepareRequest) async throws -> VoidResponse
+  func waitForBuildSystemUpdates(request: WorkspaceWaitForBuildSystemUpdatesRequest) async -> VoidResponse
+  nonisolated func onWatchedFilesDidChange(_ notification: OnWatchedFilesDidChangeNotification) throws
+  func workspaceWaitForBuildSystemUpdatesRequest(
+    _ request: WorkspaceWaitForBuildSystemUpdatesRequest
+  ) async throws -> VoidResponse
+  nonisolated func cancelRequest(_ notification: CancelRequestNotification) throws
+}
 
+extension CustomBuildServer {
   package nonisolated func handle(_ notification: some NotificationType) {
     do {
       switch notification {
@@ -64,7 +72,7 @@ package actor TestBuildSystem: MessageHandler {
         throw ResponseError.methodNotFound(type(of: notification).method)
       }
     } catch {
-      logger.error("Error while handling BSP notification")
+      logger.error("Error while handling BSP notification: \(error.forLogging)")
     }
   }
 
@@ -102,10 +110,18 @@ package actor TestBuildSystem: MessageHandler {
       reply(.failure(ResponseError.methodNotFound(type(of: request).method)))
     }
   }
+}
 
-  func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse {
-    return InitializeBuildResponse(
-      displayName: "TestBuildSystem",
+package extension CustomBuildServer {
+  // MARK: Helper functions for the implementation of BSP methods
+
+  func initializationResponse(
+    initializeData: SourceKitInitializeBuildResponseData = SourceKitInitializeBuildResponseData(
+      sourceKitOptionsProvider: true
+    )
+  ) -> InitializeBuildResponse {
+    InitializeBuildResponse(
+      displayName: "\(type(of: self))",
       version: "",
       bspVersion: "2.2.0",
       capabilities: BuildServerCapabilities(),
@@ -114,17 +130,25 @@ package actor TestBuildSystem: MessageHandler {
     )
   }
 
-  nonisolated func onBuildInitialized(_ notification: OnBuildInitializedNotification) throws {
-    // Nothing to do
+  func dummyTargetSourcesResponse(_ files: some Sequence<DocumentURI>) -> BuildTargetSourcesResponse {
+    return BuildTargetSourcesResponse(items: [
+      SourcesItem(target: .dummy, sources: files.map { SourceItem(uri: $0, kind: .file, generated: false) })
+    ])
   }
 
+  // MARK: Default implementation for all build server methods that usually don't need customization.
+
+  func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse {
+    return initializationResponse()
+  }
+
+  nonisolated func onBuildInitialized(_ notification: OnBuildInitializedNotification) throws {}
+
   func buildShutdown(_ request: BuildShutdownRequest) async throws -> VoidResponse {
     return VoidResponse()
   }
 
-  nonisolated func onBuildExit(_ notification: OnBuildExitNotification) throws {
-    // Nothing to do
-  }
+  nonisolated func onBuildExit(_ notification: OnBuildExitNotification) throws {}
 
   func workspaceBuildTargetsRequest(
     _ request: WorkspaceBuildTargetsRequest
@@ -142,32 +166,15 @@ package actor TestBuildSystem: MessageHandler {
     ])
   }
 
-  func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse {
-    return BuildTargetSourcesResponse(items: [
-      SourcesItem(
-        target: .dummy,
-        sources: buildSettingsByFile.keys.map { SourceItem(uri: $0, kind: .file, generated: false) }
-      )
-    ])
-  }
-
-  func textDocumentSourceKitOptionsRequest(
-    _ request: TextDocumentSourceKitOptionsRequest
-  ) async throws -> TextDocumentSourceKitOptionsResponse? {
-    return buildSettingsByFile[request.textDocument.uri]
-  }
-
   func prepareTarget(_ request: BuildTargetPrepareRequest) async throws -> VoidResponse {
     return VoidResponse()
   }
 
-  package func waitForBuildSystemUpdates(request: WorkspaceWaitForBuildSystemUpdatesRequest) async -> VoidResponse {
+  func waitForBuildSystemUpdates(request: WorkspaceWaitForBuildSystemUpdatesRequest) async -> VoidResponse {
     return VoidResponse()
   }
 
-  nonisolated func onWatchedFilesDidChange(_ notification: OnWatchedFilesDidChangeNotification) throws {
-    // Not watching any files
-  }
+  nonisolated func onWatchedFilesDidChange(_ notification: OnWatchedFilesDidChangeNotification) throws {}
 
   func workspaceWaitForBuildSystemUpdatesRequest(
     _ request: WorkspaceWaitForBuildSystemUpdatesRequest
@@ -177,3 +184,40 @@ package actor TestBuildSystem: MessageHandler {
 
   nonisolated func cancelRequest(_ notification: CancelRequestNotification) throws {}
 }
+
+// MARK: - CustomBuildServerTestProject
+
+/// A test project that launches a custom build server in-process.
+///
+/// In contrast to `ExternalBuildServerTestProject`, the custom build system runs in-process and is implemented in
+/// Swift.
+package final class CustomBuildServerTestProject<BuildServer: CustomBuildServer>: MultiFileTestProject {
+  private let buildServerBox = ThreadSafeBox<BuildServer?>(initialValue: nil)
+
+  package init(
+    files: [RelativeFileLocation: String],
+    buildServer buildServerType: BuildServer.Type,
+    options: SourceKitLSPOptions? = nil,
+    enableBackgroundIndexing: Bool = false,
+    testName: String = #function
+  ) async throws {
+    let hooks: Hooks = Hooks(
+      buildSystemHooks: BuildSystemHooks(injectBuildServer: { [buildServerBox] projectRoot, connectionToSourceKitLSP in
+        let buildServer = BuildServer(projectRoot: projectRoot, connectionToSourceKitLSP: connectionToSourceKitLSP)
+        buildServerBox.value = buildServer
+        return LocalConnection(receiverName: "TestBuildSystem", handler: buildServer)
+      })
+    )
+    try await super.init(
+      files: files,
+      options: options,
+      hooks: hooks,
+      enableBackgroundIndexing: enableBackgroundIndexing,
+      testName: testName
+    )
+  }
+
+  package func buildServer(file: StaticString = #filePath, line: UInt = #line) throws -> BuildServer {
+    try XCTUnwrap(buildServerBox.value, "Accessing build server before it has been created", file: file, line: line)
+  }
+}

Sources/SKTestSupport/ExternalBuildServerTestProject.swift

@@ -59,7 +59,7 @@ private let skTestSupportInputsDirectory: URL = {
 ///
 /// The build server can contain `$SDK_ARGS`, which will replaced by `"-sdk", "/path/to/sdk"` on macOS and by an empty
 /// string on all other platforms.
-package class BuildServerTestProject: MultiFileTestProject {
+package class ExternalBuildServerTestProject: MultiFileTestProject {
   package init(
     files: [RelativeFileLocation: String],
     buildServerConfigLocation: RelativeFileLocation = ".bsp/sourcekit-lsp.json",

Tests/BuildSystemIntegrationTests/BuildServerBuildSystemTests.swift

@@ -28,7 +28,7 @@ import WinSDK
 
 final class BuildServerBuildSystemTests: XCTestCase {
   func testBuildSettingsFromBuildServer() async throws {
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": """
         #if DEBUG
@@ -85,7 +85,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
   func testBuildTargetsChanged() async throws {
     try SkipUnless.longTestsEnabled()
 
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": """
         #if DEBUG
@@ -168,7 +168,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
   func testSettingsOfSingleFileChanged() async throws {
     try SkipUnless.longTestsEnabled()
 
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": """
         #if DEBUG
@@ -251,7 +251,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
   func testCrashRecovery() async throws {
     try SkipUnless.longTestsEnabled()
 
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Crash.swift": "",
         "Test.swift": """
@@ -322,7 +322,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
   }
 
   func testBuildServerConfigAtLegacyLocation() async throws {
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": """
         #if DEBUG
@@ -378,7 +378,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
   }
 
   func testBuildSettingsDataPassThrough() async throws {
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": ""
       ],
@@ -431,7 +431,7 @@ final class BuildServerBuildSystemTests: XCTestCase {
   }
 
   func testBuildSettingsForFilePartOfMultipleTargets() async throws {
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": ""
       ],
@@ -526,48 +526,37 @@ final class BuildServerBuildSystemTests: XCTestCase {
 
   func testDontBlockBuildServerInitializationIfBuildSystemIsUnresponsive() async throws {
     // A build server that responds to the initialize request but not to any other requests.
-    final class UnresponsiveBuildServer: MessageHandler {
-      func handle(_ notification: some LanguageServerProtocol.NotificationType) {}
+    final class UnresponsiveBuildServer: CustomBuildServer {
+      init(projectRoot: URL, connectionToSourceKitLSP: any Connection) {}
 
-      func handle<Request: RequestType>(
-        _ request: Request,
-        id: RequestID,
-        reply: @escaping @Sendable (LSPResult<Request.Response>) -> Void
-      ) {
-        switch request {
-        case is InitializeBuildRequest:
-          reply(
-            .success(
-              InitializeBuildResponse(
-                displayName: "UnresponsiveBuildServer",
-                version: "",
-                bspVersion: "2.2.0",
-                capabilities: BuildServerCapabilities()
-              ) as! Request.Response
-            )
-          )
-        default:
-          #if os(Windows)
-          Sleep(60 * 60 * 1000 /*ms*/)
-          #else
-          sleep(60 * 60 /*s*/)
-          #endif
-          XCTFail("Build server should be terminated before finishing the timeout")
-        }
+      func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse {
+        #if os(Windows)
+        Sleep(60 * 60 * 1000 /*ms*/)
+        #else
+        sleep(60 * 60 /*s*/)
+        #endif
+        XCTFail("Build server should be terminated before finishing the timeout")
+        throw ResponseError.methodNotFound(BuildTargetSourcesRequest.method)
+      }
+
+      func textDocumentSourceKitOptionsRequest(
+        _ request: TextDocumentSourceKitOptionsRequest
+      ) async throws -> TextDocumentSourceKitOptionsResponse? {
+        #if os(Windows)
+        Sleep(60 * 60 * 1000 /*ms*/)
+        #else
+        sleep(60 * 60 /*s*/)
+        #endif
+        XCTFail("Build server should be terminated before finishing the timeout")
+        throw ResponseError.methodNotFound(TextDocumentSourceKitOptionsRequest.method)
       }
     }
 
-    // Creating the `MultiFileTestProject` waits for the initialize response and times out if it doesn't receive one.
+    // Creating the `CustomBuildServerTestProject` waits for the initialize response and times out if it doesn't receive one.
     // Make sure that we get that response back.
-    _ = try await MultiFileTestProject(
+    _ = try await CustomBuildServerTestProject(
       files: ["Test.swift": ""],
-      hooks: Hooks(
-        buildSystemHooks: BuildSystemHooks(injectBuildServer: { _, _ in
-          let connection = LocalConnection(receiverName: "Unresponsive build system")
-          connection.start(handler: UnresponsiveBuildServer())
-          return connection
-        })
-      )
+      buildServer: UnresponsiveBuildServer.self
     )
   }
 }

Tests/BuildSystemIntegrationTests/BuildSystemManagerTests.swift

@@ -22,6 +22,30 @@ import TSCBasic
 import ToolchainRegistry
 import XCTest
 
+fileprivate actor TestBuildSystem: CustomBuildServer {
+  private let connectionToSourceKitLSP: any Connection
+  private var buildSettingsByFile: [DocumentURI: TextDocumentSourceKitOptionsResponse] = [:]
+
+  func setBuildSettings(for uri: DocumentURI, to buildSettings: TextDocumentSourceKitOptionsResponse?) {
+    buildSettingsByFile[uri] = buildSettings
+    connectionToSourceKitLSP.send(OnBuildTargetDidChangeNotification(changes: nil))
+  }
+
+  init(projectRoot: URL, connectionToSourceKitLSP: any Connection) {
+    self.connectionToSourceKitLSP = connectionToSourceKitLSP
+  }
+
+  func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) -> BuildTargetSourcesResponse {
+    return dummyTargetSourcesResponse(buildSettingsByFile.keys)
+  }
+
+  func textDocumentSourceKitOptionsRequest(
+    _ request: TextDocumentSourceKitOptionsRequest
+  ) async throws -> TextDocumentSourceKitOptionsResponse? {
+    return buildSettingsByFile[request.textDocument.uri]
+  }
+}
+
 fileprivate extension BuildSystemManager {
   func fileBuildSettingsChanged(_ changedFiles: Set<DocumentURI>) async {
     handle(OnBuildTargetDidChangeNotification(changes: nil))
@@ -45,11 +69,9 @@ final class BuildSystemManagerTests: XCTestCase {
     let spec = BuildSystemSpec(
       kind: .injected({ projectRoot, connectionToSourceKitLSP in
         assert(testBuildSystem.value == nil, "Build system injector hook can only create a single TestBuildSystem")
-        let buildSystem = TestBuildSystem(connectionToSourceKitLSP: connectionToSourceKitLSP)
+        let buildSystem = TestBuildSystem(projectRoot: projectRoot, connectionToSourceKitLSP: connectionToSourceKitLSP)
         testBuildSystem.value = buildSystem
-        let connection = LocalConnection(receiverName: "TestBuildSystem")
-        connection.start(handler: buildSystem)
-        return connection
+        return LocalConnection(receiverName: "TestBuildSystem", handler: buildSystem)
       }),
       projectRoot: dummyPath,
       configPath: dummyPath

Tests/BuildSystemIntegrationTests/LegacyBuildServerBuildSystemTests.swift

@@ -21,7 +21,7 @@ import XCTest
 
 final class LegacyBuildServerBuildSystemTests: XCTestCase {
   func testBuildSettingsFromBuildServer() async throws {
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": """
         #if DEBUG
@@ -57,7 +57,7 @@ final class LegacyBuildServerBuildSystemTests: XCTestCase {
   }
 
   func testBuildSettingsFromBuildServerChanged() async throws {
-    let project = try await BuildServerTestProject(
+    let project = try await ExternalBuildServerTestProject(
       files: [
         "Test.swift": """
         #if DEBUG

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -1954,31 +1954,50 @@ final class BackgroundIndexingTests: XCTestCase {
   }
 
   func testIndexFileIfBuildTargetsChange() async throws {
-    let testBuildSystem = ThreadSafeBox<TestBuildSystem?>(initialValue: nil)
-    let project = try await MultiFileTestProject(
+    actor BuildServer: CustomBuildServer {
+      private let projectRoot: URL
+      private let connectionToSourceKitLSP: any Connection
+      private var buildSettingsByFile: [DocumentURI: TextDocumentSourceKitOptionsResponse] = [:]
+
+      package func setBuildSettings(for uri: DocumentURI, to buildSettings: TextDocumentSourceKitOptionsResponse?) {
+        buildSettingsByFile[uri] = buildSettings
+        connectionToSourceKitLSP.send(OnBuildTargetDidChangeNotification(changes: nil))
+      }
+
+      init(projectRoot: URL, connectionToSourceKitLSP: any Connection) {
+        self.projectRoot = projectRoot
+        self.connectionToSourceKitLSP = connectionToSourceKitLSP
+      }
+
+      func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse {
+        return initializationResponse(
+          initializeData: SourceKitInitializeBuildResponseData(
+            indexDatabasePath: try projectRoot.appendingPathComponent("index-db").filePath,
+            indexStorePath: try projectRoot.appendingPathComponent("index-store").filePath,
+            prepareProvider: true,
+            sourceKitOptionsProvider: true
+          )
+        )
+      }
+
+      func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) -> BuildTargetSourcesResponse {
+        return dummyTargetSourcesResponse(buildSettingsByFile.keys)
+      }
+
+      func textDocumentSourceKitOptionsRequest(
+        _ request: TextDocumentSourceKitOptionsRequest
+      ) -> TextDocumentSourceKitOptionsResponse? {
+        return buildSettingsByFile[request.textDocument.uri]
+      }
+    }
+
+    let project = try await CustomBuildServerTestProject(
       files: [
         "Test.swift": """
         func 1️⃣myTestFunc() {}
         """
       ],
-      hooks: Hooks(
-        buildSystemHooks: BuildSystemHooks(injectBuildServer: { projectRoot, connectionToSourceKitLSP in
-          assert(testBuildSystem.value == nil, "Build system injector hook can only create a single TestBuildSystem")
-          let buildSystem = TestBuildSystem(
-            initializeData: SourceKitInitializeBuildResponseData(
-              indexDatabasePath: projectRoot.appendingPathComponent("index-db").path,
-              indexStorePath: projectRoot.appendingPathComponent("index-store").path,
-              prepareProvider: true,
-              sourceKitOptionsProvider: true
-            ),
-            connectionToSourceKitLSP: connectionToSourceKitLSP
-          )
-          testBuildSystem.value = buildSystem
-          let connection = LocalConnection(receiverName: "TestBuildSystem")
-          connection.start(handler: buildSystem)
-          return connection
-        })
-      ),
+      buildServer: BuildServer.self,
       enableBackgroundIndexing: true
     )
     let fileUrl = try XCTUnwrap(project.uri(for: "Test.swift").fileURL)
@@ -1990,7 +2009,7 @@ final class BackgroundIndexingTests: XCTestCase {
 
     // We don't initially index Test.swift because we don't have build settings for it.
 
-    try await XCTUnwrap(testBuildSystem.value).setBuildSettings(
+    try await project.buildServer().setBuildSettings(
       for: DocumentURI(fileUrl),
       to: TextDocumentSourceKitOptionsResponse(compilerArguments: compilerArguments)
     )
@@ -2086,65 +2105,59 @@ final class BackgroundIndexingTests: XCTestCase {
     // The build system returns too many arguments to fit them into a command line invocation, so we need to use a
     // response file to invoke the indexer.
 
-    let project = try await BuildServerTestProject(
+    final class BuildServer: CustomBuildServer {
+      private let projectRoot: URL
+      private var testFileURL: URL { projectRoot.appendingPathComponent("Test File.swift") }
+
+      init(projectRoot: URL, connectionToSourceKitLSP: any LanguageServerProtocol.Connection) {
+        self.projectRoot = projectRoot
+      }
+
+      func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse {
+        return initializationResponse(
+          initializeData: SourceKitInitializeBuildResponseData(
+            indexDatabasePath: try projectRoot.appendingPathComponent("index-db").filePath,
+            indexStorePath: try projectRoot.appendingPathComponent("index-store").filePath,
+            prepareProvider: true,
+            sourceKitOptionsProvider: true
+          )
+        )
+      }
+
+      func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse {
+        return BuildTargetSourcesResponse(items: [
+          SourcesItem(
+            target: .dummy,
+            sources: [
+              SourceItem(uri: URI(testFileURL), kind: .file, generated: false)
+            ]
+          )
+        ])
+      }
+
+      func textDocumentSourceKitOptionsRequest(
+        _ request: TextDocumentSourceKitOptionsRequest
+      ) async throws -> TextDocumentSourceKitOptionsResponse? {
+        var arguments =
+          [try testFileURL.filePath] + (0..<50_000).map { "-DTHIS_IS_AN_OPTION_THAT_CONTAINS_MANY_BYTES_\($0)" }
+        if let defaultSDKPath {
+          arguments += ["-sdk", defaultSDKPath]
+        }
+        return TextDocumentSourceKitOptionsResponse(
+          compilerArguments: arguments
+        )
+      }
+
+    }
+
+    let project = try await CustomBuildServerTestProject(
       files: [
         // File name contains a space to ensure we escape it in the response file.
         "Test File.swift": """
         func 1️⃣myTestFunc() {}
         """
       ],
-      buildServer: """
-        class BuildServer(AbstractBuildServer):
-
-          def initialize(self, request: Dict[str, object]) -> Dict[str, object]:
-            return {
-              "displayName": "test server",
-              "version": "0.1",
-              "bspVersion": "2.0",
-              "rootUri": "blah",
-              "capabilities": {"languageIds": ["swift", "c", "cpp", "objective-c", "objective-c"]},
-              "data": {
-                "indexDatabasePath": r"$TEST_DIR/index-db",
-                "indexStorePath": r"$TEST_DIR/index",
-                "prepareProvider": True,
-                "sourceKitOptionsProvider": True,
-              },
-            }
-
-          def workspace_build_targets(self, request: Dict[str, object]) -> Dict[str, object]:
-            return {
-              "targets": [
-                {
-                  "id": {"uri": "bsp://dummy"},
-                  "tags": [],
-                  "languageIds": [],
-                  "dependencies": [],
-                  "capabilities": {},
-                }
-              ]
-            }
-
-          def buildtarget_sources(self, request: Dict[str, object]) -> Dict[str, object]:
-            return {
-              "items": [
-                {
-                  "target": {"uri": "bsp://dummy"},
-                  "sources": [
-                    {"uri": "$TEST_DIR_URL/Test.swift", "kind": 1, "generated": False}
-                  ],
-                }
-              ]
-            }
-
-          def textdocument_sourcekitoptions(self, request: Dict[str, object]) -> Dict[str, object]:
-            return {
-              "compilerArguments": [r"$TEST_DIR/Test File.swift", "-DDEBUG", $SDK_ARGS] + \
-                [f"-DTHIS_IS_AN_OPTION_THAT_CONTAINS_MANY_BYTES_{i}" for i in range(0, 50_000)]
-            }
-
-          def buildtarget_prepare(self, request: Dict[str, object]) -> Dict[str, object]:
-            return {}
-        """,
+      buildServer: BuildServer.self,
       enableBackgroundIndexing: true
     )
     try await project.testClient.send(PollIndexRequest())

Tests/SourceKitLSPTests/BuildSystemTests.swift

@@ -23,6 +23,30 @@ import TSCBasic
 import ToolchainRegistry
 import XCTest
 
+fileprivate actor TestBuildSystem: CustomBuildServer {
+  private let connectionToSourceKitLSP: any Connection
+  private var buildSettingsByFile: [DocumentURI: TextDocumentSourceKitOptionsResponse] = [:]
+
+  func setBuildSettings(for uri: DocumentURI, to buildSettings: TextDocumentSourceKitOptionsResponse?) {
+    buildSettingsByFile[uri] = buildSettings
+    connectionToSourceKitLSP.send(OnBuildTargetDidChangeNotification(changes: nil))
+  }
+
+  init(projectRoot: URL, connectionToSourceKitLSP: any Connection) {
+    self.connectionToSourceKitLSP = connectionToSourceKitLSP
+  }
+
+  func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) -> BuildTargetSourcesResponse {
+    return dummyTargetSourcesResponse(buildSettingsByFile.keys)
+  }
+
+  func textDocumentSourceKitOptionsRequest(
+    _ request: TextDocumentSourceKitOptionsRequest
+  ) async throws -> TextDocumentSourceKitOptionsResponse? {
+    return buildSettingsByFile[request.textDocument.uri]
+  }
+}
+
 final class BuildSystemTests: XCTestCase {
   /// The mock client used to communicate with the SourceKit-LSP server.p
   ///
@@ -54,11 +78,12 @@ final class BuildSystemTests: XCTestCase {
       buildSystemSpec: BuildSystemSpec(
         kind: .injected({ projectRoot, connectionToSourceKitLSP in
           assert(testBuildSystem.value == nil, "Build system injector hook can only create a single TestBuildSystem")
-          let buildSystem = TestBuildSystem(connectionToSourceKitLSP: connectionToSourceKitLSP)
+          let buildSystem = TestBuildSystem(
+            projectRoot: projectRoot,
+            connectionToSourceKitLSP: connectionToSourceKitLSP
+          )
           testBuildSystem.value = buildSystem
-          let connection = LocalConnection(receiverName: "TestBuildSystem")
-          connection.start(handler: buildSystem)
-          return connection
+          return LocalConnection(receiverName: "TestBuildSystem", handler: buildSystem)
         }),
         projectRoot: URL(fileURLWithPath: "/"),
         configPath: URL(fileURLWithPath: "/")