userver: /data/code/service_template/third_party/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
11import socket
12
13import grpc
14import pytest
15
16from testsuite.utils import callinfo
17
18DEFAULT_PORT = 8091
19
20
22 def __init__(self, servicer, methods):
23 self.servicer = servicer
24 self._known_methods = methods
25 self._methods = {}
26
27 def get(self, method, default):
28 return self._methods.get(method, default)
29
30 def reset_handlers(self):
31 self._methods = {}
32
33 @contextlib.contextmanager
34 def mock(self):
35 try:
36 yield self.install_handler
37 finally:
38 self._methods = {}
39
40 def install_handler(self, method: str):
41 def decorator(func):
42 if method not in self._known_methods:
43 raise RuntimeError(
44 f'Trying to mock unknown grpc method {method}',
45 )
46
47 wrapped = callinfo.acallqueue(func)
48 self._methods[method] = wrapped
49 return wrapped
50
51 return decorator
52
53
54@pytest.fixture(scope='session')
55def _grpc_mockserver_endpoint(pytestconfig):
56 port = pytestconfig.option.grpc_mockserver_port
57 if pytestconfig.option.service_wait or pytestconfig.option.service_disable:
58 port = port or DEFAULT_PORT
59 if port == 0:
60 port = _find_free_port()
61 return f'{pytestconfig.option.grpc_mockserver_host}:{port}'
62
63
64@pytest.fixture(scope='session')
65def grpc_mockserver_endpoint(pytestconfig, _grpc_port) -> str:
66 """
67 Returns the gRPC endpoint to start the mocking server that is set by
68 command line `--grpc-mockserver-host` and `--grpc-mockserver-port` options.
69
70 Override this fixture to change the way the gRPC endpoint
71 is detected by the testsuite.
72
73 @snippet samples/grpc_service/tests/conftest.py Prepare configs
74 @ingroup userver_testsuite_fixtures
75 """
76 return f'{pytestconfig.option.grpc_mockserver_host}:{_grpc_port}'
77
78
79def _find_free_port() -> int:
80 with contextlib.closing(
81 socket.socket(socket.AF_INET6, socket.SOCK_STREAM),
82 ) as sock:
83 sock.bind(('', 0))
84 sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
85 return sock.getsockname()[1]
86
87
88@pytest.fixture(scope='session')
89def _grpc_port(_grpc_mockserver_endpoint):
90 _, port = _grpc_mockserver_endpoint.rsplit(':', 1)
91 return int(port)
92
93
94@pytest.fixture(scope='session')
95async def grpc_mockserver(_grpc_mockserver_endpoint):
96 """
97 Returns the gRPC mocking server.
98
99 Override this fixture to change the way the gRPC
100 mocking server is started by the testsuite.
101
102 @snippet samples/grpc_service/tests/conftest.py Prepare server mock
103 @ingroup userver_testsuite_fixtures
104 """
105 server = grpc.aio.server()
106 server.add_insecure_port(_grpc_mockserver_endpoint)
107 server_task = asyncio.create_task(server.start())
108
109 try:
110 yield server
111 finally:
112 await server.stop(grace=None)
113 await server.wait_for_termination()
114 await server_task
115
116
117@pytest.fixture(scope='session')
119 """
120 Creates the gRPC mock server for the provided type.
121
122 @snippet samples/grpc_service/tests/conftest.py Prepare server mock
123 @ingroup userver_testsuite_fixtures
124 """
125 return _create_servicer_mock
126
127
128def pytest_addoption(parser):
129 group = parser.getgroup('grpc-mockserver')
130 group.addoption(
131 '--grpc-mockserver-host',
132 default='[::]',
133 help='gRPC mockserver hostname, default is [::]',
134 )
135 group.addoption(
136 '--grpc-mockserver-port',
137 type=int,
138 default=0,
139 help='gRPC mockserver port, by default random port is used',
140 )
141
142
143def _create_servicer_mock(servicer_class, stream_method_names=None):
144 def wrap_grpc_method(name, default_method):
145 @functools.wraps(default_method)
146 async def run_method(self, *args, **kwargs):
147 method = mock.get(name, None)
148 if method is not None:
149 call = method(*args, **kwargs)
150 else:
151 call = default_method(self, *args, **kwargs)
152
153 return await call
154
155 @functools.wraps(default_method)
156 async def run_stream_method(self, *args, **kwargs):
157 method = mock.get(name, None)
158 async for response in await method(*args, **kwargs):
159 yield response
160
161 if name in (stream_method_names or []):
162 return run_stream_method
163 else:
164 return run_method
165
166 methods = {}
167 for attname, value in servicer_class.__dict__.items():
168 if callable(value):
169 methods[attname] = wrap_grpc_method(attname, value)
170
171 mocked_servicer_class = type(
172 f'Mock{servicer_class.__name__}', (servicer_class,), methods,
173 )
174 servicer = mocked_servicer_class()
175 mock = GrpcServiceMock(servicer, frozenset(methods))
176 return mock