recipe_modules/presubmit_util/test_api.py

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

from recipe_engine import recipe_test_api

from PB.go.chromium.org.luci.buildbucket.proto import (
    builds_service as builds_service_pb2,
)

from PB.recipe_modules.fuchsia.presubmit_util.options import Options


class PresubmitUtilTestApi(recipe_test_api.RecipeTestApi):
    Options = Options

    def create_cl(self, step_name, test_data):
        return self.step_data(step_name, self.m.json.output(test_data))

    def wait_for_cq(self, success=True):
        return self.step_data("wait for CQ", self.m.json.output(success))

    def trigger_tryjobs(self, step_name, tryjob_msgs):
        responses = [dict(schedule_build=m) for m in tryjob_msgs]
        return self.m.buildbucket.simulated_schedule_output(
            batch_response=builds_service_pb2.BatchResponse(responses=responses),
            step_name=f"{step_name}.schedule",
        )

    def collect_tryjobs(
        self, step_name, builders=None, tryjob_msgs=None, status="SUCCESS"
    ):
        if builders:
            tryjob_msgs = []
            for idx, builder in enumerate(builders):
                project, bucket, builder_name = builder.split("/")
                msg = self.m.buildbucket.try_build_message(
                    builder=builder_name,
                    bucket=bucket,
                    project=project,
                    build_id=(idx + 1) * 1234567890,
                    status=status,
                )
                tryjob_msgs.append(msg)
            search_results = self.m.buildbucket.simulated_search_results(
                tryjob_msgs,
                step_name=f"{step_name}.search",
            )
        collect_data = self.m.buildbucket.simulated_collect_output(
            tryjob_msgs,
            step_name=f"{step_name}.collect",
        )
        if builders:
            return search_results + collect_data
        return collect_data