1# Copyright (C) 2021 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
15"""Rules used to run tests using Tradefed."""
16
17load("//bazel/rules:platform_transitions.bzl", "device_transition", "host_transition")
18load("//bazel/rules:tradefed_test_aspects.bzl", "soong_prebuilt_tradefed_test_aspect")
19load("//bazel/rules:tradefed_test_dependency_info.bzl", "TradefedTestDependencyInfo")
20load("//bazel/rules:common_settings.bzl", "BuildSettingInfo")
21load(
22    "//:constants.bzl",
23    "aapt2_label",
24    "aapt_label",
25    "adb_label",
26    "atest_script_help_sh_label",
27    "atest_tradefed_label",
28    "atest_tradefed_sh_label",
29    "bazel_result_reporter_label",
30    "compatibility_tradefed_label",
31    "tradefed_label",
32    "tradefed_test_framework_label",
33    "vts_core_tradefed_harness_label",
34)
35load("//bazel/rules:device_test.bzl", "device_test")
36
37TradefedTestInfo = provider(
38    doc = "Info about a Tradefed test module",
39    fields = {
40        "module_name": "Name of the original Tradefed test module",
41    },
42)
43
44_BAZEL_WORK_DIR = "${TEST_SRCDIR}/${TEST_WORKSPACE}/"
45_PY_TOOLCHAIN = "@bazel_tools//tools/python:toolchain_type"
46_JAVA_TOOLCHAIN = "@bazel_tools//tools/jdk:runtime_toolchain_type"
47_TOOLCHAINS = [_PY_TOOLCHAIN, _JAVA_TOOLCHAIN]
48
49_TRADEFED_TEST_ATTRIBUTES = {
50    "module_name": attr.string(),
51    "_tradefed_test_template": attr.label(
52        default = "//bazel/rules:tradefed_test.sh.template",
53        allow_single_file = True,
54    ),
55    "_tradefed_classpath_jars": attr.label_list(
56        default = [
57            atest_tradefed_label,
58            tradefed_label,
59            tradefed_test_framework_label,
60            bazel_result_reporter_label,
61        ],
62        cfg = host_transition,
63        aspects = [soong_prebuilt_tradefed_test_aspect],
64    ),
65    "_atest_tradefed_launcher": attr.label(
66        default = atest_tradefed_sh_label,
67        allow_single_file = True,
68        cfg = host_transition,
69        aspects = [soong_prebuilt_tradefed_test_aspect],
70    ),
71    "_atest_helper": attr.label(
72        default = atest_script_help_sh_label,
73        allow_single_file = True,
74        cfg = host_transition,
75        aspects = [soong_prebuilt_tradefed_test_aspect],
76    ),
77    "_adb": attr.label(
78        default = adb_label,
79        allow_single_file = True,
80        cfg = host_transition,
81        aspects = [soong_prebuilt_tradefed_test_aspect],
82    ),
83    "_extra_tradefed_result_reporters": attr.label(
84        default = "//bazel/rules:extra_tradefed_result_reporters",
85    ),
86    # This attribute is required to use Starlark transitions. It allows
87    # allowlisting usage of this rule. For more information, see
88    # https://docs.bazel.build/versions/master/skylark/config.html#user-defined-transitions
89    "_allowlist_function_transition": attr.label(
90        default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
91    ),
92}
93
94def _add_dicts(*dictionaries):
95    """Creates a new `dict` that has all the entries of the given dictionaries.
96
97    This function serves as a replacement for the `+` operator which does not
98    work with dictionaries. The implementation is inspired by Skylib's
99    `dict.add` and duplicated to avoid the dependency. See
100    https://github.com/bazelbuild/bazel/issues/6461 for more details.
101
102    Note, if the same key is present in more than one of the input dictionaries,
103    the last of them in the argument list overrides any earlier ones.
104
105    Args:
106        *dictionaries: Dictionaries to be added.
107
108    Returns:
109        A new `dict` that has all the entries of the given dictionaries.
110    """
111    result = {}
112    for d in dictionaries:
113        result.update(d)
114    return result
115
116def _tradefed_deviceless_test_impl(ctx):
117    return _tradefed_test_impl(
118        ctx,
119        tradefed_options = [
120            "-n",
121            "--prioritize-host-config",
122            "--skip-host-arch-check",
123        ],
124        test_host_deps = ctx.attr.test,
125    )
126
127tradefed_deviceless_test = rule(
128    attrs = _add_dicts(
129        _TRADEFED_TEST_ATTRIBUTES,
130        {
131            "test": attr.label(
132                mandatory = True,
133                cfg = host_transition,
134                aspects = [soong_prebuilt_tradefed_test_aspect],
135            ),
136        },
137    ),
138    test = True,
139    implementation = _tradefed_deviceless_test_impl,
140    toolchains = _TOOLCHAINS,
141    doc = "A rule used to run host-side deviceless tests using Tradefed",
142)
143
144def _tradefed_robolectric_test_impl(ctx):
145    def add_android_all_files(ctx, tradefed_test_dir):
146        android_all_files = []
147        for target in ctx.attr._android_all:
148            for f in target.files.to_list():
149                # Tradefed expects a flat `android-all` directory structure for
150                # Robolectric tests.
151                symlink = _symlink(ctx, f, "%s/android-all/%s" % (tradefed_test_dir, f.basename))
152                android_all_files.append(symlink)
153        return android_all_files
154
155    return _tradefed_test_impl(
156        ctx,
157        data = [ctx.attr.jdk],
158        tradefed_options = [
159            "-n",
160            "--prioritize-host-config",
161            "--skip-host-arch-check",
162            "--test-arg",
163            "com.android.tradefed.testtype.IsolatedHostTest:java-folder:%s" % ctx.attr.jdk.label.package,
164        ],
165        test_host_deps = ctx.attr.test,
166        add_extra_tradefed_test_files = add_android_all_files,
167    )
168
169tradefed_robolectric_test = rule(
170    attrs = _add_dicts(
171        _TRADEFED_TEST_ATTRIBUTES,
172        {
173            "test": attr.label(
174                mandatory = True,
175                cfg = host_transition,
176                aspects = [soong_prebuilt_tradefed_test_aspect],
177            ),
178            "jdk": attr.label(
179                mandatory = True,
180            ),
181            "_android_all": attr.label_list(
182                default = ["//android-all:android-all"],
183            ),
184        },
185    ),
186    test = True,
187    implementation = _tradefed_robolectric_test_impl,
188    toolchains = _TOOLCHAINS,
189    doc = "A rule used to run Robolectric tests using Tradefed",
190)
191
192def _tradefed_device_test_impl(ctx):
193    tradefed_deps = []
194    tradefed_deps.extend(ctx.attr._aapt)
195    tradefed_deps.extend(ctx.attr._aapt2)
196    tradefed_deps.extend(ctx.attr.tradefed_deps)
197
198    test_device_deps = []
199    test_host_deps = []
200
201    if ctx.attr.host_test:
202        test_host_deps.extend(ctx.attr.host_test)
203    if ctx.attr.device_test:
204        test_device_deps.extend(ctx.attr.device_test)
205
206    return _tradefed_test_impl(
207        ctx,
208        tradefed_deps = tradefed_deps,
209        test_device_deps = test_device_deps,
210        test_host_deps = test_host_deps,
211        path_additions = [
212            _BAZEL_WORK_DIR + ctx.file._aapt.dirname,
213            _BAZEL_WORK_DIR + ctx.file._aapt2.dirname,
214        ],
215    )
216
217_tradefed_device_test = rule(
218    attrs = _add_dicts(
219        _TRADEFED_TEST_ATTRIBUTES,
220        {
221            "device_test": attr.label(
222                cfg = device_transition,
223                aspects = [soong_prebuilt_tradefed_test_aspect],
224            ),
225            "host_test": attr.label(
226                cfg = host_transition,
227                aspects = [soong_prebuilt_tradefed_test_aspect],
228            ),
229            "tradefed_deps": attr.label_list(
230                cfg = host_transition,
231                aspects = [soong_prebuilt_tradefed_test_aspect],
232            ),
233            "_aapt": attr.label(
234                default = aapt_label,
235                allow_single_file = True,
236                cfg = host_transition,
237                aspects = [soong_prebuilt_tradefed_test_aspect],
238            ),
239            "_aapt2": attr.label(
240                default = aapt2_label,
241                allow_single_file = True,
242                cfg = host_transition,
243                aspects = [soong_prebuilt_tradefed_test_aspect],
244            ),
245        },
246    ),
247    test = True,
248    implementation = _tradefed_device_test_impl,
249    toolchains = _TOOLCHAINS,
250    doc = "A rule used to run device tests using Tradefed",
251)
252
253def tradefed_device_driven_test(
254        name,
255        test,
256        tradefed_deps = [],
257        suites = [],
258        **attrs):
259    tradefed_test_name = "tradefed_test_%s" % name
260    _tradefed_device_test(
261        name = tradefed_test_name,
262        device_test = test,
263        tradefed_deps = _get_tradefed_deps(suites, tradefed_deps),
264        **attrs
265    )
266    device_test(
267        name = name,
268        test = tradefed_test_name,
269    )
270
271def tradefed_host_driven_device_test(test, tradefed_deps = [], suites = [], **attrs):
272    _tradefed_device_test(
273        host_test = test,
274        tradefed_deps = _get_tradefed_deps(suites, tradefed_deps),
275        **attrs
276    )
277
278def _tradefed_test_impl(
279        ctx,
280        tradefed_options = [],
281        tradefed_deps = [],
282        test_host_deps = [],
283        test_device_deps = [],
284        path_additions = [],
285        add_extra_tradefed_test_files = lambda ctx, tradefed_test_dir: [],
286        data = []):
287    path_additions = path_additions + [_BAZEL_WORK_DIR + ctx.file._adb.dirname]
288
289    # Files required to run the host-side test.
290    test_host_runfiles = _collect_runfiles(ctx, test_host_deps)
291    test_host_runtime_jars = _collect_runtime_jars(test_host_deps)
292    test_host_runtime_shared_libs = _collect_runtime_shared_libs(test_host_deps)
293
294    # Files required to run the device-side test.
295    test_device_runfiles = _collect_runfiles(ctx, test_device_deps)
296
297    # Files required to run Tradefed.
298    all_tradefed_deps = []
299    all_tradefed_deps.extend(ctx.attr._tradefed_classpath_jars)
300    all_tradefed_deps.extend(ctx.attr._atest_tradefed_launcher)
301    all_tradefed_deps.extend(ctx.attr._atest_helper)
302    all_tradefed_deps.extend(ctx.attr._adb)
303    all_tradefed_deps.extend(tradefed_deps)
304
305    tradefed_runfiles = _collect_runfiles(ctx, all_tradefed_deps)
306    tradefed_runtime_jars = _collect_runtime_jars(all_tradefed_deps)
307    tradefed_runtime_shared_libs = _collect_runtime_shared_libs(all_tradefed_deps)
308
309    result_reporters_config_file = _generate_reporter_config(ctx)
310    tradefed_runfiles = tradefed_runfiles.merge(
311        ctx.runfiles(files = [result_reporters_config_file]),
312    )
313
314    py_paths, py_runfiles = _configure_python_toolchain(ctx)
315    java_paths, java_runfiles, java_home = _configure_java_toolchain(ctx)
316    path_additions = path_additions + java_paths + py_paths
317    tradefed_runfiles = tradefed_runfiles.merge_all([py_runfiles, java_runfiles])
318
319    tradefed_test_dir = "%s_tradefed_test_dir" % ctx.label.name
320    tradefed_test_files = []
321
322    for dep in tradefed_deps + test_host_deps + test_device_deps:
323        for f in dep[TradefedTestDependencyInfo].transitive_test_files.to_list():
324            symlink = _symlink(ctx, f, "%s/%s" % (tradefed_test_dir, f.short_path))
325            tradefed_test_files.append(symlink)
326
327    tradefed_test_files.extend(add_extra_tradefed_test_files(ctx, tradefed_test_dir))
328
329    script = ctx.actions.declare_file("tradefed_test_%s.sh" % ctx.label.name)
330    ctx.actions.expand_template(
331        template = ctx.file._tradefed_test_template,
332        output = script,
333        is_executable = True,
334        substitutions = {
335            "{module_name}": ctx.attr.module_name,
336            "{atest_tradefed_launcher}": _abspath(ctx.file._atest_tradefed_launcher),
337            "{atest_helper}": _abspath(ctx.file._atest_helper),
338            "{tradefed_test_dir}": _BAZEL_WORK_DIR + "%s/%s" % (
339                ctx.label.package,
340                tradefed_test_dir,
341            ),
342            "{tradefed_classpath}": _classpath([tradefed_runtime_jars, test_host_runtime_jars]),
343            "{shared_lib_dirs}": _ld_library_path([tradefed_runtime_shared_libs, test_host_runtime_shared_libs]),
344            "{path_additions}": ":".join(path_additions),
345            "{additional_tradefed_options}": " ".join(tradefed_options),
346            "{result_reporters_config_file}": _abspath(result_reporters_config_file),
347            "{java_home}": java_home,
348        },
349    )
350
351    return [
352        DefaultInfo(
353            executable = script,
354            runfiles = tradefed_runfiles.merge_all([
355                test_host_runfiles,
356                test_device_runfiles,
357                ctx.runfiles(tradefed_test_files),
358            ] + [ctx.runfiles(d.files.to_list()) for d in data]),
359        ),
360        TradefedTestInfo(
361            module_name = ctx.attr.module_name,
362        ),
363    ]
364
365def _get_tradefed_deps(suites, tradefed_deps = []):
366    suite_to_deps = {
367        "host-unit-tests": [],
368        "null-suite": [],
369        "device-tests": [],
370        "general-tests": [],
371        "vts": [vts_core_tradefed_harness_label],
372    }
373    all_tradefed_deps = {d: None for d in tradefed_deps}
374
375    for s in suites:
376        all_tradefed_deps.update({
377            d: None
378            for d in suite_to_deps.get(s, [compatibility_tradefed_label])
379        })
380
381    # Since `vts-core-tradefed-harness` includes `compatibility-tradefed`, we
382    # will exclude `compatibility-tradefed` if `vts-core-tradefed-harness` exists.
383    if vts_core_tradefed_harness_label in all_tradefed_deps:
384        all_tradefed_deps.pop(compatibility_tradefed_label, default = None)
385
386    return all_tradefed_deps.keys()
387
388def _generate_reporter_config(ctx):
389    result_reporters = [
390        "com.android.tradefed.result.BazelExitCodeResultReporter",
391        "com.android.tradefed.result.BazelXmlResultReporter",
392        "com.android.tradefed.result.proto.FileProtoResultReporter",
393    ]
394
395    result_reporters.extend(ctx.attr._extra_tradefed_result_reporters[BuildSettingInfo].value)
396
397    result_reporters_config_file = ctx.actions.declare_file("result-reporters-%s.xml" % ctx.label.name)
398    _write_reporters_config_file(
399        ctx,
400        result_reporters_config_file,
401        result_reporters,
402    )
403
404    return result_reporters_config_file
405
406def _write_reporters_config_file(ctx, config_file, result_reporters):
407    config_lines = [
408        "<?xml version=\"1.0\" encoding=\"utf-8\"?>",
409        "<configuration>",
410    ]
411
412    for result_reporter in result_reporters:
413        config_lines.append("    <result_reporter class=\"%s\" />" % result_reporter)
414
415    config_lines.append("</configuration>")
416
417    ctx.actions.write(config_file, "\n".join(config_lines))
418
419def _configure_java_toolchain(ctx):
420    java_runtime = ctx.toolchains[_JAVA_TOOLCHAIN].java_runtime
421    java_home_path = _BAZEL_WORK_DIR + java_runtime.java_home
422    java_runfiles = ctx.runfiles(transitive_files = java_runtime.files)
423    return ([java_home_path + "/bin"], java_runfiles, java_home_path)
424
425def _configure_python_toolchain(ctx):
426    py_toolchain_info = ctx.toolchains[_PY_TOOLCHAIN]
427    py2_interpreter = py_toolchain_info.py2_runtime.interpreter
428    py3_interpreter = py_toolchain_info.py3_runtime.interpreter
429
430    # Create `python` and `python3` symlinks in the runfiles tree and add them
431    # to the executable path. This is required because scripts reference these
432    # commands in their shebang line.
433    py_runfiles = ctx.runfiles(symlinks = {
434        "/".join([py2_interpreter.dirname, "python"]): py2_interpreter,
435        "/".join([py3_interpreter.dirname, "python3"]): py3_interpreter,
436    })
437    py_paths = [
438        _BAZEL_WORK_DIR + py2_interpreter.dirname,
439        _BAZEL_WORK_DIR + py3_interpreter.dirname,
440    ]
441    return (py_paths, py_runfiles)
442
443def _symlink(ctx, target_file, output_path):
444    symlink = ctx.actions.declare_file(output_path)
445    ctx.actions.symlink(output = symlink, target_file = target_file)
446    return symlink
447
448def _collect_runfiles(ctx, targets):
449    return ctx.runfiles().merge_all([
450        target[DefaultInfo].default_runfiles
451        for target in targets
452    ])
453
454def _collect_runtime_jars(deps):
455    return depset(
456        transitive = [
457            d[TradefedTestDependencyInfo].runtime_jars
458            for d in deps
459        ],
460    )
461
462def _collect_runtime_shared_libs(deps):
463    return depset(
464        transitive = [
465            d[TradefedTestDependencyInfo].runtime_shared_libraries
466            for d in deps
467        ],
468    )
469
470def _classpath(deps):
471    runtime_jars = depset(transitive = deps)
472    return ":".join([_abspath(f) for f in runtime_jars.to_list()])
473
474def _ld_library_path(deps):
475    runtime_shared_libs = depset(transitive = deps)
476    return ":".join(
477        [_BAZEL_WORK_DIR + f.dirname for f in runtime_shared_libs.to_list()],
478    )
479
480def _abspath(file):
481    return _BAZEL_WORK_DIR + file.short_path
482