Merge tests defined in extensions

Merge the XCTests and swift-testing tests defined in extensions into
their parent TestItems.

This is done as another pass after the TestScanner visitors have walked
the tree.

Fixes #1218

Author: plemarquand

Sources/SourceKitLSP/Swift/SwiftTestingScanner.swift

@@ -201,7 +201,7 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
     let visitor = SyntacticSwiftTestingTestScanner(snapshot: snapshot, allTestsDisabled: false, parentTypeNames: [])
     visitor.walk(syntaxTree)
-    return visitor.result
+    return visitor.result.mergeTestsInExtensions()
   }
 
   /// Visit a class/struct/... or extension declaration.

Sources/SourceKitLSP/TestDiscovery.swift

@@ -331,7 +331,7 @@ final class SyntacticSwiftXCTestScanner: SyntaxVisitor {
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
     let visitor = SyntacticSwiftXCTestScanner(snapshot: snapshot)
     visitor.walk(syntaxTree)
-    return visitor.result
+    return visitor.result.mergeTestsInExtensions()
   }
 
   private func findTestMethods(in members: MemberBlockItemListSyntax, containerName: String) -> [TestItem] {
@@ -436,6 +436,45 @@ extension TestItem {
   }
 }
 
+extension Collection where Element == TestItem {
+  /// Walks the TestItem tree of each item in the collection and merges orphaned leaf nodes
+  /// into their parent, if a parent exists.
+  ///
+  /// A node's parent is identified by the node's ID with the last component dropped.
+  func mergeTestsInExtensions() -> [TestItem] {
+    var itemDict: [String: TestItem] = [:]
+    for item in self {
+      if var existingItem = itemDict[item.id] {
+        existingItem.children = (existingItem.children + item.children)
+        itemDict[item.id] = existingItem
+      } else {
+        itemDict[item.id] = item
+      }
+    }
+
+    for item in self {
+      let parentID = item.id.components(separatedBy: "/").dropLast().joined(separator: "/")
+      // If the parent exists, add the current item to its children and remove it from the root
+      if var parent = itemDict[parentID] {
+        parent.children.append(item)
+        itemDict[parent.id] = parent
+        itemDict[item.id] = nil
+      }
+    }
+
+    // Filter out the items that have been merged into their parents, sorting the tests by location
+    var reorganizedItems = itemDict.values.compactMap { $0 }.sorted { $0.location < $1.location }
+
+    reorganizedItems = reorganizedItems.map({
+      var newItem = $0
+      newItem.children = $0.children.mergeTestsInExtensions()
+      return newItem
+    })
+
+    return reorganizedItems
+  }
+}
+
 extension SwiftLanguageService {
   public func syntacticDocumentTests(for uri: DocumentURI, in workspace: Workspace) async throws -> [TestItem] {
     let snapshot = try documentManager.latestSnapshot(uri)

Tests/SourceKitLSPTests/DocumentTestDiscoveryTests.swift

@@ -827,29 +827,118 @@ final class DocumentTestDiscoveryTests: XCTestCase {
               location: Location(uri: uri, range: positions["2️⃣"]..<positions["3️⃣"]),
               children: [],
               tags: []
-            )
+            ),
+            TestItem(
+              id: "MyTests/twoIsThree()",
+              label: "twoIsThree()",
+              disabled: false,
+              style: TestStyle.swiftTesting,
+              location: Location(uri: uri, range: positions["6️⃣"]..<positions["7️⃣"]),
+              children: [],
+              tags: []
+            ),
           ],
           tags: []
-        ),
+        )
+      ]
+    )
+  }
+
+  func testSwiftTestingTestSuitesWithExtension() async throws {
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI.for(.swift)
+
+    let positions = testClient.openDocument(
+      """
+      import XCTest
+
+      1️⃣@Suite struct MyTests {
+        2️⃣@Test func oneIsTwo() {}3️⃣
+      }4️⃣
+
+      extension MyTests {
+        5️⃣@Test func twoIsThree() {}6️⃣
+      }
+      """,
+      uri: uri
+    )
+
+    let tests = try await testClient.send(DocumentTestsRequest(textDocument: TextDocumentIdentifier(uri)))
+    XCTAssertEqual(
+      tests,
+      [
         TestItem(
           id: "MyTests",
           label: "MyTests",
           disabled: false,
           style: TestStyle.swiftTesting,
-          location: Location(uri: uri, range: positions["5️⃣"]..<positions["8️⃣"]),
+          location: Location(uri: uri, range: positions["1️⃣"]..<positions["4️⃣"]),
           children: [
+            TestItem(
+              id: "MyTests/oneIsTwo()",
+              label: "oneIsTwo()",
+              disabled: false,
+              style: TestStyle.swiftTesting,
+              location: Location(uri: uri, range: positions["2️⃣"]..<positions["3️⃣"]),
+              children: [],
+              tags: []
+            ),
             TestItem(
               id: "MyTests/twoIsThree()",
               label: "twoIsThree()",
               disabled: false,
               style: TestStyle.swiftTesting,
-              location: Location(uri: uri, range: positions["6️⃣"]..<positions["7️⃣"]),
+              location: Location(uri: uri, range: positions["5️⃣"]..<positions["6️⃣"]),
+              children: [],
+              tags: []
+            ),
+          ],
+          tags: []
+        )
+      ]
+    )
+  }
+
+  func testXCTestTestsWithExtension() async throws {
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI.for(.swift)
+
+    let positions = testClient.openDocument(
+      """
+      import XCTest
+
+      1️⃣final class MyTests: XCTestCase {}4️⃣
+
+      extension MyTests {
+        2️⃣func testOneIsTwo() {}3️⃣
+      }
+      """,
+      uri: uri
+    )
+
+    let tests = try await testClient.send(DocumentTestsRequest(textDocument: TextDocumentIdentifier(uri)))
+    XCTAssertEqual(
+      tests,
+      [
+        TestItem(
+          id: "MyTests",
+          label: "MyTests",
+          disabled: false,
+          style: TestStyle.xcTest,
+          location: Location(uri: uri, range: positions["1️⃣"]..<positions["4️⃣"]),
+          children: [
+            TestItem(
+              id: "MyTests/testOneIsTwo()",
+              label: "testOneIsTwo()",
+              disabled: false,
+              style: TestStyle.xcTest,
+              location: Location(uri: uri, range: positions["2️⃣"]..<positions["3️⃣"]),
               children: [],
               tags: []
             )
           ],
           tags: []
-        ),
+        )
       ]
     )
   }