Add optimized build features in build_test_suites

Add the optimized build feature in build_test_suites. WIP

Test: WIP
Bug: 342264003
Change-Id: I05a4ac4026c345f7ced771aa6deb1b6d1c38705c

ci/Android.bp

@@ -14,6 +14,7 @@
 
 package {
     default_applicable_licenses: ["Android-Apache-2.0"],
+    default_team: "trendy_team_adte",
 }
 
 python_test_host {
@@ -74,6 +75,7 @@ python_library_host {
     name: "build_test_suites",
     srcs: [
         "build_test_suites.py",
+        "optimized_targets.py",
     ],
 }
 

ci/build_test_suites.py

@@ -15,11 +15,19 @@
 """Build script for the CI `test_suites` target."""
 
 import argparse
+from dataclasses import dataclass
+import json
 import logging
 import os
 import pathlib
 import subprocess
 import sys
+from typing import Callable
+import optimized_targets
+
+
+REQUIRED_ENV_VARS = frozenset(['TARGET_PRODUCT', 'TARGET_RELEASE', 'TOP'])
+SOONG_UI_EXE_REL_PATH = 'build/soong/soong_ui.bash'
 
 
 class Error(Exception):
@@ -35,16 +43,54 @@ class BuildFailureError(Error):
     self.return_code = return_code
 
 
-REQUIRED_ENV_VARS = frozenset(['TARGET_PRODUCT', 'TARGET_RELEASE', 'TOP'])
-SOONG_UI_EXE_REL_PATH = 'build/soong/soong_ui.bash'
+class BuildPlanner:
+  """Class in charge of determining how to optimize build targets.
+
+  Given the build context and targets to build it will determine a final list of
+  targets to build along with getting a set of packaging functions to package up
+  any output zip files needed by the build.
+  """
+
+  def __init__(
+      self,
+      build_context: dict[str, any],
+      args: argparse.Namespace,
+      target_optimizations: dict[str, optimized_targets.OptimizedBuildTarget],
+  ):
+    self.build_context = build_context
+    self.args = args
+    self.target_optimizations = target_optimizations
+
+  def create_build_plan(self):
+
+    if 'optimized_build' not in self.build_context['enabled_build_features']:
+      return BuildPlan(set(self.args.extra_targets), set())
+
+    build_targets = set()
+    packaging_functions = set()
+    for target in self.args.extra_targets:
+      target_optimizer_getter = self.target_optimizations.get(target, None)
+      if not target_optimizer_getter:
+        build_targets.add(target)
+        continue
+
+      target_optimizer = target_optimizer_getter(
+          target, self.build_context, self.args
+      )
+      build_targets.update(target_optimizer.get_build_targets())
+      packaging_functions.add(target_optimizer.package_outputs)
+
+    return BuildPlan(build_targets, packaging_functions)
 
 
-def get_top() -> pathlib.Path:
-  return pathlib.Path(os.environ['TOP'])
+@dataclass(frozen=True)
+class BuildPlan:
+  build_targets: set[str]
+  packaging_functions: set[Callable[..., None]]
 
 
 def build_test_suites(argv: list[str]) -> int:
-  """Builds the general-tests and any other test suites passed in.
+  """Builds all test suites passed in, optimizing based on the build_context content.
 
   Args:
     argv: The command line arguments passed in.
@@ -54,9 +100,14 @@ def build_test_suites(argv: list[str]) -> int:
   """
   args = parse_args(argv)
   check_required_env()
+  build_context = load_build_context()
+  build_planner = BuildPlanner(
+      build_context, args, optimized_targets.OPTIMIZED_BUILD_TARGETS
+  )
+  build_plan = build_planner.create_build_plan()
 
   try:
-    build_everything(args)
+    execute_build_plan(build_plan)
   except BuildFailureError as e:
     logging.error('Build command failed! Check build_log for details.')
     return e.return_code
@@ -64,6 +115,16 @@ def build_test_suites(argv: list[str]) -> int:
   return 0
 
 
+def parse_args(argv: list[str]) -> argparse.Namespace:
+  argparser = argparse.ArgumentParser()
+
+  argparser.add_argument(
+      'extra_targets', nargs='*', help='Extra test suites to build.'
+  )
+
+  return argparser.parse_args(argv)
+
+
 def check_required_env():
   """Check for required env vars.
 
@@ -79,43 +140,40 @@ def check_required_env():
   raise Error(f'Missing required environment variables: {t}')
 
 
-def parse_args(argv):
-  argparser = argparse.ArgumentParser()
+def load_build_context():
+  build_context_path = pathlib.Path(os.environ.get('BUILD_CONTEXT', ''))
+  if build_context_path.is_file():
+    try:
+      with open(build_context_path, 'r') as f:
+        return json.load(f)
+    except json.decoder.JSONDecodeError as e:
+      raise Error(f'Failed to load JSON file: {build_context_path}')
 
-  argparser.add_argument(
-      'extra_targets', nargs='*', help='Extra test suites to build.'
-  )
-
-  return argparser.parse_args(argv)
+  logging.info('No BUILD_CONTEXT found, skipping optimizations.')
+  return empty_build_context()
 
 
-def build_everything(args: argparse.Namespace):
-  """Builds all tests (regardless of whether they are needed).
+def empty_build_context():
+  return {'enabled_build_features': []}
 
-  Args:
-    args: The parsed arguments.
 
-  Raises:
-    BuildFailure: If the build command fails.
-  """
-  build_command = base_build_command(args, args.extra_targets)
+def execute_build_plan(build_plan: BuildPlan):
+  build_command = []
+  build_command.append(get_top().joinpath(SOONG_UI_EXE_REL_PATH))
+  build_command.append('--make-mode')
+  build_command.extend(build_plan.build_targets)
 
   try:
     run_command(build_command)
   except subprocess.CalledProcessError as e:
     raise BuildFailureError(e.returncode) from e
 
+  for packaging_function in build_plan.packaging_functions:
+    packaging_function()
 
-def base_build_command(
-    args: argparse.Namespace, extra_targets: set[str]
-) -> list[str]:
 
-  build_command = []
-  build_command.append(get_top().joinpath(SOONG_UI_EXE_REL_PATH))
-  build_command.append('--make-mode')
-  build_command.extend(extra_targets)
-
-  return build_command
+def get_top() -> pathlib.Path:
+  return pathlib.Path(os.environ['TOP'])
 
 
 def run_command(args: list[str], stdout=None):

ci/build_test_suites_test.py

@@ -14,7 +14,9 @@
 
 """Tests for build_test_suites.py"""
 
+import argparse
 from importlib import resources
+import json
 import multiprocessing
 import os
 import pathlib
@@ -27,9 +29,11 @@ import tempfile
 import textwrap
 import time
 from typing import Callable
+import unittest
 from unittest import mock
 import build_test_suites
 import ci_test_lib
+import optimized_targets
 from pyfakefs import fake_filesystem_unittest
 
 
@@ -80,12 +84,20 @@ class BuildTestSuitesTest(fake_filesystem_unittest.TestCase):
     with self.assertRaisesRegex(SystemExit, '42'):
       build_test_suites.main([])
 
+  def test_incorrectly_formatted_build_context_raises(self):
+    build_context = self.fake_top.joinpath('build_context')
+    build_context.touch()
+    os.environ['BUILD_CONTEXT'] = str(build_context)
+
+    with self.assert_raises_word(build_test_suites.Error, 'JSON'):
+      build_test_suites.main([])
+
   def test_build_success_returns(self):
     with self.assertRaisesRegex(SystemExit, '0'):
       build_test_suites.main([])
 
   def assert_raises_word(self, cls, word):
-    return self.assertRaisesRegex(build_test_suites.Error, rf'\b{word}\b')
+    return self.assertRaisesRegex(cls, rf'\b{word}\b')
 
   def _setup_working_build_env(self):
     self.fake_top = pathlib.Path('/fake/top')
@@ -222,6 +234,171 @@ class RunCommandIntegrationTest(ci_test_lib.TestCase):
       os.kill(p.pid, signal.SIGINT)
 
 
+class BuildPlannerTest(unittest.TestCase):
+
+  class TestOptimizedBuildTarget(optimized_targets.OptimizedBuildTarget):
+
+    def __init__(self, output_targets):
+      self.output_targets = output_targets
+
+    def get_build_targets(self):
+      return self.output_targets
+
+    def package_outputs(self):
+      return f'packaging {" ".join(self.output_targets)}'
+
+  def test_build_optimization_off_builds_everything(self):
+    build_targets = {'target_1', 'target_2'}
+    build_planner = self.create_build_planner(
+        build_context=self.create_build_context(optimized_build_enabled=False),
+        build_targets=build_targets,
+    )
+
+    build_plan = build_planner.create_build_plan()
+
+    self.assertSetEqual(build_targets, build_plan.build_targets)
+
+  def test_build_optimization_off_doesnt_package(self):
+    build_targets = {'target_1', 'target_2'}
+    build_planner = self.create_build_planner(
+        build_context=self.create_build_context(optimized_build_enabled=False),
+        build_targets=build_targets,
+    )
+
+    build_plan = build_planner.create_build_plan()
+
+    self.assertEqual(len(build_plan.packaging_functions), 0)
+
+  def test_build_optimization_on_optimizes_target(self):
+    build_targets = {'target_1', 'target_2'}
+    build_planner = self.create_build_planner(
+        build_targets=build_targets,
+        build_context=self.create_build_context(
+            enabled_build_features={self.get_target_flag('target_1')}
+        ),
+    )
+
+    build_plan = build_planner.create_build_plan()
+
+    expected_targets = {self.get_optimized_target_name('target_1'), 'target_2'}
+    self.assertSetEqual(expected_targets, build_plan.build_targets)
+
+  def test_build_optimization_on_packages_target(self):
+    build_targets = {'target_1', 'target_2'}
+    build_planner = self.create_build_planner(
+        build_targets=build_targets,
+        build_context=self.create_build_context(
+            enabled_build_features={self.get_target_flag('target_1')}
+        ),
+    )
+
+    build_plan = build_planner.create_build_plan()
+
+    optimized_target_name = self.get_optimized_target_name('target_1')
+    self.assertIn(
+        f'packaging {optimized_target_name}',
+        self.run_packaging_functions(build_plan),
+    )
+
+  def test_individual_build_optimization_off_doesnt_optimize(self):
+    build_targets = {'target_1', 'target_2'}
+    build_planner = self.create_build_planner(
+        build_targets=build_targets,
+    )
+
+    build_plan = build_planner.create_build_plan()
+
+    self.assertSetEqual(build_targets, build_plan.build_targets)
+
+  def test_individual_build_optimization_off_doesnt_package(self):
+    build_targets = {'target_1', 'target_2'}
+    build_planner = self.create_build_planner(
+        build_targets=build_targets,
+    )
+
+    build_plan = build_planner.create_build_plan()
+
+    expected_packaging_function_outputs = {None, None}
+    self.assertSetEqual(
+        expected_packaging_function_outputs,
+        self.run_packaging_functions(build_plan),
+    )
+
+  def create_build_planner(
+      self,
+      build_targets: set[str],
+      build_context: dict[str, any] = None,
+      args: argparse.Namespace = None,
+      target_optimizations: dict[
+          str, optimized_targets.OptimizedBuildTarget
+      ] = None,
+  ) -> build_test_suites.BuildPlanner:
+    if not build_context:
+      build_context = self.create_build_context()
+    if not args:
+      args = self.create_args(extra_build_targets=build_targets)
+    if not target_optimizations:
+      target_optimizations = self.create_target_optimizations(
+          build_context, build_targets
+      )
+    return build_test_suites.BuildPlanner(
+        build_context, args, target_optimizations
+    )
+
+  def create_build_context(
+      self,
+      optimized_build_enabled: bool = True,
+      enabled_build_features: set[str] = set(),
+      test_context: dict[str, any] = {},
+  ) -> dict[str, any]:
+    build_context = {}
+    build_context['enabled_build_features'] = enabled_build_features
+    if optimized_build_enabled:
+      build_context['enabled_build_features'].add('optimized_build')
+    build_context['test_context'] = test_context
+    return build_context
+
+  def create_args(
+      self, extra_build_targets: set[str] = set()
+  ) -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('extra_targets', nargs='*')
+    return parser.parse_args(extra_build_targets)
+
+  def create_target_optimizations(
+      self, build_context: dict[str, any], build_targets: set[str]
+  ):
+    target_optimizations = dict()
+    for target in build_targets:
+      target_optimizations[target] = (
+          lambda target, build_context, args: optimized_targets.get_target_optimizer(
+              target,
+              self.get_target_flag(target),
+              build_context,
+              self.TestOptimizedBuildTarget(
+                  {self.get_optimized_target_name(target)}
+              ),
+          )
+      )
+
+    return target_optimizations
+
+  def get_target_flag(self, target: str):
+    return f'{target}_enabled'
+
+  def get_optimized_target_name(self, target: str):
+    return f'{target}_optimized'
+
+  def run_packaging_functions(
+      self, build_plan: build_test_suites.BuildPlan
+  ) -> set[str]:
+    output = set()
+    for packaging_function in build_plan.packaging_functions:
+      output.add(packaging_function())
+
+    return output
+
+
 def wait_until(
     condition_function: Callable[[], bool],
     timeout_secs: float = 3.0,

ci/optimized_targets.py

@@ -0,0 +1,69 @@
+#
+# Copyright 2024, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC
+
+
+class OptimizedBuildTarget(ABC):
+  """A representation of an optimized build target.
+
+  This class will determine what targets to build given a given build_cotext and
+  will have a packaging function to generate any necessary output zips for the
+  build.
+  """
+
+  def __init__(self, build_context, args):
+    self.build_context = build_context
+    self.args = args
+
+  def get_build_targets(self):
+    pass
+
+  def package_outputs(self):
+    pass
+
+
+class NullOptimizer(OptimizedBuildTarget):
+  """No-op target optimizer.
+
+  This will simply build the same target it was given and do nothing for the
+  packaging step.
+  """
+
+  def __init__(self, target):
+    self.target = target
+
+  def get_build_targets(self):
+    return {self.target}
+
+  def package_outputs(self):
+    pass
+
+
+def get_target_optimizer(target, enabled_flag, build_context, optimizer):
+  if enabled_flag in build_context['enabled_build_features']:
+    return optimizer
+
+  return NullOptimizer(target)
+
+
+# To be written as:
+#    'target': lambda target, build_context, args: get_target_optimizer(
+#        target,
+#        'target_enabled_flag',
+#        build_context,
+#        TargetOptimizer(build_context, args),
+#    )
+OPTIMIZED_BUILD_TARGETS = dict()