userver: /home/antonyzhilin/arcadia/taxi/uservices/userver/testsuite/pytest_plugins/pytest_userver/utils/net.py Source File
All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Modules Pages Concepts
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 [asyncio.Task(_check_tcp_port_availability(val)) for val in checks.tcp],
58 timeout=25.0,
59 return_when=asyncio.ALL_COMPLETED,
60 )
61
62 for task in pending:
63 task.cancel()
64
65 for task in pending:
66 await task
67
68 return not pending and all(task.result() for task in done)
69
70
71def get_health_checks_info(service_config: dict) -> HealthChecks:
72 """
73 Returns a health checks info that for server.listener, grpc-server.port
74 and server.listener-monitor.
75
76 @see pytest_userver.plugins.base.service_non_http_health_checks()
77
78 @ingroup userver_testsuite
79 """
80 checks = HealthChecks()
81
82 components = service_config['components_manager']['components']
83 server = components.get('server')
84 if server:
85 for listener_name in ('listener-monitor', 'listener'):
86 listener_config = server.get(listener_name, {})
87 port = int(listener_config.get('port', 0))
88 if port:
89 host = listener_config.get('address', 'localhost')
90 if host == '::':
91 host = 'localhost'
92 checks.tcp.append(HostPort(host, port))
93
94 port = int(components.get('grpc-server', {}).get('port', 0))
95 if port:
96 checks.tcp.append(HostPort('localhost', port))
97
98 return checks