xref: /aosp_15_r20/build/make/ci/test_discovery_agent.py (revision 9e94795a3d4ef5c1d47486f9a02bb378756cea8a)
1# Copyright 2024, The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Test discovery agent that uses TradeFed to discover test artifacts."""
15import glob
16import json
17import logging
18import os
19import subprocess
20
21
22class TestDiscoveryAgent:
23  """Test discovery agent."""
24
25  _TRADEFED_PREBUILT_JAR_RELATIVE_PATH = (
26      "vendor/google_tradefederation/prebuilts/filegroups/google-tradefed/"
27  )
28
29  _TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY = "NoPossibleTestDiscovery"
30
31  _TRADEFED_TEST_ZIP_REGEXES_LIST_KEY = "TestZipRegexes"
32
33  _TRADEFED_DISCOVERY_OUTPUT_FILE_NAME = "test_discovery_agent.txt"
34
35  def __init__(
36      self,
37      tradefed_args: list[str],
38      test_mapping_zip_path: str = "",
39      tradefed_jar_revelant_files_path: str = _TRADEFED_PREBUILT_JAR_RELATIVE_PATH,
40  ):
41    self.tradefed_args = tradefed_args
42    self.test_mapping_zip_path = test_mapping_zip_path
43    self.tradefed_jar_relevant_files_path = tradefed_jar_revelant_files_path
44
45  def discover_test_zip_regexes(self) -> list[str]:
46    """Discover test zip regexes from TradeFed.
47
48    Returns:
49      A list of test zip regexes that TF is going to try to pull files from.
50    """
51    test_discovery_output_file_name = os.path.join(
52        os.environ.get('TOP'), 'out', self._TRADEFED_DISCOVERY_OUTPUT_FILE_NAME
53    )
54    with open(
55        test_discovery_output_file_name, mode="w+t"
56    ) as test_discovery_output_file:
57      java_args = []
58      java_args.append("prebuilts/jdk/jdk21/linux-x86/bin/java")
59      java_args.append("-cp")
60      java_args.append(
61          self.create_classpath(self.tradefed_jar_relevant_files_path)
62      )
63      java_args.append(
64          "com.android.tradefed.observatory.TestZipDiscoveryExecutor"
65      )
66      java_args.extend(self.tradefed_args)
67      env = os.environ.copy()
68      env.update({"DISCOVERY_OUTPUT_FILE": test_discovery_output_file.name})
69      logging.info(f"Calling test discovery with args: {java_args}")
70      try:
71        result = subprocess.run(args=java_args, env=env, text=True, check=True)
72        logging.info(f"Test zip discovery output: {result.stdout}")
73      except subprocess.CalledProcessError as e:
74        raise TestDiscoveryError(
75            f"Failed to run test discovery, strout: {e.stdout}, strerr:"
76            f" {e.stderr}, returncode: {e.returncode}"
77        )
78      data = json.loads(test_discovery_output_file.read())
79      logging.info(f"Test discovery result file content: {data}")
80      if (
81          self._TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY in data
82          and data[self._TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY]
83      ):
84        raise TestDiscoveryError("No possible test discovery")
85      if (
86          data[self._TRADEFED_TEST_ZIP_REGEXES_LIST_KEY] is None
87          or data[self._TRADEFED_TEST_ZIP_REGEXES_LIST_KEY] is []
88      ):
89        raise TestDiscoveryError("No test zip regexes returned")
90      return data[self._TRADEFED_TEST_ZIP_REGEXES_LIST_KEY]
91
92  def discover_test_modules(self) -> list[str]:
93    """Discover test modules from TradeFed.
94
95    Returns:
96      A list of test modules that TradeFed is going to execute based on the
97      TradeFed test args.
98    """
99    return []
100
101  def create_classpath(self, directory):
102    """Creates a classpath string from all .jar files in the given directory.
103
104    Args:
105      directory: The directory to search for .jar files.
106
107    Returns:
108      A string representing the classpath, with jar files separated by the
109      OS-specific path separator (e.g., ':' on Linux/macOS, ';' on Windows).
110    """
111    jar_files = glob.glob(os.path.join(directory, "*.jar"))
112    return os.pathsep.join(jar_files)
113
114
115class TestDiscoveryError(Exception):
116  """A TestDiscoveryErrorclass."""
117
118  def __init__(self, message):
119    super().__init__(message)
120    self.message = message
121