tests-scan

#!/usr/bin/env python3

# This file is part of Cockpit.
#
# Copyright (C) 2015 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import argparse
import json
import logging
import shlex
import sys
import time
from typing import Iterable, NamedTuple

from lib import ALLOWLIST, testmap
from task import distributed_queue, github, labels_of_pull

sys.dont_write_bytecode = True
logging.basicConfig(level=logging.INFO)

no_amqp = False
try:
    import pika
except ImportError:
    no_amqp = True


def build_policy(repo, context):
    policy = testmap.tests_for_project(repo)
    if context:
        short_contexts = []
        for c in context:
            short_contexts.append(c.split("@")[0])
        new_policy = {}
        for (branch, contexts) in policy.items():
            branch_context = []
            for c in short_contexts:
                if c in contexts:
                    branch_context.append(c)
            if branch_context:
                new_policy[branch] = branch_context
        policy = new_policy
    return policy


def main():
    parser = argparse.ArgumentParser(description='Bot: scan and update status of pull requests on GitHub')
    parser.add_argument('-v', '--human-readable', action="store_true", default=False,
                        help='Display human readable output rather than tasks')
    parser.add_argument('-n', '--dry', action="store_true", default=False,
                        help="Don't actually change anything on GitHub")
    parser.add_argument('-f', '--force', action="store_true", default=False,
                        help='Perform all actions, even those that should be skipped')
    parser.add_argument('--repo', default=None,
                        help='Repository to scan and checkout.')
    parser.add_argument('-c', '--context', action="append", default=[],
                        help='Test contexts to use.')
    parser.add_argument('-p', '--pull-number', default=None,
                        help='Single pull request to scan for tasks')
    parser.add_argument('--pull-data', default=None,
                        help='pull_request event GitHub JSON data to evaluate; mutualy exclusive with -p and -s')
    parser.add_argument('-s', '--sha', default=None,
                        help='Trigger for a specific SHA; file issue on failure if SHA is not on a PR')
    parser.add_argument('--amqp', default=None,
                        help='The host:port of the AMQP server to publish to (format host:port)')

    opts = parser.parse_args()
    if opts.amqp and no_amqp:
        logging.error("AMQP host:port specified but python-amqp not available")
        return 1
    if opts.pull_data and (opts.pull_number or opts.sha):
        parser.error("--pull-data and --pull-number/--sha are mutually exclusive")

    if opts.force:
        if not opts.repo or not opts.context or not (opts.pull_number or opts.sha):
            parser.error('--force requires --repo, --context, and one of --pull-number or --sha')

    api = github.GitHub(repo=opts.repo)

    # HACK: The `repo` option is used throughout the code, for example repo from
    # opts is needed in `tests_invoke`, `tests_human`, `queue_test` etc.
    # Would be better to use api.repo everywhere
    opts.repo = api.repo

    try:
        results = scan_for_pull_tasks(api, opts.context, opts, opts.repo)
    except RuntimeError as ex:
        logging.error("tests-scan: %s", ex)
        return 1

    for result in results:
        if result:
            sys.stdout.write(result + "\n")

    return 0


class Job(NamedTuple):
    priority: str
    name: str
    number: int
    revision: str
    ref: str
    context: str
    base: 'str | None'
    repo: str
    bots_ref: 'str | None'
    github_context: str
    container: 'str | None'


# Prepare a human readable output
def tests_human(job: Job, options: argparse.Namespace) -> str:
    if not job.priority and not options.force:
        return ""
    return "{name:11} {context:25} {revision:10} {priority:2}{repo}{bots_ref}{branches}".format(
        priority=job.priority,
        revision=job.revision[0:7],
        context=job.context,
        name=job.name,
        repo=job.repo and "  (%s)" % job.repo or "",
        bots_ref=f" [bots@{job.bots_ref}]" if job.bots_ref else "",
        branches=f"   {{{job.base}}}" if job.base else "",
    )


def is_internal_context(context):
    # HACK: CentOS CI doesn't like firefox, browser fails to start way too often
    for pattern in ["rhel", "edge", "vmware", "openstack", "/firefox"]:
        if pattern in context:
            return True
    return False


