Merge pull request #1793 from plemarquand/document-tests-index-extension-fix

XCTests in extensions produced by symbols have no parent ID

Author: ahoppen

Sources/SemanticIndex/CheckedIndex.swift

@@ -161,6 +161,28 @@ package final class CheckedIndex {
   package func fileHasInMemoryModifications(_ uri: DocumentURI) -> Bool {
     return checker.fileHasInMemoryModifications(uri)
   }
+
+  /// If there are any definition occurrences of the given USR, return these.
+  /// Otherwise return declaration occurrences.
+  package func definitionOrDeclarationOccurrences(ofUSR usr: String) -> [SymbolOccurrence] {
+    let definitions = occurrences(ofUSR: usr, roles: [.definition])
+    if !definitions.isEmpty {
+      return definitions
+    }
+    return occurrences(ofUSR: usr, roles: [.declaration])
+  }
+
+  /// Find a `SymbolOccurrence` that is considered the primary definition of the symbol with the given USR.
+  ///
+  /// If the USR has an ambiguous definition, the most important role of this function is to deterministically return
+  /// the same result every time.
+  package func primaryDefinitionOrDeclarationOccurrence(ofUSR usr: String) -> SymbolOccurrence? {
+    let result = definitionOrDeclarationOccurrences(ofUSR: usr).sorted().first
+    if result == nil {
+      logger.error("Failed to find definition of \(usr) in index")
+    }
+    return result
+  }
 }
 
 /// A wrapper around `IndexStoreDB` that allows the retrieval of a `CheckedIndex` with a specified check level or the

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -2402,28 +2402,6 @@ private let maxWorkspaceSymbolResults = 4096
 package typealias Diagnostic = LanguageServerProtocol.Diagnostic
 
 fileprivate extension CheckedIndex {
-  /// If there are any definition occurrences of the given USR, return these.
-  /// Otherwise return declaration occurrences.
-  func definitionOrDeclarationOccurrences(ofUSR usr: String) -> [SymbolOccurrence] {
-    let definitions = occurrences(ofUSR: usr, roles: [.definition])
-    if !definitions.isEmpty {
-      return definitions
-    }
-    return occurrences(ofUSR: usr, roles: [.declaration])
-  }
-
-  /// Find a `SymbolOccurrence` that is considered the primary definition of the symbol with the given USR.
-  ///
-  /// If the USR has an ambiguous definition, the most important role of this function is to deterministically return
-  /// the same result every time.
-  func primaryDefinitionOrDeclarationOccurrence(ofUSR usr: String) -> SymbolOccurrence? {
-    let result = definitionOrDeclarationOccurrences(ofUSR: usr).sorted().first
-    if result == nil {
-      logger.error("Failed to find definition of \(usr) in index")
-    }
-    return result
-  }
-
   /// Get the name of the symbol that is a parent of this symbol, if one exists
   func containerName(of symbol: SymbolOccurrence) -> String? {
     // The container name of accessors is the container of the surrounding variable.

Sources/SourceKitLSP/TestDiscovery.swift

@@ -102,19 +102,34 @@ extension SourceKitLSPServer {
   /// provide ranges for the test cases in source code instead of only the test's location that we get from the index.
   private func testItems(
     for testSymbolOccurrences: [SymbolOccurrence],
+    index: CheckedIndex?,
     resolveLocation: (DocumentURI, Position) -> Location
   ) -> [AnnotatedTestItem] {
     // Arrange tests by the USR they are contained in. This allows us to emit test methods as children of test classes.
     // `occurrencesByParent[nil]` are the root test symbols that aren't a child of another test symbol.
     var occurrencesByParent: [String?: [SymbolOccurrence]] = [:]
 
-    let testSymbolUsrs = Set(testSymbolOccurrences.map(\.symbol.usr))
+    var testSymbolUsrs = Set(testSymbolOccurrences.map(\.symbol.usr))
+
+    // Gather any extension declarations that contains tests and add them to `occurrencesByParent` so we can properly
+    // arrange their test items as the extension's children.
+    for testSymbolOccurrence in testSymbolOccurrences {
+      for parentSymbol in testSymbolOccurrence.relations.filter({ $0.roles.contains(.childOf) }).map(\.symbol) {
+        guard parentSymbol.kind == .extension else {
+          continue
+        }
+        guard let definition = index?.primaryDefinitionOrDeclarationOccurrence(ofUSR: parentSymbol.usr) else {
+          logger.fault("Unable to find primary definition of extension '\(parentSymbol.usr)' containing tests")
+          continue
+        }
+        testSymbolUsrs.insert(parentSymbol.usr)
+        occurrencesByParent[nil, default: []].append(definition)
+      }
+    }
 
     for testSymbolOccurrence in testSymbolOccurrences {
       let childOfUsrs = testSymbolOccurrence.relations
-        .filter { $0.roles.contains(.childOf) }
-        .map(\.symbol.usr)
-        .filter { testSymbolUsrs.contains($0) }
+        .filter { $0.roles.contains(.childOf) }.map(\.symbol.usr).filter { testSymbolUsrs.contains($0) }
       if childOfUsrs.count > 1 {
         logger.fault(
           "Test symbol \(testSymbolOccurrence.symbol.usr) is child or multiple symbols: \(childOfUsrs.joined(separator: ", "))"
@@ -168,7 +183,7 @@ extension SourceKitLSPServer {
           children: children.map(\.testItem),
           tags: []
         ),
-        isExtension: false
+        isExtension: testSymbolOccurrence.symbol.kind == .extension
       )
     }
 
@@ -233,6 +248,7 @@ extension SourceKitLSPServer {
     let testsFromSyntacticIndex = await workspace.syntacticTestIndex.tests()
     let testsFromSemanticIndex = testItems(
       for: semanticTestSymbolOccurrences,
+      index: index,
       resolveLocation: { uri, position in Location(uri: uri, range: Range(position)) }
     )
     let filesWithTestsFromSemanticIndex = Set(testsFromSemanticIndex.map(\.testItem.location.uri))
@@ -337,6 +353,7 @@ extension SourceKitLSPServer {
         // swift-testing tests, which aren't part of the semantic index.
         return testItems(
           for: testSymbols,
+          index: index,
           resolveLocation: { uri, position in
             if uri == snapshot.uri, let documentSymbols,
               let range = findInnermostSymbolRange(containing: position, documentSymbolsResponse: documentSymbols)

Tests/SourceKitLSPTests/DocumentTestDiscoveryTests.swift

@@ -1340,6 +1340,92 @@ final class DocumentTestDiscoveryTests: XCTestCase {
     )
   }
 
+  func testXCTestIndexedTestsWithExtension() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      import XCTest
+
+      1️⃣final class MyTests: XCTestCase {}2️⃣
+
+      extension MyTests {
+        3️⃣func testOneIsTwo() {}4️⃣
+      }
+      """,
+      allowBuildFailure: true
+    )
+
+    let tests = try await project.testClient.send(
+      DocumentTestsRequest(textDocument: TextDocumentIdentifier(project.fileURI))
+    )
+    XCTAssertEqual(
+      tests,
+      [
+        TestItem(
+          id: "MyTests",
+          label: "MyTests",
+          location: Location(uri: project.fileURI, range: project.positions["1️⃣"]..<project.positions["2️⃣"]),
+          children: [
+            TestItem(
+              id: "MyTests/testOneIsTwo()",
+              label: "testOneIsTwo()",
+              location: Location(uri: project.fileURI, range: project.positions["3️⃣"]..<project.positions["4️⃣"])
+            )
+          ]
+        )
+      ]
+    )
+  }
+
+  func testXCTestIndexedTestsWithExtensionInSeparateFile() async throws {
+    let project = try await SwiftPMTestProject(
+      files: [
+        "Tests/MyLibraryTests/MyTests.swift": """
+        import XCTest
+
+        class MyTests: XCTestCase {
+        }
+        """,
+        "Tests/MyLibraryTests/MoreTests.swift": """
+        import XCTest
+
+        1️⃣extension MyTests {
+          3️⃣func testMe() {}4️⃣
+        }2️⃣
+        """,
+      ],
+      manifest: """
+        let package = Package(
+          name: "MyLibrary",
+          targets: [.testTarget(name: "MyLibraryTests")]
+        )
+        """,
+      enableBackgroundIndexing: true
+    )
+
+    let (uri, positions) = try project.openDocument("MoreTests.swift")
+
+    let tests = try await project.testClient.send(
+      DocumentTestsRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+    XCTAssertEqual(
+      tests,
+      [
+        TestItem(
+          id: "MyLibraryTests.MyTests",
+          label: "MyTests",
+          location: Location(uri: uri, range: positions["1️⃣"]..<positions["2️⃣"]),
+          children: [
+            TestItem(
+              id: "MyLibraryTests.MyTests/testMe()",
+              label: "testMe()",
+              location: Location(uri: uri, range: positions["3️⃣"]..<positions["4️⃣"])
+            )
+          ]
+        )
+      ]
+    )
+  }
+
   func testXCTestInvalidXCTestSuiteConstructions() async throws {
     let testClient = try await TestSourceKitLSPClient()
     let uri = DocumentURI(for: .swift)