Add a request to get all tests within a file

Author: ahoppen

Sources/LanguageServerProtocol/Messages.swift

@@ -45,6 +45,7 @@ public let builtinRequests: [_RequestType.Type] = [
   DocumentSemanticTokensRangeRequest.self,
   DocumentSemanticTokensRequest.self,
   DocumentSymbolRequest.self,
+  DocumentTestsRequest.self,
   ExecuteCommandRequest.self,
   FoldingRangeRequest.self,
   HoverRequest.self,

Sources/LanguageServerProtocol/Requests/DocumentTestsRequest.swift

@@ -0,0 +1,25 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 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
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+/// A request that returns symbols for all the test classes and test methods within a file.
+///
+/// **(LSP Extension)**
+public struct DocumentTestsRequest: TextDocumentRequest, Hashable {
+  public static let method: String = "workspace/tests"
+  public typealias Response = [WorkspaceSymbolItem]?
+
+  public var textDocument: TextDocumentIdentifier
+
+  public init(textDocument: TextDocumentIdentifier) {
+    self.textDocument = textDocument
+  }
+}

Sources/SourceKitLSP/SourceKitServer.swift

@@ -840,6 +840,8 @@ extension SourceKitServer: MessageHandler {
       await request.reply { try await workspaceSymbols(request.params) }
     case let request as RequestAndReply<WorkspaceTestsRequest>:
       await request.reply { try await workspaceTests(request.params) }
+      case let request as RequestAndReply<DocumentTestsRequest>:
+      await self.handleRequest(for: request, requestHandler: self.documentTests)
     case let request as RequestAndReply<PollIndexRequest>:
       await request.reply { try await pollIndex(request.params) }
     case let request as RequestAndReply<BarrierRequest>:
@@ -1500,14 +1502,20 @@ extension SourceKitServer {
 
   func workspaceTests(_ req: WorkspaceTestsRequest) async throws -> [WorkspaceSymbolItem]? {
     let testSymbols = workspaces.flatMap { (workspace) -> [SymbolOccurrence] in
-      guard let index = workspace.index else {
-        return []
-      }
-      return index.unitTests()
+      return workspace.index?.unitTests() ?? []
     }
     return testSymbols.map(WorkspaceSymbolItem.init)
   }
 
+  func documentTests(
+    _ req: DocumentTestsRequest, 
+    workspace: Workspace,
+    languageService: ToolchainLanguageServer
+  ) async throws -> [WorkspaceSymbolItem]? {
+    let testSymbols = workspace.index?.unitTests(referencedByMainFiles: [req.textDocument.uri.pseudoPath]) ?? []
+    return testSymbols.map(WorkspaceSymbolItem.init)
+  }
+
   /// Forwards a SymbolInfoRequest to the appropriate toolchain service for this document.
   func symbolInfo(
     _ req: SymbolInfoRequest,

Tests/SourceKitLSPTests/WorkspaceTestsTests.swift renamed to Tests/SourceKitLSPTests/TestDiscoveryTests.swift

@@ -14,7 +14,7 @@ import LanguageServerProtocol
 import SKTestSupport
 import XCTest
 
-final class WorkspaceTestsTests: XCTestCase {
+final class TestDiscoveryTests: XCTestCase {
   func testWorkspaceTests() async throws {
     try XCTSkipIf(longTestsDisabled)
 
@@ -71,4 +71,69 @@ final class WorkspaceTestsTests: XCTestCase {
       ]
     )
   }
+
+  func testDocumentTests() async throws {
+    try XCTSkipIf(longTestsDisabled)
+
+    let ws = try await SwiftPMTestWorkspace(
+      files: [
+        "Tests/MyLibraryTests/MyTests.swift": """
+        import XCTest
+
+        class 1️⃣MyTests: XCTestCase {
+          func 2️⃣testMyLibrary() {}
+          func unrelatedFunc() {}
+          var testVariable: Int = 0
+        }
+        """,
+        "Tests/MyLibraryTests/MoreTests.swift": """
+        import XCTest
+
+        class MoreTests: XCTestCase {
+          func testSomeMore() {}
+        }
+        """
+      ],
+      manifest: """
+        // swift-tools-version: 5.7
+
+        import PackageDescription
+
+        let package = Package(
+          name: "MyLibrary",
+          targets: [.testTarget(name: "MyLibraryTests")]
+        )
+        """,
+      build: true
+    )
+
+    let (uri, positions) = try ws.openDocument("MyTests.swift")
+    let tests = try await ws.testClient.send(DocumentTestsRequest(textDocument: TextDocumentIdentifier(uri)))
+    XCTAssertEqual(
+      tests,
+      [
+        WorkspaceSymbolItem.symbolInformation(
+          SymbolInformation(
+            name: "MyTests",
+            kind: .class,
+            location: Location(
+              uri: uri,
+              range: Range(positions["1️⃣"])
+            )
+          )
+        ),
+        WorkspaceSymbolItem.symbolInformation(
+          SymbolInformation(
+            name: "testMyLibrary()",
+            kind: .method,
+            location: Location(
+              uri: try ws.uri(for: "MyTests.swift"),
+              range: Range(positions["2️⃣"])
+            ),
+            containerName: "MyTests"
+          )
+        ),
+      ]
+    )
+  }
 }