# Prepare a test invocation command
def tests_invoke(job: Job, options: argparse.Namespace) -> str:
    try:
        priority = int(job.priority)
    except (ValueError, TypeError):
        priority = 0
    if priority <= 0 and not options.force:
        return ""
    current = time.strftime('%Y%m%d-%H%M%S')
    (image, _, scenario) = job.context.partition("/")

    checkout = "PRIORITY={priority:04d} ./make-checkout --verbose --repo={repo}"
    # Special case for when running tests without a PR
    if job.number == 0:
        invoke = "../tests-invoke --revision {revision} --repo {github_base}"
    else:
        invoke = "../tests-invoke --pull-number {pull_number} --revision {revision} --repo {github_base}"
    test_env = "TEST_OS={image}"
    wrapper = "./s3-streamer --repo {github_base} --test-name {name}-{current} " \
              "--github-context {github_context} --revision {revision}"

    if job.base:
        quoted = shlex.quote(job.base)
        test_env += f" BASE_BRANCH={quoted}"
        checkout += f" --rebase={quoted}"

    if job.bots_ref:
        # we are checking the external repo on a cockpit PR, so stay on the project's main branch
        checkout += " {ref}"
        test_env += " COCKPIT_BOTS_REF={bots_ref}"
    else:
        # we are testing the repo itself, checkout revision from the PR
        checkout += " {ref} {revision}"

    if scenario:
        test_env += " TEST_SCENARIO={scenario}"

    cmd = " ".join([checkout, "&& cd make-checkout-workdir &&", test_env, invoke])
    return (wrapper + ' -- /bin/sh -c "' + cmd + '"').format(
        priority=priority,
        name=shlex.quote(job.name),
        revision=shlex.quote(job.revision),
        ref=shlex.quote(job.ref),
        bots_ref=shlex.quote(job.bots_ref or ''),
        image=shlex.quote(image),
        scenario=(shlex.quote(scenario)),
        github_context=job.github_context,
        current=current,
        pull_number=job.number,
        repo=shlex.quote(job.repo),
        github_base=shlex.quote(options.repo),
    )


def queue_test(job: Job, channel, options: argparse.Namespace) -> None:
    command = tests_invoke(job, options)

    if command:
        priority = min(job.priority, distributed_queue.MAX_PRIORITY)

        slug = f"{job.name}-{job.revision[:8]}-{time.strftime('%Y%m%d-%H%M%S')}-{job.context.replace('/', '-')}"
        (image, _, scenario) = job.context.partition("/")

        env = {
            "TEST_OS": image,
        }
        if scenario:
            env["TEST_SCENARIO"] = scenario
        if job.bots_ref:
            env["COCKPIT_BOTS_REF"] = job.bots_ref
        if job.base:
            env["BASE_BRANCH"] = job.base

        body = {
            # old-style command line jobs
            "command": command,
            "type": "test",
            "sha": job.revision,
            "ref": job.ref,
            "name": job.name,
            # future job-runner format; once everything moved over, remove the stuff above
            "job": {
                "repo": job.repo,
                "sha": job.revision,
                "context": job.context,
                "target": job.base,
                "slug": slug,
                "env": env,
                "container": job.container,
                # for updating naughty trackers
                "secrets": ["github-token"],
            }
        }

        # report issue for nightly runs
        if job.number == 0:
            body["job"]["report"] = {
                "title": f"Tests failed on {job.revision}",
                "labels": ["nightly"],
            }
        else:
            body["job"]["pull"] = job.number

        queue = 'rhel' if is_internal_context(job.context) else 'public'
        channel.basic_publish('', queue, json.dumps(body), properties=pika.BasicProperties(priority=priority))
        logging.info("Published '%s' on '%s' with command: '%s'", job.name, job.revision, command)


def prioritize(status, title, labels, priority, context, number, direct):
    state = status.get("state", None)
    update = {"state": "pending"}

    # This commit definitively succeeded or failed
    if state in ["success", "failure"]:
        logging.info("Skipping '%s' on #%s because it has already finished", context, number)
        priority = 0
        update = None

    # This test errored, we try again but low priority
    elif state in ["error"]:
        priority -= 2

    elif state in ["pending"]:
        logging.info("Not updating status for '%s' on #%s because it is pending", context, number)
        update = None

    # Ignore context when the PR has [no-test] in the title or as label, unless
    # the context was directly triggered
    if (('no-test' in labels or '[no-test]' in title) and status.get("description", "") != github.NOT_TESTED_DIRECT):
        logging.info("Skipping '%s' on #%s because it is no-test", context, number)
        priority = 0
        update = None

    if priority > 0:
        if "priority" in labels:
            priority += 2
        if "blocked" in labels:
            priority -= 1

        # Pull requests where the title starts with WIP get penalized
        if title.startswith("WIP") or "needswork" in labels:
            priority -= 1

        # Is testing already in progress?
        description = status.get("description", "")
        if description.startswith(github.TESTING):
            logging.info("Skipping '%s' on #%s because it is already running", context, number)
            priority = description
            update = None

    if update:
        if priority <= 0:
            logging.info("Not updating status for '%s' on #%s because of low priority", context, number)
            update = None
        else:
            update["description"] = github.NOT_TESTED_DIRECT if direct else github.NOT_TESTED

    return [priority, update]


def dict_is_subset(full, check):
    for (key, value) in check.items():
        if key not in full or full[key] != value:
            return False
    return True


def update_status(api, revision, context, last, changes):
    if changes:
        changes["context"] = context
    if changes and not dict_is_subset(last, changes):
        response = api.post("statuses/" + revision, changes, accept=[422])  # 422 Unprocessable Entity
        errors = response.get("errors", None)
        if not errors:
            return True
        for error in response.get("errors", []):
            sys.stderr.write("{0}: {1}\n".format(revision, error.get('message', json.dumps(error))))
            sys.stderr.write(json.dumps(changes))
        return False
    return True


