From 70a53bd6cf801d2aeb60e168e9ddf650d9126e5b Mon Sep 17 00:00:00 2001 From: Luca Farsi Date: Wed, 7 Aug 2024 17:29:16 -0700 Subject: [PATCH] Refactor OptimizedBuildTarget and add general-tests optimization Refactor the OptimizedBuildTarget base class to more simply handle enabling/disabling build optimizations based on flags. Now all optimizers (aside from the no-op NullOptimizer will have their enabling built into the base class, and also necessitate their get_build_targets and package_outputs functions to be defined in order to not error out. Add the GeneralTestsOptimizer class as well, unimplemented for now. Test: atest build_test_suites_test Bug: 358215235 Change-Id: Ide22c64fc238d754db8d8e76484947401e54e680 --- ci/build_test_suites_test.py | 67 ++++++++++++++++++----------------- ci/optimized_targets.py | 68 +++++++++++++++++++++++++++--------- 2 files changed, 87 insertions(+), 48 deletions(-) diff --git a/ci/build_test_suites_test.py b/ci/build_test_suites_test.py index 7efb154487..25c072e2b0 100644 --- a/ci/build_test_suites_test.py +++ b/ci/build_test_suites_test.py @@ -15,6 +15,7 @@ """Tests for build_test_suites.py""" import argparse +import functools from importlib import resources import json import multiprocessing @@ -238,14 +239,21 @@ class BuildPlannerTest(unittest.TestCase): class TestOptimizedBuildTarget(optimized_targets.OptimizedBuildTarget): - def __init__(self, output_targets): + def __init__( + self, target, build_context, args, output_targets, packaging_outputs + ): + super().__init__(target, build_context, args) self.output_targets = output_targets + self.packaging_outputs = packaging_outputs - def get_build_targets(self): + def get_build_targets_impl(self): return self.output_targets - def package_outputs(self): - return f'packaging {" ".join(self.output_targets)}' + def package_outputs_impl(self): + self.packaging_outputs.add(f'packaging {" ".join(self.output_targets)}') + + def get_enabled_flag(self): + return f'{self.target}_enabled' def test_build_optimization_off_builds_everything(self): build_targets = {'target_1', 'target_2'} @@ -285,20 +293,20 @@ class BuildPlannerTest(unittest.TestCase): def test_build_optimization_on_packages_target(self): build_targets = {'target_1', 'target_2'} + packaging_outputs = set() 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')} + enabled_build_features={self.get_target_flag('target_1')}, ), + packaging_outputs=packaging_outputs, ) build_plan = build_planner.create_build_plan() + self.run_packaging_functions(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), - ) + self.assertIn(f'packaging {optimized_target_name}', packaging_outputs) def test_individual_build_optimization_off_doesnt_optimize(self): build_targets = {'target_1', 'target_2'} @@ -312,17 +320,16 @@ class BuildPlannerTest(unittest.TestCase): def test_individual_build_optimization_off_doesnt_package(self): build_targets = {'target_1', 'target_2'} + packaging_outputs = set() build_planner = self.create_build_planner( build_targets=build_targets, + packaging_outputs=packaging_outputs, ) build_plan = build_planner.create_build_plan() + self.run_packaging_functions(build_plan) - expected_packaging_function_outputs = {None, None} - self.assertSetEqual( - expected_packaging_function_outputs, - self.run_packaging_functions(build_plan), - ) + self.assertFalse(packaging_outputs) def test_target_output_used_target_built(self): build_target = 'test_target' @@ -381,6 +388,7 @@ class BuildPlannerTest(unittest.TestCase): target_optimizations: dict[ str, optimized_targets.OptimizedBuildTarget ] = None, + packaging_outputs: set[str] = set(), ) -> build_test_suites.BuildPlanner: if not build_context: build_context = self.create_build_context() @@ -388,7 +396,9 @@ class BuildPlannerTest(unittest.TestCase): args = self.create_args(extra_build_targets=build_targets) if not target_optimizations: target_optimizations = self.create_target_optimizations( - build_context, build_targets + build_context, + build_targets, + packaging_outputs, ) return build_test_suites.BuildPlanner( build_context, args, target_optimizations @@ -415,19 +425,17 @@ class BuildPlannerTest(unittest.TestCase): return parser.parse_args(extra_build_targets) def create_target_optimizations( - self, build_context: dict[str, any], build_targets: set[str] + self, + build_context: dict[str, any], + build_targets: set[str], + packaging_outputs: set[str] = set(), ): 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)} - ), - ) + target_optimizations[target] = functools.partial( + self.TestOptimizedBuildTarget, + output_targets={self.get_optimized_target_name(target)}, + packaging_outputs=packaging_outputs, ) return target_optimizations @@ -438,14 +446,9 @@ class BuildPlannerTest(unittest.TestCase): 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() + def run_packaging_functions(self, build_plan: build_test_suites.BuildPlan): for packaging_function in build_plan.packaging_functions: - output.add(packaging_function()) - - return output + packaging_function() def get_test_context(self, target: str): return { diff --git a/ci/optimized_targets.py b/ci/optimized_targets.py index 632171906d..8a529c7420 100644 --- a/ci/optimized_targets.py +++ b/ci/optimized_targets.py @@ -14,6 +14,9 @@ # limitations under the License. from abc import ABC +from typing import Self +import argparse +import functools class OptimizedBuildTarget(ABC): @@ -24,15 +27,41 @@ class OptimizedBuildTarget(ABC): build. """ - def __init__(self, build_context, args): + def __init__( + self, + target: str, + build_context: dict[str, any], + args: argparse.Namespace, + ): + self.target = target self.build_context = build_context self.args = args - def get_build_targets(self): - pass + def get_build_targets(self) -> set[str]: + features = self.build_context.get('enabledBuildFeatures', []) + if self.get_enabled_flag() in features: + return self.get_build_targets_impl() + return {self.target} def package_outputs(self): - pass + features = self.build_context.get('enabledBuildFeatures', []) + if self.get_enabled_flag() in features: + return self.package_outputs_impl() + + def package_outputs_impl(self): + raise NotImplementedError( + f'package_outputs_impl not implemented in {type(self).__name__}' + ) + + def get_enabled_flag(self): + raise NotImplementedError( + f'get_enabled_flag not implemented in {type(self).__name__}' + ) + + def get_build_targets_impl(self) -> set[str]: + raise NotImplementedError( + f'get_build_targets_impl not implemented in {type(self).__name__}' + ) class NullOptimizer(OptimizedBuildTarget): @@ -52,18 +81,25 @@ class NullOptimizer(OptimizedBuildTarget): pass -def get_target_optimizer(target, enabled_flag, build_context, optimizer): - if enabled_flag in build_context['enabledBuildFeatures']: - return optimizer +class GeneralTestsOptimizer(OptimizedBuildTarget): + """general-tests optimizer - return NullOptimizer(target) + 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 + builds those modules and packages them in the same way general-tests.zip is + normally built. + """ + + def get_enabled_flag(self): + return 'general-tests-optimized' + + @classmethod + def get_optimized_targets(cls) -> dict[str, OptimizedBuildTarget]: + return {'general-tests': functools.partial(cls)} -# 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() +OPTIMIZED_BUILD_TARGETS = {} +OPTIMIZED_BUILD_TARGETS.update(GeneralTestsOptimizer.get_optimized_targets())