blob: 3c1caf45d91f865972dc2b65568dec8ec702e9c8 [file] [log] [blame]
Mingjun Yangd448be22024-10-15 05:37:00 +00001# 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.
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000014"""Test discovery agent that uses TradeFed to discover test artifacts."""
15import glob
16import json
17import logging
18import os
19import subprocess
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000020
21
Mingjun Yangd448be22024-10-15 05:37:00 +000022class TestDiscoveryAgent:
23 """Test discovery agent."""
24
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000025 _TRADEFED_PREBUILT_JAR_RELATIVE_PATH = (
26 "vendor/google_tradefederation/prebuilts/filegroups/google-tradefed/"
Mingjun Yangd448be22024-10-15 05:37:00 +000027 )
28
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000029 _TRADEFED_NO_POSSIBLE_TEST_DISCOVERY_KEY = "NoPossibleTestDiscovery"
30
31 _TRADEFED_TEST_ZIP_REGEXES_LIST_KEY = "TestZipRegexes"
32
Mingjun Yang3db37902025-01-26 23:41:37 +000033 _TRADEFED_TEST_MODULES_LIST_KEY = "TestModules"
34
35 _TRADEFED_TEST_DEPENDENCIES_LIST_KEY = "TestDependencies"
36
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000037 _TRADEFED_DISCOVERY_OUTPUT_FILE_NAME = "test_discovery_agent.txt"
38
Mingjun Yangd448be22024-10-15 05:37:00 +000039 def __init__(
40 self,
41 tradefed_args: list[str],
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000042 test_mapping_zip_path: str = "",
43 tradefed_jar_revelant_files_path: str = _TRADEFED_PREBUILT_JAR_RELATIVE_PATH,
Mingjun Yangd448be22024-10-15 05:37:00 +000044 ):
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 """
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000055 test_discovery_output_file_name = os.path.join(
Mingjun Yang3db37902025-01-26 23:41:37 +000056 os.environ.get("TOP"), "out", self._TRADEFED_DISCOVERY_OUTPUT_FILE_NAME
Mingjun Yangd2cc6a62024-11-05 23:44:38 +000057 )
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]
Mingjun Yangd448be22024-10-15 05:37:00 +000095
Mingjun Yang3db37902025-01-26 23:41:37 +000096 def discover_test_mapping_test_modules(self) -> (list[str], list[str]):
97 """Discover test mapping test modules and dependencies from TradeFed.
Mingjun Yangd448be22024-10-15 05:37:00 +000098
99 Returns:
Mingjun Yang3db37902025-01-26 23:41:37 +0000100 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
Mingjun Yangd448be22024-10-15 05:37:00 +0000102 TradeFed test args.
103 """
Mingjun Yang3db37902025-01-26 23:41:37 +0000104 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()
Julien Desprezbe518142025-03-21 11:31:01 -0700121 env.update({"SKIP_JAVA_QUERY": "1"})
122 env.update({"ALLOW_EMPTY_TEST_MAPPING": "1"})
Mingjun Yang3db37902025-01-26 23:41:37 +0000123 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:
Julien Desprezbe518142025-03-21 11:31:01 -0700127 result = subprocess.run(args=java_args, env=env, text=True, check=True, stdout = subprocess.PIPE,
128 stderr = subprocess.PIPE)
Mingjun Yang3db37902025-01-26 23:41:37 +0000129 logging.info(f"Test discovery agent output: {result.stdout}")
130 except subprocess.CalledProcessError as e:
131 raise TestDiscoveryError(
Julien Desprezbe518142025-03-21 11:31:01 -0700132 f"Failed to run test discovery, stdout: {e.stdout}, stderr:"
Mingjun Yang3db37902025-01-26 23:41:37 +0000133 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 )
Mingjun Yangd2cc6a62024-11-05 23:44:38 +0000151
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