scripts/led-edit-recipe-wrapper.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_wrapper using led.

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

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

import json
import os
import subprocess
import sys
import tempfile


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

    # led populates the buildbucket.name field with a string corresponding to
    # the led command that produced the job definition, generally of the form
    # "led: get-build <bid>" or "led: get-builder <builder>".
    if not job["buildbucket"]["name"].startswith("led: get-build "):
        # TODO(olivernewman): Figure out why
        # `led get-builder <builder> | led edit-isolated` causes jobs to fail to
        # find the luciexe executable.
        print(
            "led-edit-recipe-wrapper.py: input must come from `led get-build <bid>`",
            file=sys.stderr,
        )
        sys.exit(1)

    # 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_wrapper 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", "-mod=vendor", "-o", exe_path, "."],
            env={**os.environ, **go_env},
            # `go build` must be run from within the module containing
            # recipe_wrapper.
            cwd=os.path.join(infra_dir, "cmd", "recipe_wrapper"),
        )

        # `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()