api/faux.py

# Copyright 2019 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Decorators to help handle mock calls and errors in the Build API."""

import functools
from typing import Callable, ParamSpec, TYPE_CHECKING, TypeVar

from chromite.api import controller


if TYPE_CHECKING:
    from chromite.third_party.google.protobuf import message as protobuf_message

    from chromite.api import api_config


P = ParamSpec("P")  # Represents the parameters of the decorated function
R = TypeVar("R")  # Represents the return type of the decorated function

BuildAPIFunction = Callable[P, R]


def all_responses(
    faux_result_factory: BuildAPIFunction,
) -> Callable[[BuildAPIFunction], BuildAPIFunction]:
    """A decorator to handle all mock responses.

    This is syntactic sugar for handling all the mock response types in a
    single function.

    Args:
        faux_result_factory: A function with the same signature as a regular
            API endpoint function that populates the output with success or
            error results, as requested, without executing the endpoint.
    """
    # Get the decorators for each of the mock types, so we can compose them.
    success_fn = success(faux_result_factory)
    err_fn = error(faux_result_factory)

    def _decorator(func: BuildAPIFunction) -> BuildAPIFunction:
        f = success_fn(func)
        return err_fn(f)

    return _decorator


def all_empty(func: BuildAPIFunction) -> BuildAPIFunction:
    """Decorator to handle all mock responses with an empty output."""
    return empty_error(empty_success(func))


def success(
    faux_result_factory: BuildAPIFunction,
) -> Callable[[BuildAPIFunction], BuildAPIFunction]:
    """A decorator to handle mock call responses.

    Args:
        faux_result_factory: A function with the same signature as a regular
            API endpoint function that populates the output
    """

    def decorator(func: BuildAPIFunction) -> BuildAPIFunction:
        @functools.wraps(func)
        def _success(
            *args: P.args,
            **kwargs: P.kwargs,
        ) -> R:
            if args[2].mock_call:
                faux_result_factory(*args, **kwargs)
                return controller.RETURN_CODE_SUCCESS

            return func(*args, **kwargs)

        return _success

    return decorator


def empty_success(func: BuildAPIFunction) -> BuildAPIFunction:
    """A decorator to handle mock success responses with empty outputs."""

    @functools.wraps(func)
    def _empty_success(
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> R:
        if args[2].mock_call:
            return controller.RETURN_CODE_SUCCESS

        return func(*args, **kwargs)

    return _empty_success


def error(
    faux_error_factory: BuildAPIFunction,
) -> Callable[[BuildAPIFunction], BuildAPIFunction]:
    """A decorator to handle mock error responses."""

    def decorator(func: BuildAPIFunction) -> BuildAPIFunction:
        @functools.wraps(func)
        def _error(
            *args: P.args,
            **kwargs: P.kwargs,
        ) -> R:
            if args[2].mock_error:
                faux_error_factory(*args, **kwargs)
                return controller.RETURN_CODE_UNSUCCESSFUL_RESPONSE_AVAILABLE

            return func(*args, **kwargs)

        return _error

    return decorator


def empty_error(func: BuildAPIFunction) -> BuildAPIFunction:
    """A decorator to handle mock error responses with empty outputs."""

    @functools.wraps(func)
    def _empty_error(
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> R:
        if args[2].mock_error:
            return controller.RETURN_CODE_UNRECOVERABLE

        return func(*args, **kwargs)

    return _empty_error


def empty_completed_unsuccessfully_error(
    func: BuildAPIFunction,
) -> BuildAPIFunction:
    """A decorator to handle mock unsuccessful response with empty outputs."""

    @functools.wraps(func)
    def _empty_error(
        *args: P.args,
        **kwargs: P.kwargs,
    ) -> R:
        if args[2].mock_error:
            return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY

        return func(*args, **kwargs)

    return _empty_error