fix a bug in test discovery where tests define in extensions break the discovery

motivation: less bugs, happier users

changes: modify test discovery code to key of class name rather than blindly iterate over tetst

rdar://58435666

Author: tomerd

Fixtures/Miscellaneous/TestDiscovery/Extensions/Package.swift

@@ -0,0 +1,10 @@
+// swift-tools-version:4.2
+import PackageDescription
+
+let package = Package(
+    name: "Simple",
+    targets: [
+        .target(name: "Simple"),
+        .testTarget(name: "SimpleTests", dependencies: ["Simple"]),
+    ]
+)

Fixtures/Miscellaneous/TestDiscovery/Extensions/Sources/Simple/Simple.swift

@@ -0,0 +1,3 @@
+struct Simple {
+
+}

Fixtures/Miscellaneous/TestDiscovery/Extensions/Tests/SimpleTests/SwiftTests1.swift

@@ -0,0 +1,7 @@
+import XCTest
+@testable import Simple
+
+class SimpleTests1: XCTestCase {
+    func testExample1() {
+    }
+}

Fixtures/Miscellaneous/TestDiscovery/Extensions/Tests/SimpleTests/SwiftTests2.swift

@@ -0,0 +1,7 @@
+import XCTest
+@testable import Simple
+
+class SimpleTests2: XCTestCase {
+    func testExample2() {
+    }
+}

Fixtures/Miscellaneous/TestDiscovery/Extensions/Tests/SimpleTests/SwiftTests3.swift

@@ -0,0 +1,12 @@
+import XCTest
+@testable import Simple
+
+extension SimpleTests1 {
+    func testExample1_a() {
+    }
+}
+
+extension SimpleTests2 {
+    func testExample2_a() {
+    }
+}

Fixtures/Miscellaneous/TestDiscovery/Extensions/Tests/SimpleTests/SwiftTests4.swift

@@ -0,0 +1,13 @@
+import XCTest
+@testable import Simple
+
+class SimpleTests4: XCTestCase {
+    func testExample() {
+    }
+
+    func testExample1() {
+    }
+
+    func testExample2() {
+    }
+}

Sources/Build/BuildOperationBuildSystemDelegateHandler.swift

@@ -54,14 +54,18 @@ final class TestDiscoveryCommand: CustomLLBuildCommand {
     ) throws {
         let stream = try LocalFileOutputByteStream(path)
 
+        let testsByClass = Dictionary(grouping: tests, by: { $0.name })
+
         stream <<< "import XCTest" <<< "\n"
         stream <<< "@testable import " <<< module <<< "\n"
 
-        for klass in tests {
+        for classTests in testsByClass {
+            let className = classTests.key
+            let testMethods = classTests.value.flatMap{ $0.methods }
             stream <<< "\n"
-            stream <<< "fileprivate extension " <<< klass.name <<< " {" <<< "\n"
-            stream <<< indent(4) <<< "static let __allTests__\(klass.name) = [" <<< "\n"
-            for method in klass.methods {
+            stream <<< "fileprivate extension " <<< className <<< " {" <<< "\n"
+            stream <<< indent(4) <<< "static let __allTests__\(className) = [" <<< "\n"
+            for method in testMethods {
                 let method = method.hasSuffix("()") ? String(method.dropLast(2)) : method
                 stream <<< indent(8) <<< "(\"\(method)\", \(method))," <<< "\n"
             }
@@ -74,8 +78,8 @@ final class TestDiscoveryCommand: CustomLLBuildCommand {
             return [\n
         """
 
-        for klass in tests {
-            stream <<< indent(8) <<< "testCase(\(klass.name).__allTests__\(klass.name)),\n"
+        for className in testsByClass.keys {
+            stream <<< indent(8) <<< "testCase(\(className).__allTests__\(className)),\n"
         }
 
         stream <<< """

Tests/FunctionalTests/TestDiscoverTests.swift renamed to Tests/FunctionalTests/TestDiscoveryTests.swift

@@ -83,4 +83,23 @@ class TestDiscoveryTests: XCTestCase {
         }
         #endif
     }
+
+    func testTestExtensions() throws {
+        #if os(macOS)
+        try XCTSkipIf(true)
+        #else
+        fixture(name: "Miscellaneous/TestDiscovery/Extensions") { path in
+            let (stdout, _) = try executeSwiftTest(path, extraArgs: ["--enable-test-discovery"])
+            XCTAssertTrue(stdout.contains("Merging module Simple"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests1.testExample1"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests1.testExample1_a"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests2.testExample2"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests2.testExample2_a"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests4.testExample"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests4.testExample1"), stdout)
+            XCTAssertTrue(stdout.contains("SimpleTests4.testExample2"), stdout)
+            XCTAssertTrue(stdout.contains("Executed 7 tests"), stdout)
+        }
+        #endif
+    }
 }