add BuildContext class and fix enabled features

Add a BuildContext class to simplify parsing the build context dict, and
fix the parsing of enabledBuildFeatures, which was being parsed as a
list of strings when it's actually a list of dicts like:
[{'name': '<feature_name>'}]

Test: atest build_test_suites_test
Bug: 361605425
Change-Id: I6424c444daf1582e92313c39f43207cb274aa78f

ci/Android.bp

@@ -76,6 +76,7 @@ python_library_host {
     srcs: [
         "build_test_suites.py",
         "optimized_targets.py",
+        "build_context.py",
     ],
 }
 

ci/build_context.py

@@ -0,0 +1,64 @@
+# 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.
+
+"""Container class for build context with utility functions."""
+
+import re
+
+
+class BuildContext:
+
+  def __init__(self, build_context_dict: dict[str, any]):
+    self.enabled_build_features = set()
+    for opt in build_context_dict.get('enabledBuildFeatures', []):
+      self.enabled_build_features.add(opt.get('name'))
+    self.test_infos = set()
+    for test_info_dict in build_context_dict.get('testContext', dict()).get(
+        'testInfos', []
+    ):
+      self.test_infos.add(self.TestInfo(test_info_dict))
+
+  def build_target_used(self, target: str) -> bool:
+    return any(test.build_target_used(target) for test in self.test_infos)
+
+  class TestInfo:
+
+    _DOWNLOAD_OPTS = {
+        'test-config-only-zip',
+        'test-zip-file-filter',
+        'extra-host-shared-lib-zip',
+        'sandbox-tests-zips',
+        'additional-files-filter',
+        'cts-package-name',
+    }
+
+    def __init__(self, test_info_dict: dict[str, any]):
+      self.is_test_mapping = False
+      self.test_mapping_test_groups = set()
+      self.file_download_options = set()
+      for opt in test_info_dict.get('extraOptions', []):
+        key = opt.get('key')
+        if key == 'test-mapping-test-group':
+          self.is_test_mapping = True
+          self.test_mapping_test_groups.update(opt.get('values', set()))
+
+        if key in self._DOWNLOAD_OPTS:
+          self.file_download_options.update(opt.get('values', set()))
+
+    def build_target_used(self, target: str) -> bool:
+      # For all of a targets' outputs, check if any of the regexes used by tests
+      # to download artifacts would match it. If any of them do then this target
+      # is necessary.
+      regex = r'\b(%s)\b' % re.escape(target)
+      return any(re.search(regex, opt) for opt in self.file_download_options)

ci/build_test_suites.py

@@ -24,6 +24,7 @@ import re
 import subprocess
 import sys
 from typing import Callable
+from build_context import BuildContext
 import optimized_targets
 
 
@@ -53,18 +54,9 @@ class BuildPlanner:
   any output zip files needed by the build.
   """
 
-  _DOWNLOAD_OPTS = {
-      'test-config-only-zip',
-      'test-zip-file-filter',
-      'extra-host-shared-lib-zip',
-      'sandbox-tests-zips',
-      'additional-files-filter',
-      'cts-package-name',
-  }
-
   def __init__(
       self,
-      build_context: dict[str, any],
+      build_context: BuildContext,
       args: argparse.Namespace,
       target_optimizations: dict[str, optimized_targets.OptimizedBuildTarget],
   ):
@@ -74,18 +66,15 @@ class BuildPlanner:
 
   def create_build_plan(self):
 
-    if 'optimized_build' not in self.build_context.get(
+    if 'optimized_build' not in self.build_context.enabled_build_features:
-        'enabledBuildFeatures', []
-    ):
       return BuildPlan(set(self.args.extra_targets), set())
 
     build_targets = set()
     packaging_functions = set()
-    self.file_download_options = self._aggregate_file_download_options()
     for target in self.args.extra_targets:
       if self._unused_target_exclusion_enabled(
           target
-      ) and not self._build_target_used(target):
+      ) and not self.build_context.build_target_used(target):
         continue
 
       target_optimizer_getter = self.target_optimizations.get(target, None)
@@ -102,34 +91,11 @@ class BuildPlanner:
     return BuildPlan(build_targets, packaging_functions)
 
   def _unused_target_exclusion_enabled(self, target: str) -> bool:
-    return f'{target}_unused_exclusion' in self.build_context.get(
+    return (
-        'enabledBuildFeatures', []
+        f'{target}_unused_exclusion'
+        in self.build_context.enabled_build_features
     )
 
-  def _build_target_used(self, target: str) -> bool:
-    """Determines whether this target's outputs are used by the test configurations listed in the build context."""
-    # For all of a targets' outputs, check if any of the regexes used by tests
-    # to download artifacts would match it. If any of them do then this target
-    # is necessary.
-    regex = r'\b(%s)\b' % re.escape(target)
-    return any(re.search(regex, opt) for opt in self.file_download_options)
-
-  def _aggregate_file_download_options(self) -> set[str]:
-    """Lists out all test config options to specify targets to download.
-
-    These come in the form of regexes.
-    """
-    all_options = set()
-    for test_info in self._get_test_infos():
-      for opt in test_info.get('extraOptions', []):
-        # check the known list of options for downloading files.
-        if opt.get('key') in self._DOWNLOAD_OPTS:
-          all_options.update(opt.get('values', []))
-    return all_options
-
-  def _get_test_infos(self):
-    return self.build_context.get('testContext', dict()).get('testInfos', [])
-
 
 @dataclass(frozen=True)
 class BuildPlan:
