• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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_TEST_MODULES_LIST_KEY = "TestModules"
34
35  _TRADEFED_TEST_DEPENDENCIES_LIST_KEY = "TestDependencies"
36
37  _TRADEFED_DISCOVERY_OUTPUT_FILE_NAME = "test_discovery_agent.txt"
38
39  def __init__(
40      self,
41      tradefed_args: list[str],
42      test_mapping_zip_path: str = "",
43      tradefed_jar_revelant_files_path: str = _TRADEFED_PREBUILT_JAR_RELATIVE_PATH,
44  ):
45    self.tradefed_args = tradefed_args
46    self.test_mapping_zip_path = test_mapping_zip_path
47    self.tradefed_jar_relevant_files_path = tradefed_jar_revelant_files_path
48
49  def discover_test_zip_regexes(self) -> list[str]:
50    """Discover test zip regexes from TradeFed.
51
52    Returns:
53      A list of test zip regexes that TF is going to try to pull files from.
54    """
55    test_discovery_output_file_name = os.path.join(
56        os.environ.get("TOP"), "out", self._TRADEFED_DISCOVERY_OUTPUT_FILE_NAME
57    )
58    with open(
59        test_discovery_output_file_name, mode="w+t"
60    ) as test_discovery_output_file:
61      java_args = []
62      java_args.append("prebuilts/jdk/jdk21/linux-x86/bin/java")
63      java_args.append("-cp")
64      java_args.append(
65          self.create_classpath(self.tradefed_jar_relevant_files_path)
66      )
67      java_args.append(
68          "com.android.tradefed.observatory.TestZipDiscoveryExecutor"
69      )
70      java_args.extend(self.tradefed_args)
71      env = os.environ.copy()
72      env.update({"DISCOVERY_OUTPUT_FILE": test_discovery_output_file.name})
73      logging.info(f"Calling test discovery with args: {java_args}")
74      try:
75        result = subprocess.run(args=java_args, env=env, text=True, check=True)
76        logging.info(f"Test zip discovery output: {result.stdout}")
77      except subprocess.CalledProcessError as e:
78        raise TestDiscoveryError(
79            f"Failed to run test discovery, strout: {e.stdout}, strerr:"
80            f" {e.stderr}, returncode: {e.returncode}"
81        )
82      data = json.loads(test_discovery_output_file.read())
83      logging.info(f"Test discovery result file content: {data}")
84      if (
85          self._TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY in data
86          and data[self._TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY]
87      ):
88        raise TestDiscoveryError("No possible test discovery")
89      if (
90          data[self._TRADEFED_TEST_ZIP_REGEXES_LIST_KEY] is None
91          or data[self._TRADEFED_TEST_ZIP_REGEXES_LIST_KEY] is []
92      ):
93        raise TestDiscoveryError("No test zip regexes returned")
94      return data[self._TRADEFED_TEST_ZIP_REGEXES_LIST_KEY]
95
96  def discover_test_mapping_test_modules(self) -> (list[str], list[str]):
97    """Discover test mapping test modules and dependencies from TradeFed.
98
99    Returns:
100      A tuple that contains a list of test modules and a list of test
101      dependencies that TradeFed is going to execute based on the
102      TradeFed test args.
103    """
104    test_discovery_output_file_name = os.path.join(
105        os.environ.get("TOP"), "out", self._TRADEFED_DISCOVERY_OUTPUT_FILE_NAME
106    )
107    with open(
108        test_discovery_output_file_name, mode="w+t"
109    ) as test_discovery_output_file:
110      java_args = []
111      java_args.append("prebuilts/jdk/jdk21/linux-x86/bin/java")
112      java_args.append("-cp")
113      java_args.append(
114          self.create_classpath(self.tradefed_jar_relevant_files_path)
115      )
116      java_args.append(
117          "com.android.tradefed.observatory.TestMappingDiscoveryAgent"
118      )
119      java_args.extend(self.tradefed_args)
120      env = os.environ.copy()
121      env.update({"SKIP_JAVA_QUERY": "1"})
122      env.update({"ALLOW_EMPTY_TEST_MAPPING": "1"})
123      env.update({"TF_TEST_MAPPING_ZIP_FILE": self.test_mapping_zip_path})
124      env.update({"DISCOVERY_OUTPUT_FILE": test_discovery_output_file.name})
125      logging.info(f"Calling test discovery with args: {java_args}")
126      try:
127        result = subprocess.run(args=java_args, env=env, text=True, check=True, stdout = subprocess.PIPE,
128    stderr = subprocess.PIPE)
129        logging.info(f"Test discovery agent output: {result.stdout}")
130      except subprocess.CalledProcessError as e:
131        raise TestDiscoveryError(
132            f"Failed to run test discovery, stdout: {e.stdout}, stderr:"
133            f" {e.stderr}, returncode: {e.returncode}"
134        )
135      data = json.loads(test_discovery_output_file.read())
136      logging.info(f"Test discovery result file content: {data}")
137      if (
138          self._TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY in data
139          and data[self._TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY]
140      ):
141        raise TestDiscoveryError("No possible test discovery")
142      if (
143          data[self._TRADEFED_TEST_MODULES_LIST_KEY] is None
144          or data[self._TRADEFED_TEST_MODULES_LIST_KEY] is []
145      ):
146        raise TestDiscoveryError("No test modules returned")
147      return (
148          data[self._TRADEFED_TEST_MODULES_LIST_KEY],
149          data[self._TRADEFED_TEST_DEPENDENCIES_LIST_KEY],
150      )
151
152  def create_classpath(self, directory):
153    """Creates a classpath string from all .jar files in the given directory.
154
155    Args:
156      directory: The directory to search for .jar files.
157
158    Returns:
159      A string representing the classpath, with jar files separated by the
160      OS-specific path separator (e.g., ':' on Linux/macOS, ';' on Windows).
161    """
162    jar_files = glob.glob(os.path.join(directory, "*.jar"))
163    return os.pathsep.join(jar_files)
164
165
166class TestDiscoveryError(Exception):
167  """A TestDiscoveryErrorclass."""
168
169  def __init__(self, message):
170    super().__init__(message)
171    self.message = message
172