diff --git a/ci/build_test_suites.py b/ci/build_test_suites.py
index d81248b..7636f6a 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 @@
     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 @@
         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 @@
       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
 
diff --git a/ci/build_test_suites_test.py b/ci/build_test_suites_test.py
index 190740f..e4501d3 100644
--- a/ci/build_test_suites_test.py
+++ b/ci/build_test_suites_test.py
@@ -257,9 +257,9 @@
   class TestOptimizedBuildTarget(optimized_targets.OptimizedBuildTarget):
 
     def __init__(
-        self, target, build_context, args, output_targets, packaging_commands
+        self, target, build_context, args, test_infos, output_targets, packaging_commands
     ):
-      super().__init__(target, build_context, args)
+      super().__init__(target, build_context, args, test_infos)
       self.output_targets = output_targets
       self.packaging_commands = packaging_commands
 
diff --git a/ci/optimized_targets.py b/ci/optimized_targets.py
index 688bdd8..4b8b453 100644
--- a/ci/optimized_targets.py
+++ b/ci/optimized_targets.py
@@ -42,10 +42,12 @@
       target: str,
       build_context: BuildContext,
       args: argparse.Namespace,
+      test_infos
   ):
     self.target = target
     self.build_context = build_context
     self.args = args
+    self.test_infos = test_infos
 
   def get_build_targets(self) -> set[str]:
     features = self.build_context.enabled_build_features
diff --git a/ci/optimized_targets_test.py b/ci/optimized_targets_test.py
index 0b0c0ec..fd9e17c 100644
--- a/ci/optimized_targets_test.py
+++ b/ci/optimized_targets_test.py
@@ -224,7 +224,7 @@
     if not build_context:
       build_context = self._create_build_context()
     return optimized_targets.GeneralTestsOptimizer(
-        'general-tests', build_context, None
+        'general-tests', build_context, None, []
     )
 
   def _create_build_context(