@@ -148,7 +114,7 @@ def build_test_suites(argv: list[str]) -> int:
   """
   args = parse_args(argv)
   check_required_env()
-  build_context = load_build_context()
+  build_context = BuildContext(load_build_context())
   build_planner = BuildPlanner(
       build_context, args, optimized_targets.OPTIMIZED_BUILD_TARGETS
   )

ci/build_test_suites_test.py

@@ -32,6 +32,7 @@ import time
 from typing import Callable
 import unittest
 from unittest import mock
+from build_context import BuildContext
 import build_test_suites
 import ci_test_lib
 import optimized_targets
@@ -282,7 +283,7 @@ class BuildPlannerTest(unittest.TestCase):
     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=[{'name': self.get_target_flag('target_1')}]
         ),
     )
 
@@ -297,7 +298,7 @@ class BuildPlannerTest(unittest.TestCase):
     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=[{'name': self.get_target_flag('target_1')}]
         ),
         packaging_outputs=packaging_outputs,
     )
@@ -337,7 +338,7 @@ class BuildPlannerTest(unittest.TestCase):
         build_targets={build_target},
         build_context=self.create_build_context(
             test_context=self.get_test_context(build_target),
-            enabled_build_features={'test_target_unused_exclusion'},
+            enabled_build_features=[{'name': 'test_target_unused_exclusion'}],
         ),
     )
 
@@ -356,7 +357,7 @@ class BuildPlannerTest(unittest.TestCase):
         build_targets={build_target},
         build_context=self.create_build_context(
             test_context=test_context,
-            enabled_build_features={'test_target_unused_exclusion'},
+            enabled_build_features=[{'name': 'test_target_unused_exclusion'}],
         ),
     )
 
@@ -372,7 +373,7 @@ class BuildPlannerTest(unittest.TestCase):
         build_targets={build_target},
         build_context=self.create_build_context(
             test_context=test_context,
-            enabled_build_features={'test_target_unused_exclusion'},
+            enabled_build_features=[{'name': 'test_target_unused_exclusion'}],
         ),
     )
 
@@ -391,7 +392,7 @@ class BuildPlannerTest(unittest.TestCase):
         build_targets={build_target},
         build_context=self.create_build_context(
             test_context=test_context,
-            enabled_build_features={'test_target_unused_exclusion'},
+            enabled_build_features=[{'name': 'test_target_unused_exclusion'}],
         ),
     )
 
@@ -402,7 +403,7 @@ class BuildPlannerTest(unittest.TestCase):
   def create_build_planner(
       self,
       build_targets: set[str],
-      build_context: dict[str, any] = None,
+      build_context: BuildContext = None,
       args: argparse.Namespace = None,
       target_optimizations: dict[
           str, optimized_targets.OptimizedBuildTarget
@@ -426,15 +427,17 @@ class BuildPlannerTest(unittest.TestCase):
   def create_build_context(
       self,
       optimized_build_enabled: bool = True,
-      enabled_build_features: set[str] = set(),
+      enabled_build_features: list[dict[str, str]] = [],
       test_context: dict[str, any] = {},
-  ) -> dict[str, any]:
+  ) -> BuildContext:
-    build_context = {}
+    build_context_dict = {}
-    build_context['enabledBuildFeatures'] = enabled_build_features
+    build_context_dict['enabledBuildFeatures'] = enabled_build_features
     if optimized_build_enabled:
-      build_context['enabledBuildFeatures'].add('optimized_build')
+      build_context_dict['enabledBuildFeatures'].append(
-    build_context['testContext'] = test_context
+          {'name': 'optimized_build'}
-    return build_context
+      )
+    build_context_dict['testContext'] = test_context
+    return BuildContext(build_context_dict)
 
   def create_args(
       self, extra_build_targets: set[str] = set()
@@ -445,7 +448,7 @@ class BuildPlannerTest(unittest.TestCase):
 
   def create_target_optimizations(
       self,
-      build_context: dict[str, any],
+      build_context: BuildContext,
       build_targets: set[str],
       packaging_outputs: set[str] = set(),
   ):

ci/optimized_targets.py

@@ -14,9 +14,10 @@
 # limitations under the License.
 
 from abc import ABC
-from typing import Self
 import argparse
 import functools
+from typing import Self
+from build_context import BuildContext
 
 
 class OptimizedBuildTarget(ABC):
@@ -30,7 +31,7 @@ class OptimizedBuildTarget(ABC):
   def __init__(
       self,
       target: str,
-      build_context: dict[str, any],
+      build_context: BuildContext,
       args: argparse.Namespace,
   ):
     self.target = target
@@ -38,13 +39,13 @@ class OptimizedBuildTarget(ABC):
     self.args = args
 
   def get_build_targets(self) -> set[str]:
-    features = self.build_context.get('enabledBuildFeatures', [])
+    features = self.build_context.enabled_build_features
     if self.get_enabled_flag() in features:
       return self.get_build_targets_impl()
     return {self.target}
 
   def package_outputs(self):
-    features = self.build_context.get('enabledBuildFeatures', [])
+    features = self.build_context.enabled_build_features
     if self.get_enabled_flag() in features:
       return self.package_outputs_impl()