Implement package_outputs in GeneralTestsOptimizer

Implement the output packaging step in the GeneralTestsOptimizer. This
step takes all the outputs built by the general-tests target (in the
case where it was optimized) and packages them into the necessary zips
generated by the target normally.

Test: atest build_test_suites_test; atest optimized_targets_test
Bug: 358215235
Change-Id: I5d27eef4e37137cc9b6e235f52f3856ba0b30460

ci/optimized_targets.py

@@ -16,12 +16,13 @@
 from abc import ABC
 import argparse
 import functools
-from build_context import BuildContext
 import json
 import logging
 import os
-from typing import Self
+import pathlib
+import subprocess
 
+from build_context import BuildContext
 import test_mapping_module_retriever
 
 
@@ -33,6 +34,9 @@ class OptimizedBuildTarget(ABC):
   build.
   """
 
+  _SOONG_UI_BASH_PATH = 'build/soong/soong_ui.bash'
+  _PREBUILT_SOONG_ZIP_PATH = 'prebuilts/build-tools/linux-x86/bin/soong_zip'
+
   def __init__(
       self,
       target: str,
@@ -75,6 +79,88 @@ class OptimizedBuildTarget(ABC):
         f'get_build_targets_impl not implemented in {type(self).__name__}'
     )
 
+  def _generate_zip_options_for_items(
+      self,
+      prefix: str = '',
+      relative_root: str = '',
+      list_files: list[str] | None = None,
+      files: list[str] | None = None,
+      directories: list[str] | None = None,
+  ) -> list[str]:
+    if not list_files and not files and not directories:
+      raise RuntimeError(
+          f'No items specified to be added to zip! Prefix: {prefix}, Relative'
+          f' root: {relative_root}'
+      )
+    command_segment = []
+    # These are all soong_zip options so consult soong_zip --help for specifics.
+    if prefix:
+      command_segment.append('-P')
+      command_segment.append(prefix)
+    if relative_root:
+      command_segment.append('-C')
+      command_segment.append(relative_root)
+    if list_files:
+      for list_file in list_files:
+        command_segment.append('-l')
+        command_segment.append(list_file)
+    if files:
+      for file in files:
+        command_segment.append('-f')
+        command_segment.append(file)
+    if directories:
+      for directory in directories:
+        command_segment.append('-D')
+        command_segment.append(directory)
+
+    return command_segment
+
+  def _query_soong_vars(
+      self, src_top: pathlib.Path, soong_vars: list[str]
+  ) -> dict[str, str]:
+    process_result = subprocess.run(
+        args=[
+            f'{src_top / self._SOONG_UI_BASH_PATH}',
+            '--dumpvar-mode',
+            '--abs',
+            soong_vars,
+        ],
+        env=os.environ,
+        check=False,
+        capture_output=True,
+    )
+    if not process_result.returncode == 0:
+      logging.error('soong dumpvars command failed! stderr:')
+      logging.error(process_result.stderr)
+      raise RuntimeError('Soong dumpvars failed! See log for stderr.')
+
+    if not process_result.stdout:
+      raise RuntimeError(
+          'Necessary soong variables ' + soong_vars + ' not found.'
+      )
+
+    try:
+      return {
+          line.split('=')[0]: line.split('=')[1].strip("'")
+          for line in process_result.stdout.split('\n')
+      }
+    except IndexError as e:
+      raise RuntimeError(
+          'Error parsing soong dumpvars output! See output here:'
+          f' {process_result.stdout}',
+          e,
+      )
+
+  def _base_zip_command(
+      self, src_top: pathlib.Path, dist_dir: pathlib.Path, name: str
+  ) -> list[str]:
+    return [
+        f'{src_top / self._PREBUILT_SOONG_ZIP_PATH }',
+        '-d',
+        '-o',
+        f'{dist_dir / name}',
+    ]
+
 
 class NullOptimizer(OptimizedBuildTarget):
   """No-op target optimizer.
