async def test_04_multiple(self, counters):
counters.increment('abc')
counters.increment('xyz')
assert counters.total_increments('abc') == 1
assert counters.total_increments('xyz') == 1
task_abc = create_task(counters.wait_for_total_increments('abc', 2))
task_xyz = create_task(counters.wait_for_total_increments('xyz', 2))
await asyncio.sleep(0)
assert not task_abc.done()
assert not task_xyz.done()
counters.decrement('abc')
assert counters.total_increments('abc') == 1
assert counters.total_increments('xyz') == 1
await asyncio.sleep(0)
try:
assert not task_abc.done()
assert not task_xyz.done()
counters.increment('xyz')
assert counters.total_increments('xyz') == 2
await asyncio.sleep(0)
assert task_xyz.done()
assert not task_abc.done()
finally:
with pytest.raises(asyncio.TimeoutError):
await asyncio.wait_for(task_abc, timeout=1)
await asyncio.wait_for(task_xyz, timeout=1)
def __init__(self,
protocol_factory: DatagramServerProtocolFactory,
host: str,
port: int,
family: int = 0,
proto: int = 0,
flags: int = 0,
sock=None,
start_serving: bool = True,
reuse_port: bool = None,
allow_broadcast: bool = None,
loop=None):
self._loop = loop or asyncio.get_event_loop()
self._sock = sock
self._host = host
self._port = port
self._family = family
self._proto = proto
self._flags = flags
self._reuse_port = reuse_port
self._allow_broadcast = allow_broadcast
self._protocol_factory = protocol_factory
self._transport: Optional[asyncio.DatagramTransport] = None
self._protocol: Optional[UDPServerConnection] = None
self._status = StatusWaiter()
if start_serving:
create_task(self.start_serving())
self._status.set_starting()
async def task() -> asyncio.Task:
async def coro():
...
if supports_task_name():
task = create_task(coro(), name="Task-99")
else:
task = create_task(coro())
yield task
await task
async def test_00_connection_lifecycle(
self, protocol_factory_started, connection, echo_response_encoded,
transport, echo_encoded, server_sock, client_sock,
connections_manager, queue, echo_response_object, client_sock_str,
server_sock_str):
protocol_factory = protocol_factory_started()
assert protocol_factory == protocol_factory_started
protocol_factory.connection_made(transport)
transport.set_protocol(protocol_factory)
conn = protocol_factory.new_peer()
assert connections_manager.total == 1
task = create_task(conn.echo())
msg = await queue.get()
assert msg == (server_sock, echo_encoded)
protocol_factory.datagram_received(echo_response_encoded, server_sock)
response = await asyncio.wait_for(task, timeout=1)
assert response == echo_response_object
assert connections_manager.total == 1
full_peername = f"udp_{client_sock_str}_{server_sock_str}"
new_connection = connections_manager.get(full_peername)
assert new_connection.is_connected()
assert protocol_factory_started.logger == new_connection.logger
assert protocol_factory_started.is_owner(new_connection)
assert id(connections_manager.get(full_peername)) == id(new_connection)
transport.close()
await asyncio.wait_for(protocol_factory.close(), timeout=1)
await asyncio.wait_for(new_connection.wait_closed(), timeout=1)
assert connections_manager.total == 0
def finish_connection(self, exc: Optional[BaseException]) -> None:
if self._adaptor:
self._adaptor.logger.info('Finishing connection')
self._delete_connection()
self._status.set_stopping()
close_task = create_task(self._close(exc))
set_task_name(close_task, f"Close:{self.peer}")
async def test_04_server_stop_wait(self, server_started):
assert server_started.is_started()
task = create_task(server_started.wait_stopped())
await asyncio.sleep(0)
assert not task.done()
await server_started.close()
assert server_started.is_closing()
await asyncio.wait_for(task, timeout=4)
async def test_00_send_data_and_wait(self, connection_connected, echo_encoded, echo_response_encoded,
echo_response_object, transport, queue):
task = create_task(connection_connected.send_data_and_wait(1, echo_encoded))
receiver, msg = await asyncio.wait_for(queue.get(), timeout=1)
assert msg == echo_encoded
connection_connected.data_received(echo_response_encoded)
result = await asyncio.wait_for(task, timeout=1)
assert result == echo_response_object
async def test_02_requester(self, connection_connected, echo_encoded, echo_response_encoded, echo_response_object,
queue):
task = create_task(connection_connected.echo())
receiver, msg = await asyncio.wait_for(queue.get(), timeout=1)
assert msg == echo_encoded
connection_connected.data_received(echo_response_encoded)
result = await asyncio.wait_for(task, timeout=1)
assert result == echo_response_object
async def signal_server_manager_started(
tmp_config_file) -> SignalServerManager:
server_manager = SignalServerManager(tmp_config_file)
task = create_task(server_manager.serve_until_stopped())
await server_manager.wait_server_started()
yield server_manager
server_manager.close()
await server_manager.wait_server_stopped()
await task
async def test_00_server_start(self, tcp_server_ssl, capsys):
assert not tcp_server_ssl.is_started()
task = create_task(tcp_server_ssl.start())
await asyncio.wait_for(tcp_server_ssl.wait_started(), timeout=4)
assert tcp_server_ssl.is_started()
captured = capsys.readouterr()
assert captured.out.startswith("Storing")
assert "Serving" in captured.out
await asyncio.wait_for(task, timeout=4)
def close(self, immediate: bool = False):
if not self.transport.is_closing():
if immediate:
self.transport.abort()
elif self._adaptor:
task = create_task(self.wait_current_tasks())
task.add_done_callback(self._close_transport)
else:
self.transport.close()
async def test_03_server_wait_started(self, server, capsys):
assert not server.is_started()
task = create_task(server.wait_started())
await asyncio.sleep(0)
assert not task.done()
await asyncio.wait_for(server.start(), 3)
assert server.is_started()
await asyncio.wait_for(task, timeout=4)
captured = capsys.readouterr()
assert "Serving" in captured.out
async def test_02_future_lifecycle(self, task_scheduler):
await asyncio.wait_for(task_scheduler.close(), timeout=1)
fut = task_scheduler.create_future(1)
task = create_task(task_scheduler.close())
await asyncio.sleep(0)
assert not task.done()
task_scheduler.set_result(1, "abc")
await fut
assert fut.result() == 'abc'
task_scheduler.future_done(1)
def context(self) -> Optional[SSLContext]:
if self.ssl:
self.logger.info("Setting up SSL")
context = SSLContext(PROTOCOL_TLS)
if self.cert and self.key:
self.logger.info("Using SSL Cert: %s", self.cert)
try:
context.load_cert_chain(str(self.cert),
str(self.key),
password=self.key_password)
except FileNotFoundError as e:
raise FileNotFoundError(
better_file_not_found_error(
self.cert, self.key, purpose='ssl cert loading'))
if self.warn_if_expires_before_days:
self._warn_expiry_task = create_task(
self.check_cert_expiry())
set_task_name(self._warn_expiry_task,
'CheckSSLCertValidity')
context.verify_mode = CERT_REQUIRED if self.cert_required else CERT_NONE
context.check_hostname = self.check_hostname
self.logger.info('%s, Check Hostname: %s' %
(context.verify_mode, context.check_hostname))
if context.verify_mode != CERT_NONE:
if self.cafile or self.capath or self.cadata:
locations = {
'cafile': str(self.cafile) if self.cafile else None,
'capath': str(self.capath) if self.capath else None,
'cadata': self.cadata
}
try:
context.load_verify_locations(**locations)
self.logger.info("Verifying SSL certs with: %s",
locations)
except FileNotFoundError:
raise FileNotFoundError(
better_file_not_found_error(
*locations.values(),
purpose='CA ssl cert validation'))
else:
context.load_default_certs(self.purpose)
self.logger.info("Verifying SSL certs with: %s",
get_default_verify_paths())
self.logger.info("SSL Context loaded")
# OpenSSL 1.1.1 keylog file
if hasattr(context, 'keylog_filename'):
keylogfile = os.environ.get('SSLKEYLOGFILE')
if keylogfile and not sys.flags.ignore_environment:
self.logger.warning(
"TLS encryption secrets are being stored in %s",
keylogfile)
context.keylog_filename = keylogfile
return context
return None
async def test_08_serve_until_close_signal(self, server_quiet, signal_num, patch_systemd, server_port):
assert not server_quiet.is_started()
task = create_task(server_quiet.serve_until_close_signal())
await asyncio.wait_for(server_quiet.wait_started(), timeout=4)
status = next(server_quiet.listening_on_msgs)
os.kill(os.getpid(), signal_num)
await asyncio.wait_for(task, timeout=4)
assert server_quiet.is_closed()
if patch_systemd:
daemon, journal = patch_systemd
daemon.notify.assert_has_calls([status_call(status), ready_call, stopping_call])
assert daemon.notify.call_count == 3
async def test_03_counter_wait_decrement(self, counter):
counter.increment()
assert counter.num == 1
task = create_task(counter.wait_for(0))
await asyncio.sleep(0)
assert not task.done()
counter.decrement()
await asyncio.sleep(0)
assert counter.num == 0
try:
assert task.done()
finally:
await asyncio.wait_for(task, timeout=1)
async def test_03_counters_wait_decrement(self, counters):
counters.increment('abc')
assert counters.get_num('abc') == 1
task = create_task(counters.wait_for('abc', 0))
await asyncio.sleep(0)
assert not task.done()
counters.decrement('abc')
await asyncio.sleep(0)
assert counters.get_num('abc') == 0
try:
assert task.done()
finally:
await asyncio.wait_for(task, timeout=1)
async def test_02_counter_wait_increment(self, counter):
task = create_task(counter.wait_for_total_increments(2))
counter.increment()
assert counter.total_increments == 1
await asyncio.sleep(0)
assert not task.done()
counter.increment()
assert counter.total_increments == 2
await asyncio.sleep(0)
try:
assert task.done()
finally:
await asyncio.wait_for(task, timeout=1)
async def test_04_close_signal(self, patch_systemd, signal_server_manager,
signal_num, capsys, reset_logging):
task = create_task(self.send_signal(signal_server_manager, signal_num))
await signal_server_manager.serve_until_stopped()
await task
await asyncio.wait_for(signal_server_manager.server.wait_stopped(),
timeout=10)
if patch_systemd:
out = capsys.readouterr().out
port = port_from_out(out)
daemon, journal = patch_systemd
daemon.notify.assert_has_calls(
[status_call(port), ready_call, stopping_call])
assert daemon.notify.call_count == 3
async def test_04_multiple(self, counters):
counters.increment('abc')
counters.increment('xyz')
assert counters.get_num('abc') == 1
assert counters.get_num('xyz') == 1
task_abc = create_task(counters.wait_for('abc', 0))
task_xyz = create_task(counters.wait_for('xyz', 0))
await asyncio.sleep(0)
assert not task_abc.done()
assert not task_xyz.done()
counters.decrement('abc')
assert counters.get_num('abc') == 0
await asyncio.sleep(0)
try:
assert task_abc.done()
assert not task_xyz.done()
counters.decrement('xyz')
assert counters.get_num('xyz') == 0
await asyncio.sleep(0)
assert task_xyz.done()
finally:
await asyncio.wait_for(task_abc, timeout=1)
await asyncio.wait_for(task_xyz, timeout=1)
async def test_01_serve_forever(self, server, capsys):
assert not server.is_started()
task = create_task(server.serve_forever())
await asyncio.wait_for(server.wait_started(), timeout=4)
assert server.is_started()
captured = capsys.readouterr()
assert "Serving" in captured.out
await asyncio.sleep(2)
assert not task.done()
assert server.is_started()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert server.is_closed()
async def test_00_multiple_connections(self, connections_manager,
simple_network_connection):
parent_name = simple_network_connection.parent_name
peer = simple_network_connection.peer
task1 = create_task(
connections_manager.wait_num_connections(parent_name, 1))
await simple_network_connection.wait_all_messages_processed()
await asyncio.sleep(0)
assert not task1.done()
connections_manager.add_connection(simple_network_connection)
assert connections_manager.num_connections(parent_name) == 1
connection = connections_manager.get(peer)
assert connection == simple_network_connection
await asyncio.sleep(0)
assert task1.done()
task = create_task(
connections_manager.wait_num_connections(parent_name, 0))
await asyncio.sleep(0)
assert not task.done()
connections_manager.remove_connection(simple_network_connection)
connections_manager.decrement(simple_network_connection)
assert connections_manager.num_connections(parent_name) == 0
await asyncio.wait_for(task, timeout=1)
await connections_manager.wait_num_has_connected(parent_name, 1)
def create_task(self,
coro: Awaitable,
name: str = None,
include_hierarchy: bool = True,
separator: str = ':',
continuous: bool = False) -> asyncio.Future:
self._counter.increment()
task = create_task(coro)
set_task_name(task,
name,
include_hierarchy=include_hierarchy,
separator=separator)
if not continuous:
self._current_tasks.append(task)
return task
async def process_msgs(self, msgs: AsyncIterator[MessageObjectType],
buffer: bytes) -> None:
tasks = []
try:
async for msg_obj in msgs:
if not self.action.filter(msg_obj):
task = create_task(self._process_msg(msg_obj))
set_task_name(task, f'Process {msg_obj}')
tasks.append(task)
else:
self.logger.on_msg_filtered(msg_obj)
await asyncio.wait_for(asyncio.gather(*tasks),
timeout=self.action.task_timeout)
except Exception as exc:
self._on_decoding_error(buffer, exc)
raise
async def test_03_counter_wait_decrement(self, counters):
counters.increment('abc')
assert counters.total_increments('abc') == 1
task = create_task(counters.wait_for_total_increments('abc', 2))
await asyncio.sleep(0)
assert not task.done()
counters.decrement('abc')
await asyncio.sleep(0)
assert counters.total_increments('abc') == 1
assert not task.done()
counters.increment('abc')
assert counters.total_increments('abc') == 2
await asyncio.sleep(0)
try:
assert task.done()
finally:
await asyncio.wait_for(task, timeout=1)
def call_coro_periodic(self,
interval: Union[int, float],
async_callback: Callable,
*args,
fixed_start_time: bool = False,
immediate: bool = False,
task_name: str = None,
**kwargs) -> None:
start_time_interval = self.get_start_interval(fixed_start_time,
immediate, interval)
task = create_task(
self._call_coro_periodic(interval,
async_callback,
start_time_interval=start_time_interval,
*args,
**kwargs))
set_task_name(task, task_name)
self._periodic_tasks.append(task)
async def test_00_start_close(self, signal_server_manager, patch_systemd,
server_sock, capsys, reset_logging):
task = create_task(signal_server_manager.serve_until_stopped())
try:
await asyncio.wait_for(signal_server_manager.wait_server_started(),
timeout=2)
signal_server_manager.close()
except asyncio.TimeoutError:
await asyncio.wait_for(task, timeout=2)
else:
await asyncio.wait_for(task, timeout=1)
await asyncio.wait_for(signal_server_manager.wait_server_stopped(),
timeout=1)
if patch_systemd:
daemon, journal = patch_systemd
out = capsys.readouterr().out
port = port_from_out(out)
daemon.notify.assert_has_calls(
[status_call(port), ready_call, stopping_call])
assert daemon.notify.call_count == 3
async def test_07_reload_no_file_signal(self, patch_systemd,
signal_server_manager, signal_num,
tmp_config_file, capsys,
reset_logging):
task = create_task(signal_server_manager.serve_until_stopped())
await signal_server_manager.wait_server_started()
tmp_config_file.unlink()
await self.send_signal(signal_server_manager, signal_num)
await signal_server_manager._stop_event.wait()
assert not signal_server_manager._restart_event.is_set()
await asyncio.wait_for(signal_server_manager.wait_server_stopped(),
timeout=2)
await asyncio.wait_for(task, timeout=1)
if patch_systemd:
out = capsys.readouterr().out
port = port_from_out(out)
daemon, journal = patch_systemd
daemon.notify.assert_has_calls(
[status_call(port), ready_call, stopping_call])
assert daemon.notify.call_count == 3
async def test_08_run_method_and_wait(self, adaptor, echo_encoded, timestamp,
echo_response_encoded, echo_response_object, queue):
task1 = create_task(adaptor.echo())
await self.assert_response(queue, echo_encoded, adaptor, echo_response_encoded, timestamp, task1,
echo_response_object)
async def test_07_encode_send_wait(self, adaptor, echo, echo_encoded, timestamp,
echo_response_encoded, echo_response_object, queue):
task1 = create_task(adaptor.encode_send_wait(echo))
await self.assert_response(queue, echo_encoded, adaptor, echo_response_encoded, timestamp, task1,
echo_response_object)
