userver: /data/code/userver/testsuite/pytest_plugins/pytest_userver/plugins/grpc/mockserver.py Source File
Loading...
Searching...
No Matches
mockserver.py
1"""
2Mocks for the gRPC servers.
3
4@sa @ref scripts/docs/en/userver/tutorial/grpc_service.md
5"""
6
7# pylint: disable=no-member
8import asyncio
9import contextlib
10import functools
11from typing import Callable
12from typing import List
13from typing import Optional
14
15import grpc
16import pytest
17
18from testsuite.utils import callinfo
19
20DEFAULT_PORT = 8091
21
22
24 def __init__(self, servicer, methods):
25 self.servicer = servicer
26 self._known_methods = methods
27 self._methods = {}
28
29 def get(self, method, default):
30 return self._methods.get(method, default)
31
32 def reset_handlers(self):
33 self._methods = {}
34
35 @contextlib.contextmanager
36 def mock(self):
37 try:
38 yield self.install_handler
39 finally:
40 self._methods = {}
41
42 def install_handler(self, method: str):
43 def decorator(func):
44 if method not in self._known_methods:
45 raise RuntimeError(
46 f'Trying to mock unknown grpc method {method}',
47 )
48
49 wrapped = callinfo.acallqueue(func)
50 self._methods[method] = wrapped
51 return wrapped
52
53 return decorator
54
55
56@pytest.fixture(scope='session')
57def grpc_mockserver_endpoint(pytestconfig, get_free_port) -> str:
58 """
59 Returns the gRPC endpoint to start the mocking server that is set by
60 command line `--grpc-mockserver-host` and `--grpc-mockserver-port` options.
61
62 For port 0, picks some free port.
63
64 Override this fixture to customize the endpoint used by gRPC mockserver.
65
66 @snippet samples/grpc_service/tests/conftest.py Prepare configs
67 @ingroup userver_testsuite_fixtures
68 """
69 port = pytestconfig.option.grpc_mockserver_port
70 if pytestconfig.option.service_wait or pytestconfig.option.service_disable:
71 port = port or DEFAULT_PORT
72 if port == 0:
73 port = get_free_port()
74 return f'{pytestconfig.option.grpc_mockserver_host}:{port}'
75
76
77@pytest.fixture(scope='session')
78async def grpc_mockserver(grpc_mockserver_endpoint) -> grpc.aio.Server:
79 """
80 Returns the gRPC mocking server.
81
82 @snippet samples/grpc_service/tests/conftest.py Prepare server mock
83 @ingroup userver_testsuite_fixtures
84 """
85 server = grpc.aio.server()
86 server.add_insecure_port(grpc_mockserver_endpoint)
87 server_task = asyncio.create_task(server.start())
88
89 try:
90 yield server
91 finally:
92 await asyncio.shield(server.stop(grace=None))
93 await asyncio.shield(server.wait_for_termination())
94 await asyncio.shield(server_task)
95
96
97@pytest.fixture(scope='session')
99 """
100 Creates the gRPC mock server for the provided type.
101
102 @snippet samples/grpc_service/tests/conftest.py Prepare server mock
103 @ingroup userver_testsuite_fixtures
104 """
105 return _create_servicer_mock
106
107
108def pytest_addoption(parser):
109 group = parser.getgroup('grpc-mockserver')
110 group.addoption(
111 '--grpc-mockserver-host',
112 default='[::]',
113 help='gRPC mockserver hostname, default is [::]',
114 )
115 group.addoption(
116 '--grpc-mockserver-port',
117 type=int,
118 default=0,
119 help='gRPC mockserver port, by default random port is used',
120 )
121
122
123def _create_servicer_mock(
124 servicer_class: type,
125 *,
126 stream_method_names: Optional[List[str]] = None,
127 before_call_hook: Optional[Callable[..., None]] = None,
128) -> GrpcServiceMock:
129 def wrap_grpc_method(name, default_method, is_stream):
130 @functools.wraps(default_method)
131 async def run_method(self, *args, **kwargs):
132 if before_call_hook is not None:
133 before_call_hook(*args, **kwargs)
134
135 method = mock.get(name, None)
136 if method is not None:
137 call = method(*args, **kwargs)
138 else:
139 call = default_method(self, *args, **kwargs)
140
141 return await call
142
143 @functools.wraps(default_method)
144 async def run_stream_method(self, *args, **kwargs):
145 if before_call_hook is not None:
146 before_call_hook(*args, **kwargs)
147
148 method = mock.get(name, None)
149 async for response in await method(*args, **kwargs):
150 yield response
151
152 return run_stream_method if is_stream else run_method
153
154 methods = {}
155 for attname, value in servicer_class.__dict__.items():
156 if callable(value):
157 methods[attname] = wrap_grpc_method(
158 name=attname,
159 default_method=value,
160 is_stream=attname in (stream_method_names or []),
161 )
162
163 mocked_servicer_class = type(
164 f'Mock{servicer_class.__name__}', (servicer_class,), methods,
165 )
166 servicer = mocked_servicer_class()
167 mock = GrpcServiceMock(servicer, frozenset(methods))
168 return mock