__init__.py

"""
Mocks for the gRPC servers, a.k.a. `pytest_userver.grpc`.
 
@sa @ref scripts/docs/en/userver/tutorial/grpc_service.md
@sa @ref pytest_userver.plugins.grpc.mockserver
"""
 
from __future__ import annotations
 
import asyncio
from collections.abc import Callable
from collections.abc import Iterator
import contextlib
import inspect
import types
import typing
from typing import Any
from typing import Generic
from typing import TypeAlias
from typing import TypeVar
 
import grpc
 
import testsuite.utils.callinfo
 
from ._client import PreCallClientInterceptor  # noqa: F401
from ._mocked_errors import MockedError  # noqa: F401
from ._mocked_errors import NetworkError  # noqa: F401
from ._mocked_errors import TimeoutError  # noqa: F401
from ._servicer_mock import _check_is_servicer_class
from ._servicer_mock import _create_servicer_mock
from ._servicer_mock import _MockHandlerRemovalToken
from ._servicer_mock import _ServiceMock
 
Handler: TypeAlias = Callable[[Any, grpc.aio.ServicerContext], Any]
MockDecorator: TypeAlias = Callable[[Handler], testsuite.utils.callinfo.AsyncCallQueue]
AsyncExcAppend: TypeAlias = Callable[[Exception], None]
Servicer = TypeVar('Servicer')
 
 
class MockserverSession:
    """
    Allows to install mocks that are kept active between tests, see
    @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver_session "grpc_mockserver_session".
 
    @warning This is a sharp knife, use with caution! For most use-cases, prefer
    @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver" instead.
    """
 
    def __init__(self, *, server: grpc.aio.Server, experimental: bool = False) -> None:
        """
        @warning This initializer is an **experimental API**, likely to break in the future. Consider using
        @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver_session "grpc_mockserver_session" instead.
 
        Initializes `MockserverSession`. Takes ownership of `server`.
        To properly start and stop the server, use `MockserverSession` as an async contextmanager:
 
        @code{.py}
        async with pytest_userver.grpc.mockserver.MockserverSession(server=server) as mockserver:
            # ...
        @endcode
 
        @note `MockserverSession` is usually obtained from
        @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver_session "grpc_mockserver_session".
        """
        assert experimental
        self._server = server
        self._auto_service_mocks: dict[type, _ServiceMock] = {}
        self._asyncexc_append: AsyncExcAppend | None = None
 
    def _get_auto_service_mock(self, servicer_class: type, /) -> _ServiceMock:
        existing_mock = self._auto_service_mocks.get(servicer_class)
        if existing_mock:
            return existing_mock
        new_mock = _create_servicer_mock(servicer_class)
        new_mock.set_asyncexc_append(self._asyncexc_append)
        self._auto_service_mocks[servicer_class] = new_mock
        new_mock.add_to_server(self._server)
        return new_mock
 
    def _set_asyncexc_append(self, asyncexc_append: AsyncExcAppend | None, /) -> None:
        self._asyncexc_append = asyncexc_append
        for mock in self._auto_service_mocks.values():
            mock.set_asyncexc_append(asyncexc_append)
 
    def reset_mocks(self) -> None:
        """
        @brief Removes all mocks for this mockserver that have been installed using
        `MockserverSession` or @ref pytest_userver.grpc.Mockserver "Mockserver" API.
        @note Mocks installed manually using @ref MockserverSession.server will not be removed, though.
        """
        for mock in self._auto_service_mocks.values():
            mock.reset_handlers()
 
    @contextlib.contextmanager
    def asyncexc_append_scope(self, asyncexc_append: AsyncExcAppend | None, /) -> Iterator[None]:
        """
        Sets testsuite's `asyncexc_append` for use in the returned scope.
        """
        self._set_asyncexc_append(asyncexc_append)
        try:
            yield
        finally:
            self._set_asyncexc_append(None)
 
    @property
    def server(self) -> grpc.aio.Server:
        """
        The underlying [grpc.aio.Server](https://grpc.github.io/grpc/python/grpc_asyncio.html#grpc.aio.Server).
        """
        return self._server
 
    async def start(self) -> None:
        """
        Starts the server.
        @note Prefer starting mockserver using the async contextmanager syntax if possible.
        """
        await self._server.start()
 
    async def stop(self) -> None:
        """
        Stops the server properly. Prefer this method to stopping `server` manually.
        """
        await _stop_server(self._server)
 
    async def __aenter__(self) -> MockserverSession:
        await self.start()
        return self
 
    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.stop()
 
 
