Rework general-tests optimization

Rework general-tests optimization to use the test discovery agent to
determine the modules that are needed to be included in
general-tests.zip. Also start reporting optimization decisions silently.

Test: atest build_test_suites_test; atest optimized_targets_test
Bug: 358215235
Change-Id: Iabff6729e5743805167eed87d7ef5d901b255a61
diff --git a/ci/build_test_suites b/ci/build_test_suites
index 74470a8..a63f3fc 100755
--- a/ci/build_test_suites
+++ b/ci/build_test_suites
@@ -15,5 +15,5 @@
 # limitations under the License.
 set -euo pipefail
 
-build/soong/soong_ui.bash --make-mode build_test_suites
-$(build/soong/soong_ui.bash --dumpvar-mode HOST_OUT)/bin/build_test_suites $@
+build/soong/soong_ui.bash --make-mode dist build_test_suites general-tests-files-list test_mapping || exit $?
+$(build/soong/soong_ui.bash --dumpvar-mode HOST_OUT)/bin/build_test_suites $@ || exit $?
diff --git a/ci/metrics_agent.py b/ci/metrics_agent.py
index bc2479e..85cdcbd 100644
--- a/ci/metrics_agent.py
+++ b/ci/metrics_agent.py
@@ -92,15 +92,15 @@
       size: int,
       included_modules: set[str],
   ):
-    target_result = self.target_results.get(target_name)
+    target_result = self._target_results.get(target_name)
     artifact = (
         metrics_pb2.OptimizedBuildMetrics.TargetOptimizationResult.OutputArtifact()
     )
     artifact.name = artifact_name
     artifact.size = size
     for module in included_modules:
-      artifact.included_modules.add(module)
-    target_result.output_artifacts.add(artifact)
+      artifact.included_modules.append(module)
+    target_result.output_artifact.append(artifact)
 
   def end_reporting(self):
     for target_result in self._target_results.values():
diff --git a/ci/optimized_targets.py b/ci/optimized_targets.py
index 4b8b453..0e9723c 100644
--- a/ci/optimized_targets.py
+++ b/ci/optimized_targets.py
@@ -23,7 +23,9 @@
 import subprocess
 
 from build_context import BuildContext
+import metrics_agent
 import test_mapping_module_retriever
+import test_discovery_agent
 
 
 class OptimizedBuildTarget(ABC):
@@ -42,7 +44,7 @@
       target: str,
       build_context: BuildContext,
       args: argparse.Namespace,
-      test_infos
+      test_infos,
   ):
     self.target = target
     self.build_context = build_context
@@ -55,6 +57,8 @@
       self.modules_to_build = self.get_build_targets_impl()
       return self.modules_to_build
 
+    if self.target == 'general-tests':
+      self._report_info_metrics_silently('general-tests.zip')
     self.modules_to_build = {self.target}
     return {self.target}
 
@@ -163,6 +167,16 @@
         f'{dist_dir / name}',
     ]
 
+  def _report_info_metrics_silently(self, artifact_name):
+    try:
+      metrics_agent_instance = metrics_agent.MetricsAgent.instance()
+      targets = self.get_build_targets_impl()
+      metrics_agent_instance.report_optimized_target(self.target)
+      metrics_agent_instance.add_target_artifact(self.target, artifact_name, 0, targets)
+    except Exception as e:
+      logging.error(f'error while silently reporting metrics: {e}')
+
+
 
 class NullOptimizer(OptimizedBuildTarget):
   """No-op target optimizer.
@@ -209,11 +223,7 @@
 class GeneralTestsOptimizer(OptimizedBuildTarget):
   """general-tests optimizer
 
