generated from mononoki/pyproj
347 lines
16 KiB
Python
347 lines
16 KiB
Python
"""Integration checks using the same client as examples/client.py."""
|
||
|
||
import asyncio
|
||
from copy import deepcopy
|
||
|
||
import pytest
|
||
|
||
|
||
class RunningServer:
|
||
def __init__(self, client_module):
|
||
self.client_module = client_module
|
||
self.clients = []
|
||
self.handlers = []
|
||
self.watchers = []
|
||
self.records = []
|
||
self.errors = []
|
||
self.expected_errors = []
|
||
self.peer_errors = {}
|
||
|
||
async def __aenter__(self):
|
||
self.initial_tasks = asyncio.all_tasks()
|
||
# Lazy imports keep client-only tests runnable if the server has an error.
|
||
from router import Router
|
||
from session import Sessions
|
||
from transport.tcp.handler import handle
|
||
|
||
errors = self.errors
|
||
watchers = self.watchers
|
||
records = self.records
|
||
peer_errors = self.peer_errors
|
||
|
||
class ObservedSessions(Sessions):
|
||
async def new(self, connection):
|
||
uid = await super().new(connection)
|
||
watchers.append(self[uid].watcher_task)
|
||
records.append(self[uid])
|
||
return uid
|
||
|
||
async def feed(self, uid, packet):
|
||
try:
|
||
await super().feed(uid, packet)
|
||
except Exception as error:
|
||
# Preserve the original error before the handler enters
|
||
# its delayed cleanup and closes the client's transport.
|
||
errors.append(error)
|
||
peer_errors.setdefault(asyncio.current_task(), []).append(error)
|
||
raise
|
||
|
||
self.router = Router()
|
||
self.sessions = ObservedSessions(self.router)
|
||
self.router.sessions = self.sessions
|
||
|
||
def connected(reader, writer):
|
||
task = asyncio.create_task(handle(reader, writer, self.sessions))
|
||
self.handlers.append(task)
|
||
task.add_done_callback(self.handler_finished)
|
||
|
||
self.server = await asyncio.start_server(connected, "127.0.0.1", 0)
|
||
self.port = self.server.sockets[0].getsockname()[1]
|
||
return self
|
||
|
||
def handler_finished(self, task):
|
||
if not task.cancelled() and task.exception() is not None:
|
||
self.errors.append(task.exception())
|
||
|
||
def accept_peer_rejection(self, task, error):
|
||
# Accept only the checked exception from this particular client handler.
|
||
assert task in self.handlers and task.done() and not task.cancelled()
|
||
assert task.exception() is error or any(error is item for item in self.peer_errors.get(task, []))
|
||
self.expected_errors.append(error)
|
||
|
||
def assert_no_unexpected_errors(self):
|
||
for error in self.errors:
|
||
if not any(error is expected for expected in self.expected_errors):
|
||
raise error
|
||
|
||
async def __aexit__(self, exc_type, exc, traceback):
|
||
self.server.close()
|
||
await self.server.wait_closed()
|
||
close_results = await asyncio.gather(
|
||
*(client.close() for client in self.clients), return_exceptions=True,
|
||
)
|
||
# Stop pending recovery-window waits after assertions, without modifying
|
||
# the production timer or waiting 60 seconds for each test connection.
|
||
tasks = set(self.handlers + self.watchers)
|
||
tasks.update(asyncio.all_tasks() - self.initial_tasks)
|
||
for task in tasks:
|
||
if not task.done():
|
||
task.cancel()
|
||
# A cancelled handler can enter finally and start the recovery wait.
|
||
# Let it close its transport, then cancel that new wait as well.
|
||
await asyncio.sleep(0)
|
||
for task in tasks:
|
||
if not task.done():
|
||
task.cancel()
|
||
results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 1)
|
||
for result in results:
|
||
if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError):
|
||
self.errors.append(result)
|
||
if exc_type is None:
|
||
self.assert_no_unexpected_errors()
|
||
for result in close_results:
|
||
if isinstance(result, BaseException):
|
||
raise result
|
||
|
||
async def client(
|
||
self, name, inputs=("control",), outputs=("telemetry",),
|
||
node_type="test", **recovery,
|
||
):
|
||
client = self.client_module.Client(
|
||
name, port=self.port, timeout=1,
|
||
inputs=inputs, outputs=outputs, node_type=node_type,
|
||
)
|
||
self.clients.append(client)
|
||
try:
|
||
await client.connect(**recovery)
|
||
except (TimeoutError, self.client_module.ProtocolError) as error:
|
||
if self.errors:
|
||
raise self.errors[0] from error
|
||
raise
|
||
return client
|
||
|
||
async def disconnect(self, client):
|
||
# A half-close proves that the SERVER closes the transport immediately.
|
||
client.writer.write_eof()
|
||
assert await asyncio.wait_for(client.reader.read(), 1) == b""
|
||
await client.close()
|
||
assert self.sessions[client.uid].state == "stop"
|
||
assert client.uid in self.sessions.sessions
|
||
|
||
|
||
def test_server_connection_telemetry_heartbeat_and_disconnect(client_module):
|
||
"""Проверяет handshake, метаданные, 10 пакетов, heartbeat и сохранение после EOF."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
client = await server.client("rover")
|
||
session = server.sessions[client.uid]
|
||
assert session.state == "air"
|
||
assert session.upf == client.upf
|
||
assert session.static_meta["name"] == "rover"
|
||
assert session.static_meta["type"] == "test"
|
||
assert session.static_meta["in"] == {"control"}
|
||
assert session.static_meta["out"] == {"telemetry"}
|
||
assert session.dynamic_meta["active_in"] == set()
|
||
assert session.dynamic_meta["active_out"] == set()
|
||
assert client.uid in server.router.session_routes
|
||
|
||
deadline = session.timeout
|
||
for index in range(10):
|
||
await client.telemetry(temperature=42.0 + index, battery=87 - index)
|
||
assert await client.heartbeat() == {"command": "heartbeat_answer"}
|
||
assert session.last_data == {
|
||
"telemetry": {"temperature": 42.0 + index, "battery": 87 - index},
|
||
}
|
||
assert session.timeout >= deadline
|
||
deadline = session.timeout
|
||
# Rapid heartbeats may share one monotonic clock tick. Exact renewal
|
||
# is checked separately with a controlled clock.
|
||
|
||
await server.disconnect(client)
|
||
assert session.last_data == {"telemetry": {"temperature": 51.0, "battery": 78}}
|
||
assert client.uid in server.router.session_routes
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 10))
|
||
|
||
|
||
def test_logical_errors_disconnect_before_handshake(client_module):
|
||
"""Проверяет закрытие TCP после пяти команд в неправильном состоянии протокола."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
client = client_module.Client(port=server.port, timeout=1)
|
||
server.clients.append(client)
|
||
client.reader, client.writer = await asyncio.open_connection("127.0.0.1", server.port)
|
||
# Deliberately answer a handshake that the server has not requested.
|
||
for _ in range(5):
|
||
await client.send("handshake_answer", {"status": True, "state": None})
|
||
assert await asyncio.wait_for(client.reader.read(), 1) == b""
|
||
assert len(server.sessions.sessions) == 1
|
||
session = next(iter(server.sessions.sessions.values()))
|
||
assert session.state == "stop"
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 5))
|
||
|
||
|
||
def test_recovery_preserves_data_and_moves_both_route_directions(client_module):
|
||
"""Проверяет восстановление данных и фактическую передачу по маршрутам в обе стороны."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
old = await server.client("old-rover")
|
||
peer = await server.client(
|
||
"peer", inputs=["telemetry"], outputs=["control"], node_type="operator",
|
||
)
|
||
# Route creation has no API yet; explicitly seed both directions.
|
||
server.router.session_routes[old.uid].add(peer.uid)
|
||
server.router.session_routes[peer.uid].add(old.uid)
|
||
await old.telemetry()
|
||
await old.heartbeat()
|
||
assert await peer.recv() == {
|
||
"command": "telemetry", "payload": {"temperature": 42.0, "battery": 87},
|
||
}
|
||
server.sessions[old.uid].dynamic_meta = {
|
||
"active_in": {"control"}, "active_out": {"telemetry"},
|
||
}
|
||
last_data = deepcopy(server.sessions[old.uid].last_data)
|
||
dynamic_meta = deepcopy(server.sessions[old.uid].dynamic_meta)
|
||
await server.disconnect(old)
|
||
|
||
restored = await server.client("restored-rover", old_token=old.uid, old_upf=old.upf)
|
||
assert restored.uid != old.uid
|
||
assert server.sessions[restored.uid].state == "air"
|
||
assert server.sessions[restored.uid].last_data == last_data
|
||
assert server.sessions[restored.uid].dynamic_meta == dynamic_meta
|
||
assert old.uid not in server.router.session_routes
|
||
assert set(server.router.session_routes[restored.uid]) == {peer.uid}
|
||
assert set(server.router.session_routes[peer.uid]) == {restored.uid}
|
||
|
||
await peer.send("control", {"turn": 1})
|
||
await peer.heartbeat()
|
||
assert await restored.recv() == {"command": "control", "payload": {"turn": 1}}
|
||
await restored.telemetry(43.0, 86)
|
||
await restored.heartbeat()
|
||
assert await peer.recv() == {
|
||
"command": "telemetry", "payload": {"temperature": 43.0, "battery": 86},
|
||
}
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 10))
|
||
|
||
|
||
@pytest.mark.parametrize("reason", ["active", "wrong_upf", "missing"])
|
||
def test_recovery_does_not_take_an_ineligible_session(client_module, reason):
|
||
"""Проверяет отказ переноса для активной/отсутствующей сессии и неверного upf."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
old = await server.client("old-rover")
|
||
await old.telemetry()
|
||
await old.heartbeat()
|
||
if reason != "active":
|
||
await server.disconnect(old)
|
||
uid = -1 if reason == "missing" else old.uid
|
||
upf = old.upf + 1 if reason == "wrong_upf" else old.upf
|
||
fresh = await server.client("fresh-rover", old_token=uid, old_upf=upf)
|
||
assert server.sessions[fresh.uid].last_data == {}
|
||
assert old.uid in server.router.session_routes
|
||
assert server.sessions[old.uid].state == ("air" if reason == "active" else "stop")
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 10))
|
||
|
||
|
||
def test_old_recovery_proof_can_only_be_used_once(client_module):
|
||
"""Проверяет, что повторный старый UID/upf не забирает уже восстановленные данные."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
old = await server.client("old-rover")
|
||
await old.telemetry()
|
||
await old.heartbeat()
|
||
await server.disconnect(old)
|
||
restored = await server.client("restored-rover", old_token=old.uid, old_upf=old.upf)
|
||
expected = {"telemetry": {"temperature": 42.0, "battery": 87}}
|
||
assert server.sessions[restored.uid].last_data == expected
|
||
repeated = await server.client("repeated-rover", old_token=old.uid, old_upf=old.upf)
|
||
assert server.sessions[repeated.uid].last_data == {}
|
||
assert server.sessions[restored.uid].last_data == expected
|
||
assert restored.uid in server.router.session_routes
|
||
assert old.uid not in server.router.session_routes
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 10))
|
||
|
||
|
||
async def expect_timeout_close(client, session):
|
||
read = asyncio.create_task(client.reader.read())
|
||
try:
|
||
done, _ = await asyncio.wait(
|
||
{read, session.watcher_task}, timeout=1,
|
||
return_when=asyncio.FIRST_COMPLETED,
|
||
)
|
||
assert done, "Воркер не обработал истечение дедлайна"
|
||
if session.watcher_task in done:
|
||
# Expose a crashed watcher rather than just a later socket timeout.
|
||
await session.watcher_task
|
||
assert await asyncio.wait_for(read, 1) == b""
|
||
assert session.state == "stop"
|
||
finally:
|
||
if not read.done():
|
||
read.cancel()
|
||
await asyncio.gather(read, return_exceptions=True)
|
||
|
||
|
||
@pytest.mark.parametrize("state", ["new", "token", "ready", "air"])
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_server_timeout_closes_transport_in_each_phase(client_module, timing, state):
|
||
"""Проверяет закрытие настоящего TCP по таймауту на каждом этапе подключения."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
client = client_module.Client(port=server.port, timeout=1)
|
||
server.clients.append(client)
|
||
client.reader, client.writer = await asyncio.open_connection("127.0.0.1", server.port)
|
||
async with asyncio.timeout(1):
|
||
while not server.sessions.sessions:
|
||
await asyncio.sleep(0)
|
||
session = next(iter(server.sessions.sessions.values()))
|
||
if state in {"token", "ready", "air"}:
|
||
await client.send("connection", {"name": "idle-rover", "type": "test", "in": [], "out": ["telemetry"]})
|
||
token = await client.recv()
|
||
assert token["command"] == "token_ask"
|
||
if state in {"ready", "air"}:
|
||
await client.send("token_answer", {"token": None, "upf": None})
|
||
assert await client.recv() == {"command": "handshake_ask"}
|
||
if state == "air":
|
||
await client.send("handshake_answer", {"status": True, "state": None})
|
||
await client.heartbeat()
|
||
assert session.state == state
|
||
request = await timing.sleeping("session", session.watcher_task)
|
||
timing.clock.advance(session.timeout - timing.clock())
|
||
request.release()
|
||
await expect_timeout_close(client, session)
|
||
assert server.sessions[session.uid] is session
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 5))
|
||
|
||
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_server_heartbeat_prevents_stale_deadline_disconnect(client_module, timing):
|
||
"""Проверяет по TCP продление heartbeat и последующее отключение при тишине."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
client = await server.client("rover")
|
||
session = server.sessions[client.uid]
|
||
request = await timing.sleeping("session", session.watcher_task)
|
||
old_deadline = session.timeout
|
||
timing.clock.advance(4)
|
||
await client.heartbeat()
|
||
assert session.timeout > old_deadline
|
||
timing.clock.advance(old_deadline - timing.clock())
|
||
request.release()
|
||
await asyncio.sleep(0)
|
||
if session.watcher_task.done():
|
||
await session.watcher_task
|
||
assert session.state == "air"
|
||
await client.heartbeat()
|
||
request = await timing.sleeping("session", session.watcher_task)
|
||
timing.clock.advance(session.timeout - timing.clock())
|
||
request.release()
|
||
await expect_timeout_close(client, session)
|
||
assert client.uid in server.sessions.sessions
|
||
|
||
asyncio.run(asyncio.wait_for(scenario(), 5))
|