Add a request to get all tests within a file

Author: ahoppen

Sources/LanguageServerProtocol/CMakeLists.txt

@@ -50,6 +50,7 @@ add_library(LanguageServerProtocol STATIC
   Requests/DocumentSemanticTokensRangeRequest.swift
   Requests/DocumentSemanticTokensRequest.swift
   Requests/DocumentSymbolRequest.swift
+  Requests/DocumentTestsRequest.swift
   Requests/ExecuteCommandRequest.swift
   Requests/FoldingRangeRequest.swift
   Requests/FormattingRequests.swift

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 = "document/tests"
+  public typealias Response = [WorkspaceSymbolItem]?
+
+  public var textDocument: TextDocumentIdentifier
+
+  public init(textDocument: TextDocumentIdentifier) {
+    self.textDocument = textDocument
+  }
+}

Sources/SKCore/BuildSystemManager.swift

@@ -268,7 +268,7 @@ extension BuildSystemManager: MainFilesDelegate {
   /// For Swift or normal C files, this will be the file itself. For header
   /// files, we pick a main file that includes the header since header files
   /// don't have build settings by themselves.
-  private func mainFile(for uri: DocumentURI, language: Language, useCache: Bool = true) async -> DocumentURI {
+  public func mainFile(for uri: DocumentURI, language: Language, useCache: Bool = true) async -> DocumentURI {
     if language == .swift {
       // Swift doesn't have main files. Skip the main file provider query.
       return uri

Sources/SourceKitLSP/CMakeLists.txt

@@ -7,8 +7,9 @@ add_library(SourceKitLSP STATIC
   Sequence+AsyncMap.swift
   SourceKitIndexDelegate.swift
   SourceKitLSPCommandMetadata.swift
-  SourceKitServer+Options.swift
   SourceKitServer.swift
+  SourceKitServer+Options.swift
+  TestDiscovery.swift
   ToolchainLanguageServer.swift
   Workspace.swift
 )

Sources/SourceKitLSP/SourceKitServer.swift

@@ -383,7 +383,7 @@ public actor SourceKitServer {
 
   var languageServices: [LanguageServerType: [ToolchainLanguageServer]] = [:]
 
-  private let documentManager = DocumentManager()
+  let documentManager = DocumentManager()
 
   private var packageLoadingWorkDoneProgress = WorkDoneProgressState(
     "SourceKitLSP.SourceKitServer.reloadPackage",
@@ -399,7 +399,7 @@ public actor SourceKitServer {
   /// Must only be accessed from `queue`.
   private var uriToWorkspaceCache: [DocumentURI: WeakWorkspace] = [:]
 
-  private var workspaces: [Workspace] = [] {
+  private(set) var workspaces: [Workspace] = [] {
     didSet {
       uriToWorkspaceCache = [:]
     }
@@ -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>:
@@ -1498,16 +1500,6 @@ extension SourceKitServer {
     return symbols
   }
 
-  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 testSymbols.map(WorkspaceSymbolItem.init)
-  }
-
   /// Forwards a SymbolInfoRequest to the appropriate toolchain service for this document.
   func symbolInfo(
     _ req: SymbolInfoRequest,
@@ -2204,7 +2196,7 @@ fileprivate struct DocumentNotificationRequestQueue {
     queue = []
   }
 }
-fileprivate extension WorkspaceSymbolItem {
+extension WorkspaceSymbolItem {
   init(_ symbolOccurrence: SymbolOccurrence) {
     let symbolPosition = Position(
       line: symbolOccurrence.location.line - 1,  // 1-based -> 0-based

Sources/SourceKitLSP/TestDiscovery.swift

@@ -0,0 +1,61 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 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
+//
+//===----------------------------------------------------------------------===//
+
+import IndexStoreDB
+import LanguageServerProtocol
+
+fileprivate extension SymbolOccurrence {
+  /// Assuming that this is a symbol occurrence returned by the index, return whether it can constitute the definition
+  /// of a test case.
+  ///
+  /// The primary intention for this is to filter out references to test cases and extension declarations of test cases.
+  /// The latter is important to filter so we don't include extension declarations for the derived `DiscoveredTests`
+  /// files on non-Darwin platforms.
+  var canBeTestDefinition: Bool {
+    guard roles.contains(.definition) else {
+      return false
+    }
+    guard symbol.kind == .class || symbol.kind == .instanceMethod else {
+      return false
+    }
+    return true
+  }
+}
+
+extension SourceKitServer {
+  func workspaceTests(_ req: WorkspaceTestsRequest) async throws -> [WorkspaceSymbolItem]? {
+    let testSymbols = workspaces.flatMap { (workspace) -> [SymbolOccurrence] in
+      return workspace.index?.unitTests() ?? []
+    }
+    return
+      testSymbols
+      .filter { $0.canBeTestDefinition }
+      .map(WorkspaceSymbolItem.init)
+  }
+
+  func documentTests(
+    _ req: DocumentTestsRequest,
+    workspace: Workspace,
+    languageService: ToolchainLanguageServer
+  ) async throws -> [WorkspaceSymbolItem]? {
+    let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
+    let mainFileUri = await workspace.buildSystemManager.mainFile(
+      for: req.textDocument.uri,
+      language: snapshot.language
+    )
+    let testSymbols = workspace.index?.unitTests(referencedByMainFiles: [mainFileUri.pseudoPath]) ?? []
+    return
+      testSymbols
+      .filter { $0.canBeTestDefinition }
+      .map(WorkspaceSymbolItem.init)
+  }
+}

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"
+          )
+        ),
+      ]
+    )
+  }
 }