userver: /data/code/userver/testsuite/pytest_plugins/pytest_userver/utils/net.py Source File
Loading...
Searching...
No Matches
net.py
1# pylint: disable=no-member
2import asyncio
3import dataclasses
4import logging
5import typing
6
7# @cond
8
9
10logger = logging.getLogger(__name__)
11
12
13# @endcond
14
15
16@dataclasses.dataclass(frozen=True)
18 """
19 Class that holds a host and port.
20
21 @ingroup userver_testsuite
22 """
23
24 host: str
25 port: int
26
27
28@dataclasses.dataclass(frozen=False)
30 """
31 Class that holds all the info for health checks.
32
33 @ingroup userver_testsuite
34 """
35
36 tcp: typing.List[HostPort] = dataclasses.field(default_factory=list)
37
38
39async def _check_tcp_port_availability(tcp: HostPort) -> bool:
40 try:
41 _, writer = await asyncio.open_connection(tcp.host, tcp.port)
42 writer.close()
43 await writer.wait_closed()
44 except (OSError, asyncio.TimeoutError):
45 return False
46 return True
47
48
49async def check_availability(checks: HealthChecks) -> bool:
50 """
51 Checks availability for each provided entry.
52
53 @ingroup userver_testsuite
54 """
55 assert checks.tcp
56 done, pending = await asyncio.wait(
57 [
58 asyncio.Task(_check_tcp_port_availability(val))
59 for val in checks.tcp
60 ],
61 timeout=25.0,
62 return_when=asyncio.ALL_COMPLETED,
63 )
64
65 for task in pending:
66 task.cancel()
67
68 for task in pending:
69 await task
70
71 return not pending and all(task.result() for task in done)
72
73
74def get_health_checks_info(service_config: dict) -> HealthChecks:
75 """
76 Returns a health checks info that for server.listener, grpc-server.port
77 and server.listener-monitor.
78
79 @see pytest_userver.plugins.base.service_non_http_health_checks()
80
81 @ingroup userver_testsuite
82 """
83 checks = HealthChecks()
84
85 components = service_config['components_manager']['components']
86 server = components.get('server')
87 if server:
88 for listener_name in ('listener-monitor', 'listener'):
89 listener_config = server.get(listener_name, {})
90 port = int(listener_config.get('port', 0))
91 if port:
92 host = listener_config.get('address', 'localhost')
93 if host == '::':
94 host = 'localhost'
95 checks.tcp.append(HostPort(host, port))
96
97 port = int(components.get('grpc-server', {}).get('port', 0))
98 if port:
99 checks.tcp.append(HostPort('localhost', port))
100
101 return checks