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
diff --git a/ci/build_test_suites_test.py b/ci/build_test_suites_test.py
index 7efb154..25c072e 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 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 @@
 
   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 @@
 
   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 @@
       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 @@
       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 @@
     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 @@
   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 6321719..8a529c7 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 @@
   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 @@
     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())