Prioritize tests defined in type definition over those in extensions

Sort the list of test items prioritizing those defined in the
originating type definition over those in extensions.

Author: plemarquand

Sources/LanguageServerProtocol/SupportTypes/TestItem.swift

@@ -56,12 +56,16 @@ public struct TestItem: ResponseType, Equatable {
   /// Tags associated with this test item.
   public var tags: [TestTag]
 
+  /// Whether the `TestItem` is declared in an extension.
+  public var isExtension: Bool
+
   public init(
     id: String,
     label: String,
     description: String? = nil,
     sortText: String? = nil,
     disabled: Bool,
+    isExtension: Bool,
     style: String,
     location: Location,
     children: [TestItem],
@@ -72,6 +76,7 @@ public struct TestItem: ResponseType, Equatable {
     self.description = description
     self.sortText = sortText
     self.disabled = disabled
+    self.isExtension = isExtension
     self.style = style
     self.location = location
     self.children = children

Sources/SourceKitLSP/Swift/SwiftTestingScanner.swift

@@ -163,6 +163,9 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
   /// This is the case when the scanner is looking for tests inside a disabled suite.
   private let allTestsDisabled: Bool
 
+  /// Whether the tests discovered by the scanner should be marked as being delcared in an extension.
+  private let isScanningExtension: Bool
+
   /// The names of the types that this scanner is scanning members for.
   ///
   /// For example, when scanning for tests inside `Bar` in the following, this is `["Foo", "Bar"]`
@@ -179,10 +182,16 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
   /// The discovered test items.
   private var result: [TestItem] = []
 
-  private init(snapshot: DocumentSnapshot, allTestsDisabled: Bool, parentTypeNames: [String]) {
+  private init(
+    snapshot: DocumentSnapshot,
+    allTestsDisabled: Bool,
+    isScanningExtension: Bool,
+    parentTypeNames: [String]
+  ) {
     self.snapshot = snapshot
     self.allTestsDisabled = allTestsDisabled
     self.parentTypeNames = parentTypeNames
+    self.isScanningExtension = isScanningExtension
     super.init(viewMode: .fixedUp)
   }
 
@@ -199,9 +208,14 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
       return []
     }
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
-    let visitor = SyntacticSwiftTestingTestScanner(snapshot: snapshot, allTestsDisabled: false, parentTypeNames: [])
+    let visitor = SyntacticSwiftTestingTestScanner(
+      snapshot: snapshot,
+      allTestsDisabled: false,
+      isScanningExtension: false,
+      parentTypeNames: []
+    )
     visitor.walk(syntaxTree)
-    return visitor.result.mergeTestsInExtensions()
+    return visitor.result
   }
 
   /// Visit a class/struct/... or extension declaration.
@@ -231,6 +245,7 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
     let memberScanner = SyntacticSwiftTestingTestScanner(
       snapshot: snapshot,
       allTestsDisabled: attributeData?.isDisabled ?? false,
+      isScanningExtension: node is ExtensionDeclSyntax,
       parentTypeNames: parentTypeNames + typeNames
     )
     memberScanner.walk(node.memberBlock)
@@ -245,6 +260,7 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
       id: (parentTypeNames + typeNames).joined(separator: "/"),
       label: attributeData?.displayName ?? typeNames.last!,
       disabled: (attributeData?.isDisabled ?? false) || allTestsDisabled,
+      isExtension: node is ExtensionDeclSyntax,
       style: TestStyle.swiftTesting,
       location: Location(uri: snapshot.uri, range: range),
       children: memberScanner.result,
@@ -299,6 +315,7 @@ final class SyntacticSwiftTestingTestScanner: SyntaxVisitor {
       id: (parentTypeNames + [name]).joined(separator: "/"),
       label: attributeData.displayName ?? name,
       disabled: attributeData.isDisabled || allTestsDisabled,
+      isExtension: isScanningExtension,
       style: TestStyle.swiftTesting,
       location: Location(uri: snapshot.uri, range: range),
       children: [],

Sources/SourceKitLSP/TestDiscovery.swift

@@ -134,6 +134,7 @@ extension SourceKitLSPServer {
         id: id,
         label: testSymbolOccurrence.symbol.name,
         disabled: false,
+        isExtension: false,
         style: TestStyle.xcTest,
         location: location,
         children: children,
@@ -177,7 +178,7 @@ extension SourceKitLSPServer {
         return []
       }
       return await orLog("Getting document tests for \(uri)") {
-        try await self.documentTests(
+        try await self.getDocumentTests(
           DocumentTestsRequest(textDocument: TextDocumentIdentifier(uri)),
           workspace: workspace,
           languageService: languageService
@@ -238,6 +239,7 @@ extension SourceKitLSPServer {
       .concurrentMap { await self.tests(in: $0) }
       .flatMap { $0 }
       .sorted { $0.location < $1.location }
+      .mergingTestsInExtensions()
   }
 
   /// Extracts a flat dictionary mapping test IDs to their locations from the given `testItems`.
@@ -254,6 +256,15 @@ extension SourceKitLSPServer {
     _ req: DocumentTestsRequest,
     workspace: Workspace,
     languageService: LanguageService
+  ) async throws -> [TestItem] {
+    return try await getDocumentTests(req, workspace: workspace, languageService: languageService)
+      .mergingTestsInExtensions()
+  }
+
+  private func getDocumentTests(
+    _ req: DocumentTestsRequest,
+    workspace: Workspace,
+    languageService: LanguageService
   ) async throws -> [TestItem] {
     let snapshot = try self.documentManager.latestSnapshot(req.textDocument.uri)
     let mainFileUri = await workspace.buildSystemManager.mainFile(
@@ -331,7 +342,7 @@ final class SyntacticSwiftXCTestScanner: SyntaxVisitor {
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
     let visitor = SyntacticSwiftXCTestScanner(snapshot: snapshot)
     visitor.walk(syntaxTree)
-    return visitor.result.mergeTestsInExtensions()
+    return visitor.result
   }
 
   private func findTestMethods(in members: MemberBlockItemListSyntax, containerName: String) -> [TestItem] {
@@ -361,6 +372,7 @@ final class SyntacticSwiftXCTestScanner: SyntaxVisitor {
         id: "\(containerName)/\(function.name.text)()",
         label: "\(function.name.text)()",
         disabled: false,
+        isExtension: false,
         style: TestStyle.xcTest,
         location: Location(uri: snapshot.uri, range: range),
         children: [],
@@ -392,6 +404,7 @@ final class SyntacticSwiftXCTestScanner: SyntaxVisitor {
       id: node.name.text,
       label: node.name.text,
       disabled: false,
+      isExtension: false,
       style: TestStyle.xcTest,
       location: Location(uri: snapshot.uri, range: range),
       children: testMethods,
@@ -436,7 +449,7 @@ extension TestItem {
   }
 }
 
-extension Collection where Element == TestItem {
+extension Array<TestItem> {
   /// When the test scanners discover tests in extensions they are captured in their own parent `TestItem`, not the
   /// `TestItem` generated from the class/struct's definition. This is largely because of the syntatic nature of the
   /// test scanners as they are today, which only know about tests within the context of the current file. Extensions
@@ -451,18 +464,61 @@ extension Collection where Element == TestItem {
   /// This method walks the `TestItem` tree produced by the test scanners and merges in the tests defined in extensions
   /// into the `TestItem` that represents the type definition.
   ///
+  /// This causes extensions to be merged into their type's definition if the type's definition exists in the list of
+  /// test items. If the type's definition is not a test item in this collection, the first extension of that type will
+  /// be used as the primary test location.
+  ///
+  /// For example if there are two files
+  ///
+  /// FileA.swift
+  /// ```swift
+  /// @Suite struct MyTests {
+  ///   @Test func oneIsTwo {}
+  /// }
+  /// ```
+  ///
+  /// FileB.swift
+  /// ```swift
+  /// extension MyTests {
+  ///   @Test func twoIsThree() {}
+  /// }
+  /// ```
+  ///
+  /// Then `workspace/tests` will return
+  /// - `MyTests` (FileA.swift:1)
+  ///   - `oneIsTwo`
+  ///   - `twoIsThree`
+  ///
+  /// And `textDocument/tests` for FileB.swift will return
+  /// - `MyTests` (FileB.swift:1)
+  ///   - `twoIsThree`
+  ///
   /// A node's parent is identified by the node's ID with the last component dropped.
-  func mergeTestsInExtensions() -> [TestItem] {
+  func mergingTestsInExtensions() -> [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
+      if var rootItem = itemDict[item.id] {
+        // If we've encountered an extension first, and this is the
+        // type declaration, then use the type declaration TestItem
+        // as the root item.
+        if rootItem.isExtension && !item.isExtension {
+          var newItem = item
+          newItem.children += rootItem.children
+          rootItem = newItem
+        } else {
+          rootItem.children += item.children
+        }
+
+        itemDict[item.id] = rootItem
       } else {
         itemDict[item.id] = item
       }
     }
 
+    if itemDict.isEmpty {
+      return []
+    }
+
     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
@@ -473,16 +529,22 @@ extension Collection where Element == TestItem {
       }
     }
 
-    // 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 }
+    // Filter out the items that have been merged into their parents, sorting the tests by location.
+    // TestItems not in extensions should be priotitized first.
+    var sortedItems = itemDict.values.compactMap { $0 }.sorted {
+      $0.location.uri != $1.location.uri ? !$0.isExtension : ($0.location < $1.location)
+    }
 
-    reorganizedItems = reorganizedItems.map({
+    sortedItems = sortedItems.map {
+      guard !$0.children.isEmpty else {
+        return $0
+      }
       var newItem = $0
-      newItem.children = $0.children.mergeTestsInExtensions()
+      newItem.children = $0.children.mergingTestsInExtensions()
       return newItem
-    })
+    }
 
-    return reorganizedItems
+    return sortedItems
   }
 }