class Mockserver:
    """
    Allows to install mocks that are reset between tests, see
    @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver".
 
    Acts as a scope for the installed mocks. Mocks are removed on `__exit__` or @ref reset_mocks call.
    """
 
    def __init__(self, *, mockserver_session: MockserverSession, experimental: bool = False) -> None:
        """
        @warning This initializer is an **experimental API**, likely to break in the future. Consider using
        @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver" instead.
 
        Initializes Mockserver. Should be used together with context manager syntax (`with` block).
 
        @note `Mockserver` is usually obtained from
        @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver".
 
        An example on creating a custom scope for gRPC mocks:
        @snippet grpc/functional_tests/middleware_client/tests/test_custom_mockserver_session_scope.py  global mock
        """
        assert experimental
        self._mockserver_session = mockserver_session
        self._token = _MockHandlerRemovalToken()
        self._installed_service_mocks: set[_ServiceMock] = set()
 
    def __call__(self, servicer_method, /) -> MockDecorator:
        """
        Returns a decorator to mock the specified gRPC service method implementation.
 
        Example:
 
        @snippet samples/grpc_service/testsuite/test_grpc.py  Prepare modules
        @snippet samples/grpc_service/testsuite/test_grpc.py  grpc client test
        """
        servicer_class = _get_class_from_method(servicer_method)
        mock = self._get_service_mock(servicer_class)
        return mock.install_handler(servicer_method.__name__, self._token)
 
    def mock_factory(self, servicer_class, /) -> Callable[[str], MockDecorator]:
        """
        Allows to create a fixture as a shorthand for mocking methods of the specified gRPC service.
 
        Example:
 
        @snippet grpc/functional_tests/metrics/tests/conftest.py  Prepare modules
        @snippet grpc/functional_tests/metrics/tests/conftest.py  Prepare server mock
        @snippet grpc/functional_tests/metrics/tests/test_metrics.py  grpc client test
        """
 
        def factory(method_name):
            method = getattr(servicer_class, method_name, None)
            if method is None:
                raise ValueError(f'No method "{method_name}" in servicer class "{servicer_class.__name__}"')
            return self(method)
 
        _check_is_servicer_class(servicer_class)
        return factory
 
    def install_servicer(self, servicer: Servicer, /) -> Servicer:
        """
        Installs as a mock `servicer`, the class of which should inherit from a generated `*Servicer` class.
 
        1. Write a service implementation:
 
           @snippet grpc/functional_tests/middleware_client/tests/test_install_servicer.py servicer
 
           @note Inheritance from multiple `*Servicer` classes at once is allowed.
 
        2. Install servicer instance:
 
           @snippet grpc/functional_tests/middleware_client/tests/test_install_servicer.py install servicer
 
        3. Use service mock:
 
           @snippet grpc/functional_tests/middleware_client/tests/test_install_servicer.py use mock
        """
        base_servicer_classes = [cls for cls in inspect.getmro(type(servicer)) if _is_servicer_class(cls)]
        if not base_servicer_classes:
            raise ValueError(f"Given object's type ({type(servicer)}) is not inherited from any grpc *Servicer class")
        proxy = _MockProxy(servicer)
        for servicer_class in base_servicer_classes:
            mock = self._get_service_mock(servicer_class)
            for python_method_name in mock.known_methods:
                if _get_class_that_defined_method(type(servicer), python_method_name) not in base_servicer_classes:
                    handler_func: types.MethodType = getattr(servicer, python_method_name)
                    callqueue = mock.install_handler(python_method_name, token=self._token)(handler_func)
                    object.__setattr__(proxy, python_method_name, callqueue)
        return typing.cast(Servicer, proxy)
 
    def reset_mocks(self) -> None:
        """
        Removes all mocks installed using this `Mockserver` instance.
        """
        for mock in self._installed_service_mocks:
            mock.pop_handlers_for_token(self._token)
        self._installed_service_mocks.clear()
 
    def __enter__(self) -> MockserverSession:
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.reset_mocks()
 
    def _get_service_mock(self, servicer_class: type, /) -> _ServiceMock:
        mock = self._mockserver_session._get_auto_service_mock(servicer_class)  # pylint: disable=protected-access
        self._installed_service_mocks.add(mock)
        return mock
 
 
# @cond
 
 
async def _stop_server(server: grpc.aio.Server, /) -> None:
    async def stop_server():
        await server.stop(grace=None)
        await server.wait_for_termination()
 
    stop_server_task = asyncio.shield(asyncio.create_task(stop_server()))
    # asyncio.shield does not protect our await from cancellations, and we
    # really need to wait for the server stopping before continuing.
    try:
        await stop_server_task
    except asyncio.CancelledError:
        await stop_server_task
        # Propagate cancellation when we are done
        raise
 
 
def _get_class_from_method(method) -> type:
    # https://stackoverflow.com/a/54597033/5173839
    assert inspect.isfunction(method), f'Expected an unbound method: foo(ClassName.MethodName), got: {method}'
    class_name = method.__qualname__.split('.<locals>', 1)[0].rsplit('.', 1)[0]
    try:
        cls = getattr(inspect.getmodule(method), class_name)
    except AttributeError:
        cls = method.__globals__.get(class_name)
    assert isinstance(cls, type)
    return cls
 
 
def _get_class_that_defined_method(cls: type, method_name: str) -> type | None:
    for cls in inspect.getmro(cls):
        if method_name in cls.__dict__:
            return cls
    # May happen with a method defined via __getattr__.
    return None
 
 
def _is_servicer_class(cls: type) -> bool:
    try:
        return '_pb2_grpc.' in inspect.getfile(cls)
    except TypeError:
        return False
 
 
class _MockProxy(Generic[Servicer]):
    def __init__(self, wrapped_servicer: Servicer) -> None:
        object.__setattr__(self, '_wrapped_servicer', wrapped_servicer)
 
    def __getattr__(self, name: str) -> Any:
        servicer: Servicer = object.__getattribute__(self, '_wrapped_servicer')
        return getattr(servicer, name)
 
    def __setattr__(self, name: str, value: Any) -> None:
        servicer: Servicer = object.__getattribute__(self, '_wrapped_servicer')
        setattr(servicer, name, value)
 
 
# @endcond