-  This optimizer reads in the list of changed files from the file located in
-  env[CHANGE_INFO] and uses this list alongside the normal TEST MAPPING logic to
-  determine what test mapping modules will run for the given changes. It then
-  builds those modules and packages them in the same way general-tests.zip is
-  normally built.
+  This optimizer uses test discovery to build a list of modules that are needed by all tests configured for the build. These modules are then build and packaged by the optimizer in the same way as they are in a normal build.
   """
 
   # List of modules that are built alongside general-tests as dependencies.
@@ -221,93 +231,99 @@
       'cts-tradefed',
       'vts-tradefed',
       'compatibility-host-util',
-      'general-tests-shared-libs',
   ])
 
   def get_build_targets_impl(self) -> set[str]:
-    change_info_file_path = os.environ.get('CHANGE_INFO')
-    if not change_info_file_path:
-      logging.info(
-          'No CHANGE_INFO env var found, general-tests optimization disabled.'
-      )
-      return {'general-tests'}
-
-    test_infos = self.build_context.test_infos
-    test_mapping_test_groups = set()
-    for test_info in test_infos:
-      is_test_mapping = test_info.is_test_mapping
-      current_test_mapping_test_groups = test_info.test_mapping_test_groups
-      uses_general_tests = test_info.build_target_used('general-tests')
-
-      if uses_general_tests and not is_test_mapping:
-        logging.info(
-            'Test uses general-tests.zip but is not test-mapping, general-tests'
-            ' optimization disabled.'
-        )
-        return {'general-tests'}
-
-      if is_test_mapping:
-        test_mapping_test_groups.update(current_test_mapping_test_groups)
-
-    change_info = ChangeInfo(change_info_file_path)
-    changed_files = change_info.find_changed_files()
-
-    test_mappings = test_mapping_module_retriever.GetTestMappings(
-        changed_files, set()
-    )
+    self._general_tests_outputs = self._get_general_tests_outputs()
+    test_modules = self._get_test_discovery_modules()
 
     modules_to_build = set(self._REQUIRED_MODULES)
-
-    modules_to_build.update(
-        test_mapping_module_retriever.FindAffectedModules(
-            test_mappings, changed_files, test_mapping_test_groups
-        )
-    )
+    self._build_outputs = []
+    for module in test_modules:
+      module_outputs = [output for output in self._general_tests_outputs if module in output]
+      if module_outputs:
+        modules_to_build.add(module)
+        self._build_outputs.extend(module_outputs)
 
     return modules_to_build
 
+  def _get_general_tests_outputs(self) -> list[str]:
+    src_top = pathlib.Path(os.environ.get('TOP', os.getcwd()))
+    soong_vars = self._query_soong_vars(
+        src_top,
+        [
+            'PRODUCT_OUT',
+        ],
+    )
+    product_out = pathlib.Path(soong_vars.get('PRODUCT_OUT'))
+    with open(f'{product_out / "general-tests_files"}') as general_tests_list_file:
+      general_tests_list = general_tests_list_file.readlines()
+    with open(f'{product_out / "general-tests_host_files"}') as general_tests_list_file:
+      self._general_tests_host_outputs = general_tests_list_file.readlines()
+    with open(f'{product_out / "general-tests_target_files"}') as general_tests_list_file:
+      self._general_tests_target_outputs = general_tests_list_file.readlines()
+    return general_tests_list
+
+
+  def _get_test_discovery_modules(self) -> set[str]:
+    test_modules = set()
+    for test_info in self.test_infos:
+      tf_command = self._build_tf_command(test_info)
+      discovery_agent = test_discovery_agent.TestDiscoveryAgent(tradefed_args=tf_command, test_mapping_zip_path=os.environ.get('DIST_DIR')+'/test_mappings.zip')
+      modules, dependencies = discovery_agent.discover_test_mapping_test_modules()
+      for regex in modules:
+        test_modules.add(regex)
+    return test_modules
+
+
+  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
+
   def get_package_outputs_commands_impl(self):
     src_top = pathlib.Path(os.environ.get('TOP', os.getcwd()))
     dist_dir = pathlib.Path(os.environ.get('DIST_DIR'))
+    tmp_dir = pathlib.Path(os.environ.get('TMPDIR'))
+    print(f'modules: {self.modules_to_build}')
 
+    host_outputs = [str(src_top) + '/' + file for file in self._general_tests_host_outputs if any('/'+module+'/' in file for module in self.modules_to_build)]
+    target_outputs = [str(src_top) + '/' + file for file in self._general_tests_target_outputs if any('/'+module+'/' in file for module in self.modules_to_build)]
+    host_config_files = [file for file in host_outputs if file.endswith('.config\n')]
+    target_config_files = [file for file in target_outputs if file.endswith('.config\n')]
+    logging.info(host_outputs)
+    logging.info(target_outputs)
+    with open(f"{tmp_dir / 'host.list'}", 'w') as host_list_file:
+      for output in host_outputs:
+        host_list_file.write(output)
+    with open(f"{tmp_dir / 'target.list'}", 'w') as target_list_file:
+      for output in target_outputs:
+        target_list_file.write(output)
     soong_vars = self._query_soong_vars(
         src_top,
         [
-            'HOST_OUT_TESTCASES',
-            'TARGET_OUT_TESTCASES',
             'PRODUCT_OUT',
             'SOONG_HOST_OUT',
             'HOST_OUT',
         ],
     )
-    host_out_testcases = pathlib.Path(soong_vars.get('HOST_OUT_TESTCASES'))
-    target_out_testcases = pathlib.Path(soong_vars.get('TARGET_OUT_TESTCASES'))
     product_out = pathlib.Path(soong_vars.get('PRODUCT_OUT'))
     soong_host_out = pathlib.Path(soong_vars.get('SOONG_HOST_OUT'))
     host_out = pathlib.Path(soong_vars.get('HOST_OUT'))
-
-    host_paths = []
-    target_paths = []
-    host_config_files = []
-    target_config_files = []
-    for module in self.modules_to_build:
-      # The required modules are handled separately, no need to package.
-      if module in self._REQUIRED_MODULES:
-        continue
-
-      host_path = host_out_testcases / module
-      if os.path.exists(host_path):
-        host_paths.append(host_path)
-        self._collect_config_files(src_top, host_path, host_config_files)
-
-      target_path = target_out_testcases / module
-      if os.path.exists(target_path):
-        target_paths.append(target_path)
-        self._collect_config_files(src_top, target_path, target_config_files)
-
-      if not os.path.exists(host_path) and not os.path.exists(target_path):
-        logging.info(f'No host or target build outputs found for {module}.')
-
     zip_commands = []
 
     zip_commands.extend(
@@ -322,24 +338,23 @@
     )
 
     zip_command = self._base_zip_command(src_top, dist_dir, 'general-tests.zip')
-
     # Add host testcases.
-    if host_paths:
+    if host_outputs:
       zip_command.extend(
           self._generate_zip_options_for_items(
               prefix='host',
-              relative_root=f'{src_top / soong_host_out}',
-              directories=host_paths,
+              relative_root=str(host_out),
+              list_files=[f"{tmp_dir / 'host.list'}"],
           )
       )
 
     # Add target testcases.
-    if target_paths:
+    if target_outputs:
       zip_command.extend(
           self._generate_zip_options_for_items(
               prefix='target',
-              relative_root=f'{src_top / product_out}',
-              directories=target_paths,
+              relative_root=str(product_out),
+              list_files=[f"{tmp_dir / 'target.list'}"],
           )
       )
 
@@ -359,20 +374,11 @@
         )
     )
 
+    zip_command.append('-sha256')
+
     zip_commands.append(zip_command)
     return zip_commands
 
-  def _collect_config_files(
-      self,
-      src_top: pathlib.Path,
-      root_dir: pathlib.Path,
-      config_files: list[str],
-  ):
-    for root, dirs, files in os.walk(src_top / root_dir):
-      for file in files:
-        if file.endswith('.config'):
-          config_files.append(root_dir / file)
-
   def _get_zip_test_configs_zips_commands(
       self,
       src_top: pathlib.Path,
diff --git a/ci/optimized_targets_test.py b/ci/optimized_targets_test.py
index fd9e17c..8256acd 100644
--- a/ci/optimized_targets_test.py
+++ b/ci/optimized_targets_test.py
@@ -26,6 +26,7 @@
 from build_context import BuildContext
 import optimized_targets
 from pyfakefs import fake_filesystem_unittest
+import test_discovery_agent
 
 
 class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
@@ -38,13 +39,10 @@
     self.mock_os_environ = os_environ_patcher.start()
 
     self._setup_working_build_env()
-    self._write_change_info_file()
     test_mapping_dir = pathlib.Path('/project/path/file/path')
     test_mapping_dir.mkdir(parents=True)
-    self._write_test_mapping_file()
 
   def _setup_working_build_env(self):
-    self.change_info_file = pathlib.Path('/tmp/change_info')
     self._write_soong_ui_file()
     self._host_out_testcases = pathlib.Path('/tmp/top/host_out_testcases')
     self._host_out_testcases.mkdir(parents=True)
@@ -56,14 +54,15 @@
     self._soong_host_out.mkdir(parents=True)
     self._host_out = pathlib.Path('/tmp/top/host_out')
     self._host_out.mkdir(parents=True)
+    self._write_general_tests_files_outputs()
 
     self._dist_dir = pathlib.Path('/tmp/top/out/dist')
     self._dist_dir.mkdir(parents=True)
 
     self.mock_os_environ.update({
-        'CHANGE_INFO': str(self.change_info_file),
         'TOP': '/tmp/top',
         'DIST_DIR': '/tmp/top/out/dist',
+        'TMPDIR': '/tmp/'
     })
 
   def _write_soong_ui_file(self):
@@ -72,127 +71,92 @@
     with open(os.path.join(soong_path, 'soong_ui.bash'), 'w') as f:
       f.write("""
               #/bin/bash
