add RBE device test android build downloader am: ded6ea7ad5 am: 33a3632287 am: c399aa52d4 am: 6faa48e575

Original change: https://android-review.googlesource.com/c/platform/build/bazel_common_rules/+/2720157

Change-Id: I2ce9a030d4c356602a45bc14c18479744598d5d3
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

3 files changed, 136 insertions, 0 deletions

diff --git a/rules/remote_device/downloader/BUILD.bazel b/rules/remote_device/downloader/BUILD.bazel
new file mode 100644
index 0000000..600db87
--- /dev/null
+++ b/rules/remote_device/downloader/BUILD.bazel
@@ -0,0 +1,3 @@
+package(default_visibility = ["//visibility:public"])
+
+exports_files(["download_cvd_build.sh.template"])
\ No newline at end of file
diff --git a/rules/remote_device/downloader/download_cvd_artifact.bzl b/rules/remote_device/downloader/download_cvd_artifact.bzl
new file mode 100644
index 0000000..bdc1aef
--- /dev/null
+++ b/rules/remote_device/downloader/download_cvd_artifact.bzl
@@ -0,0 +1,111 @@
+# Copyright (C) 2023 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.
+
+load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
+
+ImageProvider = provider(
+    "Provide device image and host package files",
+    fields = {
+        "image": "device image file to launch virtual device",
+        "cvd_host_package": "cvd host package to launch virtual device",
+    },
+)
+
+# TODO(b/273846592): Move to a common file, same for other files.
+_BAZEL_WORK_DIR = "${TEST_SRCDIR}/${TEST_WORKSPACE}/"
+
+# Target name in main branches are fixed, and in aosp branches it is prefixed
+# with "aosp_".
+_DEFAULT_BUILD_TARGET = "cf_x86_64_phone-userdebug"
+_IMAGE_ARTIFACT_PATH = "cf_x86_64_phone-img-{BUILD_ID}.zip"
+_HOST_PACKAGE_ARTIFACT_PATH = "cvd-host_package.tar.gz"
+
+def _download_cvd_artifact_impl(ctx):
+    build_id = ctx.attr.build_id[BuildSettingInfo].value
+    branch = ctx.attr.branch[BuildSettingInfo].value
+    image_artifact_path = _IMAGE_ARTIFACT_PATH.replace("{BUILD_ID}", build_id)
+    target = _DEFAULT_BUILD_TARGET
+    if not build_id:
+        fail("build_id must be specified to download build image.")
+    if not branch:
+        fail("branch must be specified to download build image.")
+
+    # Add "aosp_" prefix to target and artifact name if the branch is AOSP.
+    if "aosp" in branch:
+        image_artifact_path = "aosp_" + image_artifact_path
+        target = "aosp_" + target
+
+    image_out_file = _download_helper(
+        ctx,
+        image_artifact_path,
+        build_id,
+        branch,
+        target,
+    )
+    cvd_host_out_file = _download_helper(
+        ctx,
+        _HOST_PACKAGE_ARTIFACT_PATH,
+        build_id,
+        branch,
+        target,
+    )
+    return ImageProvider(
+        image = image_out_file,
+        cvd_host_package = cvd_host_out_file,
+    )
+
+def _download_helper(ctx, artifact_path, build_id, branch, target):
+    script = ctx.actions.declare_file("download_cvd_build_%s_%s.sh" %
+                                      (ctx.label.name, artifact_path))
+
+    out_file = ctx.actions.declare_file(artifact_path)
+    ctx.actions.expand_template(
+        template = ctx.file._create_script_template,
+        output = script,
+        is_executable = True,
+        substitutions = {
+            "{build_id}": build_id,
+            "{artifact_path}": artifact_path,
+            "{output_dir}": out_file.dirname,
+            "{branch}": branch,
+            "{target}": target,
+        },
+    )
+    ctx.actions.run_shell(
+        inputs = [script],
+        outputs = [out_file],
+        mnemonic = "DownloadCvd",
+        command = "source %s" % (script.path),
+        progress_message = "Downloading Android Build artifact %s for Build ID %s." % (artifact_path, build_id),
+    )
+    return out_file
+
+download_cvd_artifact = rule(
+    attrs = {
+        "_create_script_template": attr.label(
+            default = "//remote_device:download_cvd_build.sh.template",
+            allow_single_file = True,
+        ),
+        "build_id": attr.label(
+            mandatory = True,
+            doc = "sets the build id of the Android image",
+        ),
+        "branch": attr.label(
+            mandatory = True,
+            doc = "sets the branch of the Android image",
+        ),
+    },
+    implementation = _download_cvd_artifact_impl,
+    doc = "A rule used to download cuttlefish image files.",
+)
diff --git a/rules/remote_device/downloader/download_cvd_build.sh.template b/rules/remote_device/downloader/download_cvd_build.sh.template
new file mode 100644
index 0000000..6a67e6b
--- /dev/null
+++ b/rules/remote_device/downloader/download_cvd_build.sh.template
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+BUILD_ID="{build_id}"
+ARTIFACT_PATH="{artifact_path}"
+OUTPUT_DIR="{output_dir}"
+BRANCH="{branch}"
+TARGET="{target}"
+
+recipe_file=/tradefed/android_build_downloader/recipe
+cat > "$recipe_file" << EOF
+build_artifact_group {
+    build_id: "$BUILD_ID",
+    branch: "$BRANCH"
+    target: "$TARGET",
+    artifact_path: "$ARTIFACT_PATH"
+}
+
+EOF
+
+/tradefed/android_build_downloader/downloader.sh \
+  --output_dir $OUTPUT_DIR \
+  --recipe_file /tradefed/android_build_downloader/recipe