userver: /home/antonyzhilin/arcadia/taxi/uservices/userver/testsuite/pytest_plugins/pytest_userver/grpc/__init__.py Source File
All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Modules Pages Concepts
__init__.py
1"""
2Mocks for the gRPC servers, a.k.a. `pytest_userver.grpc.
3
4@sa @ref scripts/docs/en/userver/tutorial/grpc_service.md
5@sa @ref pytest_userver.plugins.grpc.mockserver
6"""
7
8from __future__ import annotations
9
10import asyncio
11import contextlib
12import inspect
13from typing import Any
14from typing import Callable
15from typing import Dict
16from typing import Optional
17
18import grpc
19
20import testsuite.utils.callinfo
21
22from ._mocked_errors import MockedError # noqa: F401
23from ._mocked_errors import NetworkError # noqa: F401
24from ._mocked_errors import TimeoutError # noqa: F401
25from ._servicer_mock import _check_is_servicer_class
26from ._servicer_mock import _create_servicer_mock
27from ._servicer_mock import _ServiceMock
28
29Handler = Callable[[Any, grpc.aio.ServicerContext], Any]
30MockDecorator = Callable[[Handler], testsuite.utils.callinfo.AsyncCallQueue]
31AsyncExcAppend = Callable[[Exception], None]
32
33
35 """
36 Allows to install mocks that are kept active between tests, see
37 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver_session "grpc_mockserver_session".
38
39 @warning This is a sharp knife, use with caution! For most use-cases, prefer
40 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver" instead.
41 """
42
43 def __init__(self, *, server: grpc.aio.Server, experimental: bool = False) -> None:
44 """
45 @warning This initializer is an **experimental API**, likely to break in the future. Consider using
46 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver_session "grpc_mockserver_session" instead.
47
48 Initializes `MockserverSession`. Takes ownership of `server`.
49 To properly start and stop the server, use `MockserverSession` as an async contextmanager:
50
51 @code{.py}
52 async with pytest_userver.grpc.mockserver.MockserverSession(server=server) as mockserver:
53 # ...
54 @endcode
55
56 @note `MockserverSession` is usually obtained from
57 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver_session "grpc_mockserver_session".
58 """
59 assert experimental
60 self._server = server
61 self._auto_service_mocks: Dict[type, _ServiceMock] = {}
62 self._asyncexc_append: Optional[AsyncExcAppend] = None
63
64 def _get_auto_service_mock(self, servicer_class: type, /) -> _ServiceMock:
65 existing_mock = self._auto_service_mocks.get(servicer_class)
66 if existing_mock:
67 return existing_mock
68 new_mock = _create_servicer_mock(servicer_class)
69 new_mock.set_asyncexc_append(self._asyncexc_append)
70 self._auto_service_mocks[servicer_class] = new_mock
71 new_mock.add_to_server(self._server)
72 return new_mock
73
74 def _set_asyncexc_append(self, asyncexc_append: Optional[AsyncExcAppend], /) -> None:
75 self._asyncexc_append = asyncexc_append
76 for mock in self._auto_service_mocks.values():
77 mock.set_asyncexc_append(asyncexc_append)
78
79 def reset_mocks(self) -> None:
80 """
81 @brief Removes all mocks for this mockserver that have been installed using
82 `MockserverSession` or @ref pytest_userver.grpc.Mockserver "Mockserver" API.
83 @note Mocks installed manually using @ref MockserverSession.server will not be removed, though.
84 """
85 for mock in self._auto_service_mocks.values():
86 mock.reset_handlers()
87
88 @contextlib.contextmanager
89 def asyncexc_append_scope(self, asyncexc_append: Optional[AsyncExcAppend], /):
90 """
91 Sets testsuite's `asyncexc_append` for use in the returned scope.
92 """
93 self._set_asyncexc_append(asyncexc_append)
94 try:
95 yield
96 finally:
97 self._set_asyncexc_append(None)
98
99 @property
100 def server(self) -> grpc.aio.Server:
101 """
102 The underlying [grpc.aio.Server](https://grpc.github.io/grpc/python/grpc_asyncio.html#grpc.aio.Server).
103 """
104 return self._server
105
106 async def start(self) -> None:
107 """
108 Starts the server.
109 @note Prefer starting mockserver using the async contextmanager syntax if possible.
110 """
111 await self._server.start()
112
113 async def stop(self) -> None:
114 """
115 Stops the server properly. Prefer this method to stopping `server` manually.
116 """
117 await _stop_server(self._server)
118
119 async def __aenter__(self) -> MockserverSession:
120 await self.start()
121 return self
122
123 async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
124 await self.stop()
125
126
128 """
129 Allows to install mocks that are reset between tests, see
130 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver".
131 """
132
133 def __init__(self, *, mockserver_session: MockserverSession, experimental: bool = False) -> None:
134 """
135 @warning This initializer is an **experimental API**, likely to break in the future. Consider using
136 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver" instead.
137
138 Initializes Mockserver.
139
140 @note `Mockserver` is usually obtained from
141 @ref pytest_userver.plugins.grpc.mockserver.grpc_mockserver "grpc_mockserver".
142 """
143 assert experimental
144 self._mockserver_session = mockserver_session
145
146 def __call__(self, servicer_method, /) -> MockDecorator:
147 """
148 Returns a decorator to mock the specified gRPC service method implementation.
149
150 Example:
151
152 @snippet samples/grpc_service/testsuite/test_grpc.py Prepare modules
153 @snippet samples/grpc_service/testsuite/test_grpc.py grpc client test
154 """
155 servicer_class = _get_class_from_method(servicer_method)
156 mock = self._mockserver_session._get_auto_service_mock(servicer_class) # pylint: disable=protected-access
157 return mock.install_handler(servicer_method.__name__)
158
159 def mock_factory(self, servicer_class, /) -> Callable[[str], MockDecorator]:
160 """
161 Allows to create a fixture as a shorthand for mocking methods of the specified gRPC service.
162
163 Example:
164
165 @snippet grpc/functional_tests/metrics/tests/conftest.py Prepare modules
166 @snippet grpc/functional_tests/metrics/tests/conftest.py Prepare server mock
167 @snippet grpc/functional_tests/metrics/tests/test_metrics.py grpc client test
168 """
169
170 def factory(method_name):
171 method = getattr(servicer_class, method_name, None)
172 if method is None:
173 raise ValueError(f'No method "{method_name}" in servicer class "{servicer_class.__name__}"')
174 return self(method)
175
176 _check_is_servicer_class(servicer_class)
177 return factory
178
179 def install_servicer(self, servicer: object, /) -> None:
180 """
181 Installs as a mock `servicer`, the class of which should inherit from a generated `*Servicer` class.
182
183 For example, @ref grpc/functional_tests/basic_chaos/tests-grpcclient/service.py "this servicer class"
184 can be installed as follows:
185
186 @snippet grpc/functional_tests/basic_chaos/tests-grpcclient/conftest.py installing mockserver servicer
187
188 @note Inheritance from multiple `*Servicer` classes at once is allowed.
189
190 @example grpc/functional_tests/basic_chaos/tests-grpcclient/service.py
191 """
192 base_servicer_classes = [cls for cls in inspect.getmro(type(servicer))[1:] if cls.__name__.endswith('Servicer')]
193 if not base_servicer_classes:
194 raise ValueError(f"Given object's type ({type(servicer)}) is not inherited from any grpc *Servicer class")
195 for servicer_class in base_servicer_classes:
196 # pylint: disable=protected-access
197 mock = self._mockserver_session._get_auto_service_mock(servicer_class)
198 for python_method_name in mock.known_methods:
199 handler_func = getattr(servicer, python_method_name)
200 mock.install_handler(python_method_name)(handler_func)
201
202
203# @cond
204
205
206async def _stop_server(server: grpc.aio.Server, /) -> None:
207 async def stop_server():
208 await server.stop(grace=None)
209 await server.wait_for_termination()
210
211 stop_server_task = asyncio.shield(asyncio.create_task(stop_server()))
212 # asyncio.shield does not protect our await from cancellations, and we
213 # really need to wait for the server stopping before continuing.
214 try:
215 await stop_server_task
216 except asyncio.CancelledError:
217 await stop_server_task
218 # Propagate cancellation when we are done
219 raise
220
221
222def _get_class_from_method(method) -> type:
223 # https://stackoverflow.com/a/54597033/5173839
224 assert inspect.isfunction(method), f'Expected an unbound method: foo(ClassName.MethodName), got: {method}'
225 class_name = method.__qualname__.split('.<locals>', 1)[0].rsplit('.', 1)[0]
226 try:
227 cls = getattr(inspect.getmodule(method), class_name)
228 except AttributeError:
229 cls = method.__globals__.get(class_name)
230 assert isinstance(cls, type)
231 return cls
232
233
234# @endcond