-              echo HOST_OUT_TESTCASES='/tmp/top/host_out_testcases'
-              echo TARGET_OUT_TESTCASES='/tmp/top/target_out_testcases'
               echo PRODUCT_OUT='/tmp/top/product_out'
               echo SOONG_HOST_OUT='/tmp/top/soong_host_out'
               echo HOST_OUT='/tmp/top/host_out'
               """)
     os.chmod(os.path.join(soong_path, 'soong_ui.bash'), 0o666)
 
-  def _write_change_info_file(self):
-    change_info_contents = {
-        'changes': [{
-            'projectPath': '/project/path',
-            'revisions': [{
-                'fileInfos': [{
-                    'path': 'file/path/file_name',
-                }],
-            }],
-        }]
-    }
+  def _write_general_tests_files_outputs(self):
+    with open(os.path.join(self._product_out, 'general-tests_files'), 'w') as f:
+      f.write("""
+              path/to/module_1/general-tests-host-file
+              path/to/module_1/general-tests-host-file.config
+              path/to/module_1/general-tests-target-file
+              path/to/module_1/general-tests-target-file.config
+              path/to/module_2/general-tests-host-file
+              path/to/module_2/general-tests-host-file.config
+              path/to/module_2/general-tests-target-file
+              path/to/module_2/general-tests-target-file.config
+              path/to/module_1/general-tests-host-file
+              path/to/module_1/general-tests-host-file.config
+              path/to/module_1/general-tests-target-file
+              path/to/module_1/general-tests-target-file.config
+              """)
+    with open(os.path.join(self._product_out, 'general-tests_host_files'), 'w') as f:
+      f.write("""
+              path/to/module_1/general-tests-host-file
+              path/to/module_1/general-tests-host-file.config
+              path/to/module_2/general-tests-host-file
+              path/to/module_2/general-tests-host-file.config
+              path/to/module_1/general-tests-host-file
+              path/to/module_1/general-tests-host-file.config
+              """)
+    with open(os.path.join(self._product_out, 'general-tests_target_files'), 'w') as f:
+      f.write("""
+              path/to/module_1/general-tests-target-file
+              path/to/module_1/general-tests-target-file.config
+              path/to/module_2/general-tests-target-file
+              path/to/module_2/general-tests-target-file.config
+              path/to/module_1/general-tests-target-file
+              path/to/module_1/general-tests-target-file.config
+              """)
 
-    with open(self.change_info_file, 'w') as f:
-      json.dump(change_info_contents, f)
-
-  def _write_test_mapping_file(self):
-    test_mapping_contents = {
-        'test-mapping-group': [
-            {
-                'name': 'test_mapping_module',
-            },
-        ],
-    }
-
-    with open('/project/path/file/path/TEST_MAPPING', 'w') as f:
-      json.dump(test_mapping_contents, f)
-
-  def test_general_tests_optimized(self):
-    optimizer = self._create_general_tests_optimizer()
-
-    build_targets = optimizer.get_build_targets()
-
-    expected_build_targets = set(
-        optimized_targets.GeneralTestsOptimizer._REQUIRED_MODULES
-    )
-    expected_build_targets.add('test_mapping_module')
-
-    self.assertSetEqual(build_targets, expected_build_targets)
-
-  def test_no_change_info_no_optimization(self):
-    del os.environ['CHANGE_INFO']
-
-    optimizer = self._create_general_tests_optimizer()
-
-    build_targets = optimizer.get_build_targets()
-
-    self.assertSetEqual(build_targets, {'general-tests'})
-
-  def test_mapping_groups_unused_module_not_built(self):
-    test_context = self._create_test_context()
-    test_context['testInfos'][0]['extraOptions'] = [
-        {
-            'key': 'additional-files-filter',
-            'values': ['general-tests.zip'],
-        },
-        {
-            'key': 'test-mapping-test-group',
-            'values': ['unused-test-mapping-group'],
-        },
-    ]
-    optimizer = self._create_general_tests_optimizer(
-        build_context=self._create_build_context(test_context=test_context)
-    )
-
-    build_targets = optimizer.get_build_targets()
-
-    expected_build_targets = set(
-        optimized_targets.GeneralTestsOptimizer._REQUIRED_MODULES
-    )
-    self.assertSetEqual(build_targets, expected_build_targets)
-
-  def test_general_tests_used_by_non_test_mapping_test_no_optimization(self):
-    test_context = self._create_test_context()
-    test_context['testInfos'][0]['extraOptions'] = [{
-        'key': 'additional-files-filter',
-        'values': ['general-tests.zip'],
-    }]
-    optimizer = self._create_general_tests_optimizer(
-        build_context=self._create_build_context(test_context=test_context)
-    )
-
-    build_targets = optimizer.get_build_targets()
-
-    self.assertSetEqual(build_targets, {'general-tests'})
-
-  def test_malformed_change_info_raises(self):
-    with open(self.change_info_file, 'w') as f:
-      f.write('not change info')
-
-    optimizer = self._create_general_tests_optimizer()
-
-    with self.assertRaises(json.decoder.JSONDecodeError):
-      build_targets = optimizer.get_build_targets()
-
-  def test_malformed_test_mapping_raises(self):
-    with open('/project/path/file/path/TEST_MAPPING', 'w') as f:
-      f.write('not test mapping')
-
-    optimizer = self._create_general_tests_optimizer()
-
-    with self.assertRaises(json.decoder.JSONDecodeError):
-      build_targets = optimizer.get_build_targets()
 
   @mock.patch('subprocess.run')
-  def test_packaging_outputs_success(self, subprocess_run):
+  @mock.patch.object(test_discovery_agent.TestDiscoveryAgent, 'discover_test_mapping_test_modules')
+  def test_general_tests_optimized(self, discover_modules, subprocess_run):
     subprocess_run.return_value = self._get_soong_vars_output()
+    discover_modules.return_value = (['module_1'], ['dependency_1'])
+
+    optimizer = self._create_general_tests_optimizer()
+
+    build_targets = optimizer.get_build_targets()
+
+    expected_build_targets = set(
+        optimized_targets.GeneralTestsOptimizer._REQUIRED_MODULES
+    )
+    expected_build_targets.add('module_1')
+
+    self.assertSetEqual(build_targets, expected_build_targets)
+
+  @mock.patch('subprocess.run')
+  @mock.patch.object(test_discovery_agent.TestDiscoveryAgent, 'discover_test_mapping_test_modules')
+  def test_module_unused_module_not_built(self, discover_modules, subprocess_run):
+    subprocess_run.return_value = self._get_soong_vars_output()
+    discover_modules.return_value = (['no_module'], ['dependency_1'])
+
+    optimizer = self._create_general_tests_optimizer()
+
+    build_targets = optimizer.get_build_targets()
+
+    expected_build_targets = set(
+        optimized_targets.GeneralTestsOptimizer._REQUIRED_MODULES
+    )
+    self.assertSetEqual(build_targets, expected_build_targets)
+
+  @mock.patch('subprocess.run')
+  @mock.patch.object(test_discovery_agent.TestDiscoveryAgent, 'discover_test_mapping_test_modules')
+  def test_packaging_outputs_success(self, discover_modules, subprocess_run):
+    subprocess_run.return_value = self._get_soong_vars_output()
+    discover_modules.return_value = (['module_1'], ['dependency_1'])
     optimizer = self._create_general_tests_optimizer()
     self._set_up_build_outputs(['test_mapping_module'])
 
     targets = optimizer.get_build_targets()
     package_commands = optimizer.get_package_outputs_commands()
 
-    self._verify_soong_zip_commands(package_commands, ['test_mapping_module'])
+    self._verify_soong_zip_commands(package_commands, ['module_1'])
 
   @mock.patch('subprocess.run')
   def test_get_soong_dumpvars_fails_raises(self, subprocess_run):
@@ -200,10 +164,8 @@
     optimizer = self._create_general_tests_optimizer()
     self._set_up_build_outputs(['test_mapping_module'])
 
-    targets = optimizer.get_build_targets()
-
     with self.assertRaisesRegex(RuntimeError, 'Soong dumpvars failed!'):
-      package_commands = optimizer.get_package_outputs_commands()
+      targets = optimizer.get_build_targets()
 
   @mock.patch('subprocess.run')
   def test_get_soong_dumpvars_bad_output_raises(self, subprocess_run):
@@ -213,18 +175,16 @@
     optimizer = self._create_general_tests_optimizer()
     self._set_up_build_outputs(['test_mapping_module'])
 
-    targets = optimizer.get_build_targets()
-
     with self.assertRaisesRegex(
         RuntimeError, 'Error parsing soong dumpvars output'
     ):
-      package_commands = optimizer.get_package_outputs_commands()
+      targets = optimizer.get_build_targets()
 
   def _create_general_tests_optimizer(self, build_context: BuildContext = None):
     if not build_context:
       build_context = self._create_build_context()
     return optimized_targets.GeneralTestsOptimizer(
-        'general-tests', build_context, None, []
+        'general-tests', build_context, None, build_context.test_infos
     )
 
   def _create_build_context(
@@ -274,11 +234,10 @@
     return_value = subprocess.CompletedProcess(args=[], returncode=return_code)
     if not stdout:
       stdout = textwrap.dedent(f"""\
-                               HOST_OUT_TESTCASES='{self._host_out_testcases}'
-                               TARGET_OUT_TESTCASES='{self._target_out_testcases}'
                                PRODUCT_OUT='{self._product_out}'
                                SOONG_HOST_OUT='{self._soong_host_out}'
-                               HOST_OUT='{self._host_out}'""")
+                               HOST_OUT='{self._host_out}'
+                               """)
 
     return_value.stdout = stdout
     return return_value