XCTests in extensions have no parent ID

If the workspace has been indexed then TestItems are pulled from
indexed symbols. In this code path XCTests defined in extensions don't
have their parent class name appended to their ID.

As a result, when the user freshly opened a test file that contained
XCTests defined in extensions then their IDs would be incorrect.
However as soon as they made a change to the file then the method to
produce the document tests would switch over to getting the tests from
`languageService.syntatciDocumentTests(for:in:)` which does handle
tests in extensions correctly, and the issue would dissapear.

Author: plemarquand

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)