Address feedback, properly handle duplicate IDs in extensions

Author: plemarquand

Sources/SourceKitLSP/TestDiscovery.swift

@@ -433,7 +433,6 @@ fileprivate extension Array<AnnotatedTestItem> {
   /// A node's parent is identified by the node's ID with the last component dropped.
   func mergingTestsInExtensions() -> [TestItem] {
     var itemDict: [String: AnnotatedTestItem] = [:]
-    var duplicatedIds: Set<String> = []
     for item in self {
       let id = item.testItem.id
       if var rootItem = itemDict[id] {
@@ -442,14 +441,13 @@ fileprivate extension Array<AnnotatedTestItem> {
         // as the root item.
         if rootItem.isExtension && !item.isExtension {
           var newItem = item
-          newItem.testItem.children += rootItem.testItem.children
+          newItem.testItem.children = (newItem.testItem.children + rootItem.testItem.children).deduplicateIds()
           rootItem = newItem
-        } else if (rootItem.testItem.children.isEmpty && item.testItem.children.isEmpty) {
-          duplicatedIds.insert(item.testItem.id)
-          itemDict[item.testItem.fullyQualifiedTestId] = item
+        } else if rootItem.testItem.children.isEmpty && item.testItem.children.isEmpty {
+          itemDict[item.testItem.ambiguousTestDifferentiator] = item
           continue
         } else {
-          rootItem.testItem.children += item.testItem.children
+          rootItem.testItem.children = (rootItem.testItem.children + item.testItem.children).deduplicateIds()
         }
 
         itemDict[id] = rootItem
@@ -480,32 +478,16 @@ fileprivate extension Array<AnnotatedTestItem> {
       .sorted { ($0.isExtension != $1.isExtension) ? !$0.isExtension : ($0.testItem.location < $1.testItem.location) }
 
     let result = sortedItems.map {
-      var newItem = $0.testItem
-
-      // If multiple testItems share the same ID we add more context to make it unique.
-      // Two tests can share the same ID when two swift testing tests accept
-      // arguments of different types, i.e:
-      // @Test(arguments: [1,2,3]) func foo(_ x: Int) {}
-      // @Test(arguments: ["a", "b", "c"]) func foo(_ x: String) {}
-      // or when tests are in separate files but don't conflict because they are marked
-      // private, i.e:
-      // File1.swift: @Test private func foo() {}
-      // File2.swift: @Test private func foo() {}
-      // If we encounter one of these cases, we need to deduplicate the ID
-      // by appending /filename:filename:lineNumber.
-      if duplicatedIds.contains(newItem.id) {
-        newItem.id = newItem.fullyQualifiedTestId
-      }
-
       guard !$0.testItem.children.isEmpty, mergedIds.contains($0.testItem.id) else {
-        return newItem
+        return $0.testItem
       }
+      var newItem = $0.testItem
       newItem.children = newItem.children
         .map { AnnotatedTestItem(testItem: $0, isExtension: false) }
         .mergingTestsInExtensions()
       return newItem
     }
-    return result
+    return result.deduplicateIds()
   }
 
   func prefixTestsWithModuleName(workspace: Workspace) async -> Self {
@@ -518,19 +500,58 @@ fileprivate extension Array<AnnotatedTestItem> {
   }
 }
 
-extension TestItem {
-  fileprivate var fullyQualifiedTestId: String {
-    return "\(self.id)/\(self.sourceLocation)"
+fileprivate extension Array<TestItem> {
+  /// If multiple testItems share the same ID we add more context to make it unique.
+  /// Two tests can share the same ID when two swift testing tests accept
+  /// arguments of different types, i.e:
+  /// ```
+  /// @Test(arguments: [1,2,3]) func foo(_ x: Int) {}
+  /// @Test(arguments: ["a", "b", "c"]) func foo(_ x: String) {}
+  /// ```
+  ///
+  /// or when tests are in separate files but don't conflict because they are marked
+  /// private, i.e:
+  /// ```
+  /// File1.swift: @Test private func foo() {}
+  /// File2.swift: @Test private func foo() {}
+  /// ```
+  ///
+  /// If we encounter one of these cases, we need to deduplicate the ID
+  /// by appending `/filename:filename:lineNumber`.
+  func deduplicateIds() -> [TestItem] {
+    var idCounts: [String: Int] = [:]
+    var result: [TestItem] = []
+
+    for element in self where element.children.isEmpty {
+      idCounts[element.id, default: 0] += 1
+    }
+
+    for element in self {
+      if idCounts[element.id] ?? 0 > 1 {
+        var newItem = element
+        newItem.id = newItem.ambiguousTestDifferentiator
+        result.append(newItem)
+      } else {
+        result.append(element)
+      }
+    }
+
+    return result
   }
+}
 
-  private var sourceLocation: String {
+extension TestItem {
+  /// A fully qualified name to disambiguate identical TestItem IDs.
+  /// This matches the IDs produced by `swift test list` when there are
+  /// tests that cannot be disambiguated by their simple ID.
+  fileprivate var ambiguousTestDifferentiator: String {
     let filename = self.location.uri.arbitrarySchemeURL.lastPathComponent
     let position = location.range.lowerBound
     // Lines and columns start at 1.
     // swift-testing tests start from _after_ the @ symbol in @Test, so we need to add an extra column.
     // see https://github.com/swiftlang/swift-testing/blob/cca6de2be617aded98ecdecb0b3b3a81eec013f3/Sources/TestingMacros/Support/AttributeDiscovery.swift#L153
     let columnOffset = self.style == TestStyle.swiftTesting ? 2 : 1
-    return "\(filename):\(position.line + 1):\(position.utf16index + columnOffset)"
+    return "\(self.id)/\(filename):\(position.line + 1):\(position.utf16index + columnOffset)"
   }
 
   fileprivate func prefixIDWithModuleName(workspace: Workspace) async -> TestItem {

Tests/SourceKitLSPTests/DocumentTestDiscoveryTests.swift

@@ -455,6 +455,55 @@ final class DocumentTestDiscoveryTests: XCTestCase {
     )
   }
 
+  func testSwiftTestingTestsWithDuplicateFunctionIdentifiersInExtension() async throws {
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI(for: .swift)
+
+    let positions = testClient.openDocument(
+      """
+      import Testing
+      1️⃣struct MySuite {
+        3️⃣@Test(arguments: [1, 2, 3])
+        func foo(_ x: Int) {}4️⃣
+      }2️⃣
+
+      extension MySuite {
+        5️⃣@Test(arguments: ["a", "b", "c"])
+        func foo(_ x: String) {}6️⃣
+      }
+      """,
+      uri: uri
+    )
+
+    let filename = uri.fileURL?.lastPathComponent ?? ""
+    let tests = try await testClient.send(DocumentTestsRequest(textDocument: TextDocumentIdentifier(uri)))
+    XCTAssertEqual(
+      tests,
+      [
+        TestItem(
+          id: "MySuite",
+          label: "MySuite",
+          style: TestStyle.swiftTesting,
+          location: Location(uri: uri, range: positions["1️⃣"]..<positions["2️⃣"]),
+          children: [
+            TestItem(
+              id: "MySuite/foo(_:)/\(filename):\(positions["3️⃣"].line + 1):\(positions["3️⃣"].utf16index + 2)",
+              label: "foo(_:)",
+              style: TestStyle.swiftTesting,
+              location: Location(uri: uri, range: positions["3️⃣"]..<positions["4️⃣"])
+            ),
+            TestItem(
+              id: "MySuite/foo(_:)/\(filename):\(positions["5️⃣"].line + 1):\(positions["5️⃣"].utf16index + 2)",
+              label: "foo(_:)",
+              style: TestStyle.swiftTesting,
+              location: Location(uri: uri, range: positions["5️⃣"]..<positions["6️⃣"])
+            ),
+          ]
+        )
+      ]
+    )
+  }
+
   func testSwiftTestingSuiteWithNoTests() async throws {
     let testClient = try await TestSourceKitLSPClient()
     let uri = DocumentURI(for: .swift)