# 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_ensure",
    "fuchsia/cipd_util",
    "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):
    previous_version_dir = api.path.mkdtemp(project.name)
    api.cipd.ensure(
        previous_version_dir,
        api.cipd.EnsureFile().add_package(project.cipd_package, CIPD_REF),
    )

    output_dir = api.path.mkdtemp(project.name + "_output")

    update_tool = api.cipd_ensure(
        api.resource("cipd.ensure"),
        "fuchsia/infra/update_test_durations/${platform}",
    )
    api.step(
        "query durations",
        [
            update_tool,
            "run",
            "-previous-version-dir",
            previous_version_dir,
            "-output-dir",
            output_dir,
            "-project",
            project.name,
        ],
    )
    if dry_run:
        api.cas_util.upload(output_dir)
    else:
        api.cipd_util.upload_package(
            project.cipd_package,
            output_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.test("basic") + properties()
    yield api.test("dry_run") + properties(dry_run=True)