Keep track of which test infos use which target

Modify build script to keep track of which test info uses which target,
which will be necessary for intra-zip test discovery later

Test: test
Bug: 358215235
Change-Id: I3d51d6527396fc1538d51602ad84084dbcdc7b0e

1 files changed, 7 insertions, 1 deletions

diff --git a/ci/build_test_suites.py b/ci/build_test_suites.py
index d81248b496..7636f6a44b 100644
--- a/ci/build_test_suites.py
+++ b/ci/build_test_suites.py
@@ -16,6 +16,7 @@
 
 import argparse
 from dataclasses import dataclass
+from collections import defaultdict
 import json
 import logging
 import os
@@ -68,6 +69,7 @@ class BuildPlanner:
     self.build_context = build_context
     self.args = args
     self.target_optimizations = target_optimizations
+    self.target_to_test_infos = defaultdict(list)
 
   def create_build_plan(self):
 
@@ -102,7 +104,7 @@ class BuildPlanner:
         continue
 
       target_optimizer = target_optimizer_getter(
-          target, self.build_context, self.args
+          target, self.build_context, self.args, self.target_to_test_infos[target]
       )
       build_targets.update(target_optimizer.get_build_targets())
       packaging_commands_getters.append(
@@ -178,6 +180,10 @@ class BuildPlanner:
       tf_command = self._build_tf_command(test_info)
       discovery_agent = test_discovery_agent.TestDiscoveryAgent(tradefed_args=tf_command)
       for regex in discovery_agent.discover_test_zip_regexes():
+        for target in self.args.extra_targets:
+          target_regex = r'\b(%s.*)\b' % re.escape(target)
+          if re.search(target_regex, regex):
+            self.target_to_test_infos[target].append(test_info)
         build_target_regexes.add(regex)
     return build_target_regexes