From c03fceee4108694e0dec44b71148eb8c4864be9e Mon Sep 17 00:00:00 2001
From: Bruno Henrique da Rocha e Silva <rochab@spotify.com>
Date: Thu, 19 Feb 2026 09:31:45 +0100
Subject: [PATCH] Improve parsing of test files

---
 .../BazelTargetQuerierParser.swift            | 18 ++++-----
 .../BazelTargetQuerierParserImplTests.swift   | 37 +++++++++++++++++++
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/Sources/SourceKitBazelBSP/RequestHandlers/BuildTargets/BazelTargetQuerierParser.swift b/Sources/SourceKitBazelBSP/RequestHandlers/BuildTargets/BazelTargetQuerierParser.swift
index 2d684b9..d6c4bb7 100644
--- a/Sources/SourceKitBazelBSP/RequestHandlers/BuildTargets/BazelTargetQuerierParser.swift
+++ b/Sources/SourceKitBazelBSP/RequestHandlers/BuildTargets/BazelTargetQuerierParser.swift
@@ -124,7 +124,7 @@ final class BazelTargetQuerierParserImpl: BazelTargetQuerierParser {
         var configurationToTopLevelLabelsMap: [String: [String]] = [:]
         var allTopLevelLabels = [(String, TopLevelRuleType)]()
         var allAliases = [BlazeQuery_Target]()
-        var allTestBundles: [(target: BlazeQuery_Target, isUITest: Bool)] = []
+        var allTestBundles: [BlazeQuery_Target] = []
         var unfilteredDependencyTargets = [Analysis_ConfiguredTarget]()
         var seenSourceFiles = Set<String>()
         var allSrcs = [BlazeQuery_Target]()
@@ -163,7 +163,7 @@ final class BazelTargetQuerierParserImpl: BazelTargetQuerierParser {
                     guard target.rule.name.hasSuffix(TopLevelRuleType.testBundleRuleSuffix) else {
                         throw BazelTargetQuerierParserError.unexpectedTestBundleRuleWithoutSuffix(target.rule.name)
                     }
-                    allTestBundles.append((target, kind.hasSuffix("_ui_test_bundle")))
+                    allTestBundles.append(target)
                     let realTopLevelName = String(
                         target.rule.name.dropLast(TopLevelRuleType.testBundleRuleSuffix.count)
                     )
@@ -269,15 +269,13 @@ final class BazelTargetQuerierParserImpl: BazelTargetQuerierParser {
 
         // Treat test bundle rules as aliases as well. This allows us to locate the "true" dependency
         // when encountering a test bundle.
-        for (target, isUITest) in allTestBundles {
+        for (target) in allTestBundles {
             let label = target.rule.name
-            let actual: String
-            if isUITest {
-                // For UI tests, the first input is the test_host.
-                // So in this case we're looking for the second input.
-                actual = target.rule.ruleInput[1]
-            } else {
-                actual = target.rule.ruleInput[0]
+            let deps = target.rule.attribute.first { $0.name == "deps" }?.stringListValue ?? []
+            // The first dependency of the generated test bundle is what we're looking for.
+            guard let actual = deps.first else {
+                logger.error("Unexpected missing dependency for test bundle \(label).")
+                continue
             }
             aliasToLabelMap[label] = actual
             registeredAliases.append(label)
diff --git a/Tests/SourceKitBazelBSPTests/BazelTargetQuerierParserImplTests.swift b/Tests/SourceKitBazelBSPTests/BazelTargetQuerierParserImplTests.swift
index a115964..1ce1590 100644
--- a/Tests/SourceKitBazelBSPTests/BazelTargetQuerierParserImplTests.swift
+++ b/Tests/SourceKitBazelBSPTests/BazelTargetQuerierParserImplTests.swift
@@ -309,6 +309,43 @@ struct BazelTargetQuerierParserImplTests {
                     "//HelloWorld:HelloWorldWatchTests",
                 ])
         )
+
+        // Verify bazelLabelToTestFilesMap - maps test targets to their source files
+        // These are the test bundle targets that should have source files mapped
+        let expectedTestTargets = Set([
+            "//HelloWorld:HelloWorldTests",
+            "//HelloWorld:HelloWorldE2ETests",
+            "//HelloWorld:HelloWorldMacTests",
+            "//HelloWorld:HelloWorldWatchTests",
+        ])
+        let actualTestTargets = Set(result.bazelLabelToTestFilesMap.keys)
+        #expect(
+            actualTestTargets == expectedTestTargets,
+            "bazelLabelToTestFilesMap keys don't match expected test targets"
+        )
+
+        // Verify each test target has non-empty source files
+        for testTarget in expectedTestTargets {
+            let sourceFiles = try #require(
+                result.bazelLabelToTestFilesMap[testTarget],
+                "Missing source files for test target: \(testTarget)"
+            )
+            #expect(
+                !sourceFiles.isEmpty,
+                "Source files should not be empty for test target: \(testTarget)"
+            )
+        }
+
+        // Verify specific source file mappings
+        // HelloWorldTests should have source files from the HelloWorldTests directory
+        // (HelloWorldTestsLib library sources are in HelloWorldTests/*.swift)
+        let unitTestFiles = try #require(result.bazelLabelToTestFilesMap["//HelloWorld:HelloWorldTests"])
+        #expect(unitTestFiles.contains { $0.stringValue.contains("HelloWorldTests/") })
+
+        // HelloWorldE2ETests should have source files from the HelloWorldE2ETests directory
+        // (HelloWorldE2ETestsLib library sources are in HelloWorldE2ETests/*.swift)
+        let e2eTestFiles = try #require(result.bazelLabelToTestFilesMap["//HelloWorld:HelloWorldE2ETests"])
+        #expect(e2eTestFiles.contains { $0.stringValue.contains("HelloWorldE2ETests/") })
     }
 
     @Test