recipes/update_test_durations.py

# Copyright 2022 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.

"""Recipe for updating test duration CIPD packages.

These packages will be rolled into a Fuchsia checkout (by a separate roller
builder) and the historical test duration data will used by testsharder to
produce shards of approximately equal duration.
"""

from PB.recipes.fuchsia.update_test_durations import InputProperties

DEPS = [
    "fuchsia/cas_util",
    "fuchsia/cipd_util",
    "fuchsia/ensure_tool",
    "fuchsia/status_check",
    "recipe_engine/cipd",
    "recipe_engine/path",
    "recipe_engine/properties",
    "recipe_engine/step",
    "recipe_engine/time",
]

PROPERTIES = InputProperties

CIPD_REF = "latest"


def RunSteps(api, props):
    # Compute a version for the CIPD packages based on the current timestamp.
    # The same version must be used for all the CIPD packages so that they can
    # be rolled in sync.
    now = api.time.utcnow()
    # E.g. 20220130115959 is 30 January 2022 at 11:59:59 AM.
    cipd_version = now.strftime("%Y%m%d%H%M%S")

    for project in props.projects:
        with api.step.nest(project.name):
            update_project_durations(api, project, cipd_version, props.dry_run)


def update_project_durations(api, project, cipd_version, dry_run):
    package_dir = api.path.mkdtemp(project.name)
    # Install the current version of the CIPD package so that the
    # updated durations can be overlaid on top.
    api.cipd.ensure(
        package_dir,
        api.cipd.EnsureFile().add_package(project.cipd_package, CIPD_REF),
    )

    update_tool = api.ensure_tool(
        "update_test_durations",
        api.resource("tool_manifest.json"),
    )
    api.step(
        "query durations",
        [update_tool, "run", "-dir", package_dir, "-project", project.name],
    )
    if dry_run:
        api.cas_util.upload(package_dir)
    else:
        api.cipd_util.upload_package(
            project.cipd_package,
            package_dir,
            search_tag={"version": cipd_version},
            refs=[CIPD_REF],
        )


def GenTests(api):
    def properties(**kwargs):
        props = {
            "projects": [
                {
                    "name": "fuchsia",
                    "cipd_package": "fuchsia/test_durations",
                },
                {
                    "name": "fuchsia-internal",
                    "cipd_package": "fuchsia-internal/test_durations",
                },
            ],
        }
        props.update(kwargs)
        return api.properties(**props)

    yield api.status_check.test("basic") + properties()
    yield api.status_check.test("dry_run") + properties(dry_run=True)