Plumb capabilities to SwiftPMTestProject and add a closure that can be executed before SourceKit-LSP is initialized

Author: ahoppen

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -80,8 +80,10 @@ public class MultiFileTestProject {
   public init(
     files: [RelativeFileLocation: String],
     workspaces: (URL) async throws -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
+    capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
     usePullDiagnostics: Bool = true,
+    preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     testName: String = #function
   ) async throws {
     scratchDirectory = try testScratchDir(testName: testName)
@@ -112,8 +114,10 @@ public class MultiFileTestProject {
 
     self.testClient = try await TestSourceKitLSPClient(
       serverOptions: serverOptions,
+      capabilities: capabilities,
       usePullDiagnostics: usePullDiagnostics,
       workspaceFolders: workspaces(scratchDirectory),
+      preInitialization: preInitialization,
       cleanUp: { [scratchDirectory] in
         if cleanScratchDirectories {
           try? FileManager.default.removeItem(at: scratchDirectory)

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -42,8 +42,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
     workspaces: (URL) async throws -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
     build: Bool = false,
     allowBuildFailure: Bool = false,
+    capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
     pollIndex: Bool = true,
+    preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     usePullDiagnostics: Bool = true,
     testName: String = #function
   ) async throws {
@@ -66,8 +68,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
     try await super.init(
       files: filesByPath,
       workspaces: workspaces,
+      capabilities: capabilities,
       serverOptions: serverOptions,
       usePullDiagnostics: usePullDiagnostics,
+      preInitialization: preInitialization,
       testName: testName
     )
 

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -83,6 +83,8 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///   - capabilities: The test client's capabilities.
   ///   - usePullDiagnostics: Whether to use push diagnostics or use push-based diagnostics
   ///   - workspaceFolders: Workspace folders to open.
+  ///   - preInitialization: A closure that is called after the test client is created but before SourceKit-LSP is
+  ///     initialized. This can be used to eg. register request handlers.
   ///   - cleanUp: A closure that is called when the `TestSourceKitLSPClient` is destructed.
   ///     This allows e.g. a `IndexedSingleSwiftFileTestProject` to delete its temporary files when they are no longer
   ///     needed.
@@ -94,6 +96,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
     capabilities: ClientCapabilities = ClientCapabilities(),
     usePullDiagnostics: Bool = true,
     workspaceFolders: [WorkspaceFolder]? = nil,
+    preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     cleanUp: @escaping () -> Void = {}
   ) async throws {
     if !useGlobalModuleCache {
@@ -135,7 +138,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
       guard capabilities.textDocument!.diagnostic == nil else {
         struct ConflictingDiagnosticsError: Error, CustomStringConvertible {
           var description: String {
-            "usePushDiagnostics = false is not supported if capabilities already contain diagnostic options"
+            "usePullDiagnostics = false is not supported if capabilities already contain diagnostic options"
           }
         }
         throw ConflictingDiagnosticsError()
@@ -145,6 +148,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
         XCTAssertEqual(request.registrations.only?.method, DocumentDiagnosticsRequest.method)
         return VoidResponse()
       }
+      preInitialization?(self)
     }
     if initialize {
       _ = try await self.send(