Begin reporting Test Discovery Agent metrics

Start running the Test Discovery Agent and reportings its results via
metrics. No changes to the actual build process yet.

Test: atest build_test_suites_test && atest optimized_targets_test
Bug: 372973116
Change-Id: I958e034985c41ede8d6a2700311a4dd6e7ac18ba

4 files changed, 58 insertions, 2 deletions

diff --git a/ci/build_context.py b/ci/build_context.py
index cc48d53992..c7a1defb57 100644
--- a/ci/build_context.py
+++ b/ci/build_context.py
@@ -47,6 +47,9 @@ class BuildContext:
       self.is_test_mapping = False
       self.test_mapping_test_groups = set()
       self.file_download_options = set()
+      self.name = test_info_dict.get('name')
+      self.command = test_info_dict.get('command')
+      self.extra_options = test_info_dict.get('extraOptions')
       for opt in test_info_dict.get('extraOptions', []):
         key = opt.get('key')
         if key == 'test-mapping-test-group':
diff --git a/ci/build_test_suites.py b/ci/build_test_suites.py
index 7d76b9afdc..cd9d76d23b 100644
--- a/ci/build_test_suites.py
+++ b/ci/build_test_suites.py
@@ -20,12 +20,14 @@ import json
 import logging
 import os
 import pathlib
+import re
 import subprocess
 import sys
 from typing import Callable
 from build_context import BuildContext
 import optimized_targets
 import metrics_agent
+import test_discovery_agent
 
 
 REQUIRED_ENV_VARS = frozenset(['TARGET_PRODUCT', 'TARGET_RELEASE', 'TOP'])
@@ -71,7 +73,24 @@ class BuildPlanner:
 
     build_targets = set()
     packaging_commands_getters = []
+    test_discovery_zip_regexes = set()
+    optimization_rationale = ''
+    try:
+      # Do not use these regexes for now, only run this to collect data on what
+      # would be optimized.
+      test_discovery_zip_regexes = self._get_test_discovery_zip_regexes()
+      logging.info(f'Discovered test discovery regexes: {test_discovery_zip_regexes}')
+    except test_discovery_agent.TestDiscoveryError as e:
+      optimization_rationale = e.message
+      logging.warning(f'Unable to perform test discovery: {optimization_rationale}')
     for target in self.args.extra_targets:
+      if optimization_rationale:
+        get_metrics_agent().report_unoptimized_target(target, optimization_rationale)
+      else:
+        regex = r'\b(%s)\b' % re.escape(target)
+        if any(re.search(regex, opt) for opt in test_discovery_zip_regexes):
+          get_metrics_agent().report_optimized_target(target)
+
       if self._unused_target_exclusion_enabled(
           target
       ) and not self.build_context.build_target_used(target):
@@ -98,6 +117,34 @@ class BuildPlanner:
         in self.build_context.enabled_build_features
     )
 
+  def _get_test_discovery_zip_regexes(self) -> set[str]:
+    build_target_regexes = set()
+    for test_info in self.build_context.test_infos:
+      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():
+        build_target_regexes.add(regex)
+    return build_target_regexes
+
+
+  def _build_tf_command(self, test_info) -> list[str]:
+    command = [test_info.command]
+    for extra_option in test_info.extra_options:
+      if not extra_option.get('key'):
+        continue
+      arg_key = '--' + extra_option.get('key')
+      if arg_key == '--build-id':
+        command.append(arg_key)
+        command.append(os.environ.get('BUILD_NUMBER'))
+        continue
+      if extra_option.get('values'):
+        for value in extra_option.get('values'):
+          command.append(arg_key)
+          command.append(value)
+      else:
+        command.append(arg_key)
+
+    return command
 
 @dataclass(frozen=True)
 class BuildPlan:
diff --git a/ci/build_test_suites_test.py b/ci/build_test_suites_test.py
index c7cd4ab1b0..26f4316f8c 100644
--- a/ci/build_test_suites_test.py
+++ b/ci/build_test_suites_test.py
@@ -38,6 +38,7 @@ import ci_test_lib
 import optimized_targets
 from pyfakefs import fake_filesystem_unittest
 import metrics_agent
+import test_discovery_agent
 
 
 class BuildTestSuitesTest(fake_filesystem_unittest.TestCase):
@@ -261,6 +262,12 @@ class BuildPlannerTest(unittest.TestCase):
     def get_enabled_flag(self):
       return f'{self.target}_enabled'
 
+  def setUp(self):
+    test_discovery_agent_patcher = mock.patch('test_discovery_agent.TestDiscoveryAgent.discover_test_zip_regexes')
+    self.addCleanup(test_discovery_agent_patcher.stop)
+    self.mock_test_discovery_agent_end = test_discovery_agent_patcher.start()
+
+
   def test_build_optimization_off_builds_everything(self):
     build_targets = {'target_1', 'target_2'}
     build_planner = self.create_build_planner(
diff --git a/ci/test_discovery_agent.py b/ci/test_discovery_agent.py
index 4eed28dcb9..008ee47f8e 100644
--- a/ci/test_discovery_agent.py
+++ b/ci/test_discovery_agent.py
@@ -17,7 +17,6 @@ import json
 import logging
 import os
 import subprocess
-import buildbot
 
 
 class TestDiscoveryAgent:
@@ -50,7 +49,7 @@ class TestDiscoveryAgent:
       A list of test zip regexes that TF is going to try to pull files from.
     """
     test_discovery_output_file_name = os.path.join(
-        buildbot.OutDir(), self._TRADEFED_DISCOVERY_OUTPUT_FILE_NAME
+        os.environ.get('TOP'), 'out', self._TRADEFED_DISCOVERY_OUTPUT_FILE_NAME
     )
     with open(
         test_discovery_output_file_name, mode="w+t"