scripts/led-edit-recipe-bootstrap.py

#!/usr/bin/env python3
# Copyright 2021 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A script for testing recipe_bootstrap using led.

To launch a led job using your local modifications to recipe_bootstrap, run this
from the root of the infra/infra repo:

    led get-build ... | ./scripts/led-edit-recipe-bootstrap.py | led launch
"""

import json
import os
import subprocess
import sys
import tempfile


def main():
    job = json.load(sys.stdin)

    # Clear the CIPD package info out of the `exe` field to force the build to
    # use the `luciexe` executable included in the CAS input instead.
    build = job["buildbucket"]["bbagent_args"]["build"]
    exe = build["exe"]
    exe.pop("cipd_package", None)
    exe.pop("cipd_version", None)

    # Make sure we build recipe_bootstrap for the correct OS and arch used by
    # the build. This is best-effort, as it encodes infra implementation details
    # (names and allowed values of swarming dimensions) that may change in the
    # future.
    dimensions = {
        dim["key"]: dim["value"]
        for dim in build["infra"]["swarming"]["task_dimensions"]
    }
    os_dimension = dimensions.get("os", "Linux").lower()
    cpu_dimension = dimensions.get("cpu", "x64").lower()
    go_env = {
        "GOOS": {"linux": "linux", "mac": "darwin", "windows": "windows"}[os_dimension],
        "GOARCH": {"x64": "amd64", "arm64": "arm64"}[cpu_dimension],
        # Don't dynamically link against GLIBC.
        "CGO_ENABLED": "0",
    }

    infra_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))

    with tempfile.TemporaryDirectory() as build_dir:
        exe_name = "luciexe"
        if os_dimension == "windows":
            exe_name += ".exe"
        exe_path = os.path.join(build_dir, exe_name)
        subprocess.check_call(
            ["go", "build", "-o", exe_path, "."],
            env={**os.environ, **go_env},
            # `go build` must be run from within the module containing
            # recipe_bootstrap.
            cwd=os.path.join(infra_dir, "cmd", "recipe_bootstrap"),
        )

        # `led edit-isolated` runs a command in a temporary directory that will
        # be used as the job's CAS input.
        subprocess.run(
            ["led", "edit-isolated", "mv", exe_path, "."],
            input=json.dumps(job),
            text=True,
            check=True,
        )


if __name__ == "__main__":
    main()