@@ -121,8 +207,6 @@ class ChangeInfo:
 class GeneralTestsOptimizer(OptimizedBuildTarget):
   """general-tests optimizer
 
-  TODO(b/358215235): Implement
-
   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
@@ -177,6 +261,208 @@ class GeneralTestsOptimizer(OptimizedBuildTarget):
 
     return modules_to_build
 
+  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'))
+
+    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:
+      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(
+        self._get_zip_test_configs_zips_commands(
+            dist_dir,
+            host_out,
+            product_out,
+            host_config_files,
+            target_config_files,
+        )
+    )
+
+    zip_command = self._base_zip_command(
+        host_out, dist_dir, 'general-tests.zip'
+    )
+
+    # Add host testcases.
+    zip_command.extend(
+        self._generate_zip_options_for_items(
+            prefix='host',
+            relative_root=f'{src_top / soong_host_out}',
+            directories=host_paths,
+        )
+    )
+
+    # Add target testcases.
+    zip_command.extend(
+        self._generate_zip_options_for_items(
+            prefix='target',
+            relative_root=f'{src_top / product_out}',
+            directories=target_paths,
+        )
+    )
+
+    # TODO(lucafarsi): Push this logic into a general-tests-minimal build command
+    # Add necessary tools. These are also hardcoded in general-tests.mk.
+    framework_path = soong_host_out / 'framework'
+
+    zip_command.extend(
+        self._generate_zip_options_for_items(
+            prefix='host/tools',
+            relative_root=str(framework_path),
+            files=[
+                f"{framework_path / 'cts-tradefed.jar'}",
+                f"{framework_path / 'compatibility-host-util.jar'}",
+                f"{framework_path / 'vts-tradefed.jar'}",
+            ],
+        )
+    )
+
+    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,
+      dist_dir: pathlib.Path,
+      host_out: pathlib.Path,
+      product_out: pathlib.Path,
+      host_config_files: list[str],
+      target_config_files: list[str],
+  ) -> tuple[list[str], list[str]]:
+    """Generate general-tests_configs.zip and general-tests_list.zip.
+
+    general-tests_configs.zip contains all of the .config files that were
+    built and general-tests_list.zip contains a text file which lists
+    all of the .config files that are in general-tests_configs.zip.
+
+    general-tests_configs.zip is organized as follows:
+    /
+      host/
+        testcases/
+          test_1.config
+          test_2.config
+          ...
+      target/
+        testcases/
+          test_1.config
+          test_2.config
+          ...
+
+    So the process is we write out the paths to all the host config files into
+    one
+    file and all the paths to the target config files in another. We also write
+    the paths to all the config files into a third file to use for
+    general-tests_list.zip.
+
+    Args:
+      dist_dir: dist directory.
+      host_out: host out directory.
+      product_out: product out directory.
+      host_config_files: list of all host config files.
+      target_config_files: list of all target config files.
+
+    Returns:
+      The commands to generate general-tests_configs.zip and
+      general-tests_list.zip
+    """
+    with open(
+        f"{host_out / 'host_general-tests_list'}", 'w'
+    ) as host_list_file, open(
+        f"{product_out / 'target_general-tests_list'}", 'w'
+    ) as target_list_file, open(
+        f"{host_out / 'general-tests_list'}", 'w'
+    ) as list_file:
+
+      for config_file in host_config_files:
+        host_list_file.write(f'{config_file}' + '\n')
+        list_file.write('host/' + os.path.relpath(config_file, host_out) + '\n')
+
+      for config_file in target_config_files:
+        target_list_file.write(f'{config_file}' + '\n')
+        list_file.write(
+            'target/' + os.path.relpath(config_file, product_out) + '\n'
+        )
+
+    zip_commands = []
+
+    tests_config_zip_command = self._base_zip_command(
+        host_out, dist_dir, 'general-tests_configs.zip'
+    )
+    tests_config_zip_command.extend(
+        self._generate_zip_options_for_items(
+            prefix='host',
+            relative_root=str(host_out),
+            list_files=[f"{host_out / 'host_general-tests_list'}"],
+        )
+    )
+
+    tests_config_zip_command.extend(
+        self._generate_zip_options_for_items(
+            prefix='target',
+            relative_root=str(product_out),
+            list_files=[
+                f"{product_out / 'target_general-tests_list'}"
+            ],
+        ),
+    )
+
+    zip_commands.append(tests_config_zip_command)
+
+    tests_list_zip_command = self._base_zip_command(
+        host_out, dist_dir, 'general-tests_list.zip'
+    )
+    tests_list_zip_command.extend(
+        self._generate_zip_options_for_items(
+            relative_root=str(host_out),
+            files=[f"{host_out / 'general-tests_list'}"],
+        )
+    )
+    zip_commands.append(tests_list_zip_command)
+
+    return zip_commands
+
   def get_enabled_flag(self):
     return 'general_tests_optimized'
 

ci/optimized_targets_test.py

@@ -19,10 +19,12 @@ import logging
 import os
 import pathlib
 import re
+import subprocess
+import textwrap
 import unittest
 from unittest import mock
-import optimized_targets
 from build_context import BuildContext
+import optimized_targets
 from pyfakefs import fake_filesystem_unittest
 
 
@@ -43,11 +45,68 @@ class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
 
   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)
+    self._target_out_testcases = pathlib.Path('/tmp/top/target_out_testcases')
+    self._target_out_testcases.mkdir(parents=True)
+    self._product_out = pathlib.Path('/tmp/top/product_out')
+    self._product_out.mkdir(parents=True)
+    self._soong_host_out = pathlib.Path('/tmp/top/soong_host_out')
+    self._soong_host_out.mkdir(parents=True)
+    self._host_out = pathlib.Path('/tmp/top/host_out')
+    self._host_out.mkdir(parents=True)
+
+    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',
     })
 
+  def _write_soong_ui_file(self):
+    soong_path = pathlib.Path('/tmp/top/build/soong')
+    soong_path.mkdir(parents=True)
+    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',
+                }],
+            }],
+        }]
+    }
+
+    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()
 
@@ -124,36 +183,56 @@ class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
     with self.assertRaises(json.decoder.JSONDecodeError):
       build_targets = optimizer.get_build_targets()
 
-  def _write_change_info_file(self):
-    change_info_contents = {
-        'changes': [{
-            'projectPath': '/project/path',
-            'revisions': [{
-                'fileInfos': [{
-                    'path': 'file/path/file_name',
-                }],
-            }],
-        }]
-    }
+  @mock.patch('subprocess.run')
+  def test_packaging_outputs_success(self, subprocess_run):
+    subprocess_run.return_value = self._get_soong_vars_output()
+    optimizer = self._create_general_tests_optimizer()
+    self._set_up_build_outputs(['test_mapping_module'])
 
-    with open(self.change_info_file, 'w') as f:
-      json.dump(change_info_contents, f)
+    targets = optimizer.get_build_targets()
+    package_commands = optimizer.get_package_outputs_commands()
 
-  def _write_test_mapping_file(self):
-    test_mapping_contents = {
-        'test-mapping-group': [
-            {
-                'name': 'test_mapping_module',
-            },
-        ],
-    }
+    self._verify_soong_zip_commands(package_commands, ['test_mapping_module'])
 
-    with open('/project/path/file/path/TEST_MAPPING', 'w') as f:
-      json.dump(test_mapping_contents, f)
+  @mock.patch('subprocess.run')
+  def test_get_soong_dumpvars_fails_raises(self, subprocess_run):
+    subprocess_run.return_value = self._get_soong_vars_output(return_code=-1)
+    optimizer = self._create_general_tests_optimizer()
+    self._set_up_build_outputs(['test_mapping_module'])
 
-  def _create_general_tests_optimizer(
-      self, build_context: BuildContext = None
-  ):
+    targets = optimizer.get_build_targets()
+
+    with self.assertRaisesRegex(RuntimeError, 'Soong dumpvars failed!'):
+      package_commands = optimizer.get_package_outputs_commands()
+
+  @mock.patch('subprocess.run')
+  def test_get_soong_dumpvars_bad_output_raises(self, subprocess_run):
+    subprocess_run.return_value = self._get_soong_vars_output(
+        stdout='This output is bad'
+    )
+    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()
+
+  @mock.patch('subprocess.run')
+  def test_no_build_outputs_packaging_fails(self, subprocess_run):
+    subprocess_run.return_value = self._get_soong_vars_output()
+    optimizer = self._create_general_tests_optimizer()
+
+    targets = optimizer.get_build_targets()
+
+    with self.assertRaisesRegex(
+        RuntimeError, 'No items specified to be added to zip'
+    ):
+      package_commands = optimizer.get_package_outputs_commands()
+
+  def _create_general_tests_optimizer(self, build_context: BuildContext = None):
     if not build_context:
       build_context = self._create_build_context()
     return optimized_targets.GeneralTestsOptimizer(
@@ -170,7 +249,9 @@ class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
     build_context_dict = {}
     build_context_dict['enabledBuildFeatures'] = [{'name': 'optimized_build'}]
     if general_tests_optimized:
-      build_context_dict['enabledBuildFeatures'].append({'name': 'general_tests_optimized'})
+      build_context_dict['enabledBuildFeatures'].append(
+          {'name': 'general_tests_optimized'}
+      )
     build_context_dict['testContext'] = test_context
     return BuildContext(build_context_dict)
 
@@ -199,6 +280,81 @@ class GeneralTestsOptimizerTest(fake_filesystem_unittest.TestCase):
         ],
     }
 
+  def _get_soong_vars_output(
+      self, return_code: int = 0, stdout: str = ''
+  ) -> subprocess.CompletedProcess:
+    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}'""")
+
+    return_value.stdout = stdout
+    return return_value
+
+  def _set_up_build_outputs(self, targets: list[str]):
+    for target in targets:
+      host_dir = self._host_out_testcases / target
+      host_dir.mkdir()
+      (host_dir / f'{target}.config').touch()
+      (host_dir / f'test_file').touch()
+
+      target_dir = self._target_out_testcases / target
+      target_dir.mkdir()
+      (target_dir / f'{target}.config').touch()
+      (target_dir / f'test_file').touch()
+
+  def _verify_soong_zip_commands(self, commands: list[str], targets: list[str]):
+    """Verify the structure of the zip commands.
+
+    Zip commands have to start with the soong_zip binary path, then are followed
+    by a couple of options and the name of the file being zipped. Depending on
+    which zip we are creating look for a few essential items being added in
+    those zips.
+
+    Args:
+      commands: list of command lists
+      targets: list of targets expected to be in general-tests.zip
+    """
+    for command in commands:
+      self.assertEqual(
+          '/tmp/top/host_out/prebuilts/build-tools/linux-x86/bin/soong_zip',
+          command[0],
+      )
+      self.assertEqual('-d', command[1])
+      self.assertEqual('-o', command[2])
+      match (command[3]):
+        case '/tmp/top/out/dist/general-tests_configs.zip':
+          self.assertIn(f'{self._host_out}/host_general-tests_list', command)
+          self.assertIn(
+              f'{self._product_out}/target_general-tests_list', command
+          )
+          return
+        case '/tmp/top/out/dist/general-tests_list.zip':
+          self.assertIn('-f', command)
+          self.assertIn(f'{self._host_out}/general-tests_list', command)
+          return
+        case '/tmp/top/out/dist/general-tests.zip':
+          for target in targets:
+            self.assertIn(f'{self._host_out_testcases}/{target}', command)
+            self.assertIn(f'{self._target_out_testcases}/{target}', command)
+          self.assertIn(
+              f'{self._soong_host_out}/framework/cts-tradefed.jar', command
+          )
+          self.assertIn(
+              f'{self._soong_host_out}/framework/compatibility-host-util.jar',
+              command,
+          )
+          self.assertIn(
+              f'{self._soong_host_out}/framework/vts-tradefed.jar', command
+          )
+          return
+        case _:
+          self.fail(f'malformed command: {command}')
+
 
 if __name__ == '__main__':
   # Setup logging to be silent so unit tests can pass through TF.