def cockpit_tasks(api: github.GitHub, contexts, opts: argparse.Namespace, repo: str) -> 'Iterable[Job]':
    pulls = []

    if opts.pull_data:
        pulls.append(json.loads(opts.pull_data)['pull_request'])
    elif opts.pull_number:
        pull = api.get(f"pulls/{opts.pull_number}")
        if pull:
            pulls.append(pull)
        else:
            sys.exit(f"Can't find pull request {opts.pull_number}")
    else:
        pulls = api.pulls()

    # Find matching PR for --sha
    if opts.sha:
        for pull in pulls:
            if pull["head"]["sha"].startswith(opts.sha):
                pulls = [pull]
                break
        else:
            logging.info("Processing revision %s without pull request", opts.sha)
            pulls = [{
                "title": f"{opts.sha}",
                "number": 0,
                "head": {
                    "sha": opts.sha,
                    "user": {
                        "login": "cockpit-project"
                    }
                },
                "labels": [],
            }]

    for pull in pulls:
        title = pull["title"]
        number = pull["number"]
        revision = pull["head"]["sha"]
        statuses = api.statuses(revision)
        login = pull["head"]["user"]["login"]
        base = pull.get("base", {}).get("ref")  # The branch this pull request targets (None for direct SHA triggers)
        merge_sha = pull.get("merge_commit_sha", revision)

        allowed = login in ALLOWLIST

        logging.info("Processing #%s titled '%s' on revision %s", number, title, revision)

        labels = labels_of_pull(pull)

        baseline: float = distributed_queue.BASELINE_PRIORITY
        # amqp automatically prioritizes on age
        if not opts.amqp:
            # modify the baseline slightly to favor older pull requests, so that we don't
            # end up with a bunch of half tested pull requests
            baseline += 1.0 - (min(100000, float(number)) / 100000)

        # Create list of statuses to process: always process the requested contexts, if given
        todos: dict[str, dict[str, object]] = {context: {} for context in contexts}
        for status in statuses:  # Firstly add all valid contexts that already exist in github
            if contexts and status not in contexts:
                continue
            if testmap.is_valid_context(status, repo):
                todos[status] = statuses[status]
        if not statuses and base:  # If none already present in PR, add basic set of contexts
            for context in build_policy(repo, contexts).get(base, []):
                todos[context] = {}

        container = api.fetch_file(merge_sha, '.cockpit-ci/container')

        # there are 3 different HEADs
        # ref:    the PR or SHA that we are testing
        # base:   the target branch of that PR (None for direct SHA trigger)
        # branch: the branch of the external project that we are testing
        #         against this PR (only applies to cockpit-project/bots PRs)
        for context in todos:
            # Get correct project and branch. Ones from test name have priority
            project = repo
            branch = base
            image_scenario, bots_pr, context_project, context_branch = testmap.split_context(context)
            if context_project:
                project = context_project
                branch = context_branch or testmap.get_default_branch(project)

            # Note: Don't use `pull/<pr_number>/head` as it may point to an old revision
            ref = revision

            # For unmarked and untested status, user must be allowed
            # Not this only applies to this specific commit. A new status
            # will apply if the user pushes a new commit.
            status = todos[context]
            if not allowed and status.get("description", github.NO_TESTING) == github.NO_TESTING:
                priority = github.NO_TESTING
                changes = {"description": github.NO_TESTING, "context": context, "state": "pending"}
            else:
                # with --amqp (as called from run-queue), trigger tests as NOT_TESTED, as they already get queued;
                # without --amqp (as called manually or from workflows), trigger tests as NOT_TESTED_DIRECT,
                # so that the webhook queues them
                (priority, changes) = prioritize(
                    status, title, labels, baseline, context, number, direct=not opts.amqp
                )
            if opts.dry or update_status(api, revision, context, status, changes):
                checkout_ref = ref
                if project != repo:
                    checkout_ref = testmap.get_default_branch(project)
                if base != branch:
                    checkout_ref = branch

                if repo == "cockpit-project/bots":
                    # bots own test doesn't need bots/ setup as there is a permanent symlink to itself there
                    # otherwise if we're testing an external project (repo != project) then checkout bots from the PR
                    bots_ref = None if repo == project else ref
                else:
                    if bots_pr:
                        # Note: Don't use `pull/<pr_number>/head` as it may point to an old revision
                        bots_api = github.GitHub(repo="cockpit-project/bots")
                        bots_ref = bots_api.getHead(bots_pr) or "xxx"  # Make sure we fail when cannot get the head
                    else:
                        bots_ref = "main"

                yield Job(
                    priority,
                    "pull-%d" % number,
                    number,
                    revision,
                    checkout_ref,
                    image_scenario,
                    branch,
                    project,
                    bots_ref,
                    context,
                    container,
                )


def scan_for_pull_tasks(api, contexts, opts, repo):
    results = list(cockpit_tasks(api, contexts, opts, repo))

    if opts.human_readable:
        results.sort(reverse=True, key=str)
        return [tests_human(job, options=opts) for job in results]
    if not opts.amqp:
        return [tests_invoke(job, options=opts) for job in results]
    with distributed_queue.DistributedQueue(opts.amqp, ['rhel', 'public']) as q:
        return [queue_test(job, channel=q.channel, options=opts) for job in results]


if __name__ == '__main__':
    sys.exit(main())