generated from mononoki/pyproj
397 lines
19 KiB
Python
397 lines
19 KiB
Python
"""Настоящий TCP: некорректные пакеты, фрейминг и фоновые отправки."""
|
||
|
||
import asyncio
|
||
import json
|
||
|
||
import pytest
|
||
|
||
from test_client_run import RunningServer
|
||
|
||
|
||
def run(scenario):
|
||
asyncio.run(asyncio.wait_for(scenario(), 5))
|
||
|
||
|
||
async def open_at_state(server, state="new", connection_payload=None):
|
||
count = len(server.records)
|
||
client = server.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 len(server.records) == count:
|
||
await asyncio.sleep(0)
|
||
session = server.records[-1]
|
||
if state in {"token", "ready", "air"}:
|
||
await client.send("connection", connection_payload if connection_payload is not None else {
|
||
"name": "raw-client", "type": "test", "in": [], "out": [],
|
||
})
|
||
token = await client.recv()
|
||
assert token["command"] == "token_ask"
|
||
client.uid, client.upf = token["payload"], token["upf"]
|
||
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})
|
||
await client.heartbeat()
|
||
assert session.state == state
|
||
return client, session
|
||
|
||
|
||
async def assert_bad_peer_isolated(
|
||
server, bad, session, healthy, timing, rejection_type=None,
|
||
rejection_key=None, allowed_responses=(), background_start=None,
|
||
):
|
||
handler = server.handlers[-1]
|
||
raw = await asyncio.wait_for(bad.reader.read(), 1)
|
||
responses = [json.loads(line) for line in raw.splitlines()]
|
||
assert responses == list(allowed_responses[:len(responses)])
|
||
assert session.state == "stop"
|
||
assert session.uid in server.sessions.sessions
|
||
if background_start is not None:
|
||
# token_answer queues handshake_ask before inspecting its payload.
|
||
# Closing this peer may make that already queued send fail normally.
|
||
tasks = timing.tasks[background_start:]
|
||
results = await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), 1)
|
||
for task, result in zip(tasks, results):
|
||
if isinstance(result, BaseException):
|
||
assert isinstance(result, ConnectionError)
|
||
timing.accept_task_error(task, result)
|
||
await healthy.telemetry(50, 90)
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
assert server.sessions[healthy.uid].last_data["telemetry"] == {"temperature": 50, "battery": 90}
|
||
cleanup = await timing.sleeping("manager")
|
||
assert cleanup.delay == 60
|
||
timing.clock.advance(60)
|
||
cleanup.release()
|
||
results = await asyncio.wait_for(asyncio.gather(
|
||
cleanup.task, handler, return_exceptions=True,
|
||
), 1)
|
||
errors = [result for result in results if isinstance(result, BaseException)]
|
||
errors.extend(server.peer_errors.get(handler, []))
|
||
for error in errors:
|
||
assert rejection_type is not None and isinstance(error, rejection_type)
|
||
if rejection_key is not None:
|
||
assert error.args == (rejection_key,)
|
||
server.accept_peer_rejection(handler, error)
|
||
assert session.uid not in server.sessions.sessions
|
||
assert session.uid not in server.router.session_routes
|
||
assert all(session.uid not in routes for routes in server.router.session_routes.values())
|
||
server.assert_no_unexpected_errors()
|
||
|
||
|
||
@pytest.mark.parametrize("raw,repetitions,rejection_type", [
|
||
(b"{broken\n", 1, json.JSONDecodeError),
|
||
(b"null\n", 1, AttributeError),
|
||
(b"[]\n", 1, AttributeError),
|
||
(b"{}\n", 1, None),
|
||
(b'{"command":[],"payload":{}}\n', 5, None),
|
||
], ids=["broken-json", "null", "array", "missing-command", "command-error-limit"])
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_invalid_packet_disconnects_only_its_peer(client_module, timing, raw, repetitions, rejection_type):
|
||
"""Проверяет изоляцию и очистку плохого пакета; неверная команда достигает лимита ошибок."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
bad, session = await open_at_state(server)
|
||
bad.writer.write(raw * repetitions)
|
||
await bad.writer.drain()
|
||
await assert_bad_peer_isolated(server, bad, session, healthy, timing, rejection_type)
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.parametrize("state,command", [("new", "connection"), ("token", "token_answer"),
|
||
("ready", "handshake_answer")])
|
||
@pytest.mark.parametrize("payload_case", ["missing", "null", "array", "string", "number", "boolean"])
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_invalid_control_payload_disconnects_only_its_peer(client_module, timing, state, command, payload_case):
|
||
"""Проверяет отключение и очистку, когда управляющий payload нельзя разобрать как объект."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
bad, session = await open_at_state(server, state)
|
||
packet = {"command": command}
|
||
if payload_case != "missing":
|
||
packet["payload"] = {"null": None, "array": [], "string": "bad", "number": 1, "boolean": True}[payload_case]
|
||
background_start = len(timing.tasks)
|
||
bad.writer.write((json.dumps(packet) + "\n").encode())
|
||
await bad.writer.drain()
|
||
await assert_bad_peer_isolated(
|
||
server, bad, session, healthy, timing, AttributeError,
|
||
allowed_responses=({"command": "handshake_ask"},) if state == "token" else (),
|
||
background_start=background_start,
|
||
)
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.parametrize("status", [
|
||
pytest.param(True, id="true"),
|
||
pytest.param(False, id="false"),
|
||
])
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_handshake_status_follows_protocol(client_module, timing, status):
|
||
"""Проверяет True/Other как успех и False как ожидание повторного handshake; Other — 0."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
client, session = await open_at_state(server, "ready")
|
||
await client.send("handshake_answer", {"status": status})
|
||
if status is False:
|
||
retry = await timing.sleeping("session", server.handlers[-1])
|
||
assert session.state == "ready"
|
||
assert session.uid not in server.router.session_routes
|
||
assert retry.delay > 0
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
retry.release()
|
||
assert await client.recv() == {"command": "handshake_ask"}
|
||
assert session.state == "ready"
|
||
await client.send("handshake_answer", {"status": True})
|
||
assert await client.heartbeat() == {"command": "heartbeat_answer"}
|
||
assert session.state == "air"
|
||
assert session.uid in server.router.session_routes
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.parametrize("token_case", ["missing", "null", "robot-token"])
|
||
def test_connection_preserves_robot_auth_token(client_module, token_case):
|
||
"""Проверяет хранение auth без whitelist-проверки и независимо от UID/upf восстановления."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
payload = {"name": "auth-client", "type": "test", "in": [], "out": []}
|
||
expected = "robot-auth-token" if token_case == "robot-token" else None
|
||
if token_case != "missing":
|
||
payload["token"] = expected
|
||
client, session = await open_at_state(server, "air", payload)
|
||
assert session.static_meta["auth"] == expected
|
||
assert client.uid == session.uid and client.upf == session.upf
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.parametrize("payload", [None, [], "unused", 1, True])
|
||
def test_heartbeat_does_not_require_an_object_payload(client_module, payload):
|
||
"""Проверяет ответ и продление heartbeat, который не использует содержимое payload."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
client = await server.client("heartbeat-payload")
|
||
healthy = await server.client("healthy")
|
||
session = server.sessions[client.uid]
|
||
deadline = session.timeout
|
||
await client.send("heartbeat_ask", payload)
|
||
assert await client.recv() == {"command": "heartbeat_answer"}
|
||
assert session.state == "air" and session.timeout >= deadline
|
||
assert session.last_data == {}
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_missing_heartbeat_payload_is_isolated_or_accepted(client_module, timing):
|
||
"""Проверяет heartbeat без payload: обработку либо локальное отключение с очисткой."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
client, session = await open_at_state(server, "air")
|
||
client.writer.write(b'{"command":"heartbeat_ask"}\n')
|
||
await client.writer.drain()
|
||
raw = await asyncio.wait_for(client.reader.readline(), 1)
|
||
if raw:
|
||
assert json.loads(raw) == {"command": "heartbeat_answer"}
|
||
assert session.state == "air"
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
else:
|
||
await assert_bad_peer_isolated(
|
||
server, client, session, healthy, timing, KeyError, "payload",
|
||
)
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.parametrize("payload", [None, [1, {"ok": True}], "telemetry", 42, False])
|
||
def test_data_payload_is_forwarded_without_dictionary_conversion(client_module, payload):
|
||
"""Проверяет сохранение и маршрутизацию JSON-значений payload без преобразования в словарь."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
sender = await server.client("sender")
|
||
receiver = await server.client("receiver")
|
||
server.router.session_routes[sender.uid].add(receiver.uid)
|
||
await sender.send("telemetry", payload)
|
||
await sender.heartbeat()
|
||
assert await receiver.recv() == {"command": "telemetry", "payload": payload}
|
||
assert server.sessions[sender.uid].last_data == {"telemetry": payload}
|
||
assert await receiver.heartbeat() == {"command": "heartbeat_answer"}
|
||
run(scenario)
|
||
|
||
|
||
def test_json_packet_split_between_tcp_writes(client_module):
|
||
"""Проверяет сборку одного JSON-пакета из нескольких записей до завершающего LF."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
client = await server.client("fragmented")
|
||
session = server.sessions[client.uid]
|
||
data = b'{"command":"telemetry","payload":{"battery":42}}\n'
|
||
for fragment in (data[:8], data[8:30], data[30:-1]):
|
||
client.writer.write(fragment)
|
||
await client.writer.drain()
|
||
# A response to a second client proves the loop had a chance to read fragments.
|
||
peer = await server.client("peer")
|
||
await peer.heartbeat()
|
||
assert "telemetry" not in session.last_data
|
||
client.writer.write(data[-1:])
|
||
await client.writer.drain()
|
||
await client.heartbeat()
|
||
assert session.last_data["telemetry"] == {"battery": 42}
|
||
run(scenario)
|
||
|
||
|
||
def test_several_json_packets_in_one_tcp_write(client_module):
|
||
"""Проверяет порядок нескольких JSON-строк, отправленных одной TCP-записью."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
client = await server.client("coalesced")
|
||
packets = [
|
||
{"command": "telemetry", "payload": {"battery": 43}},
|
||
{"command": "status", "payload": {"ok": True}},
|
||
{"command": "telemetry", "payload": {"battery": 42}},
|
||
{"command": "heartbeat_ask", "payload": {}},
|
||
]
|
||
client.writer.write(b"".join((json.dumps(p) + "\n").encode() for p in packets))
|
||
await client.writer.drain()
|
||
assert await client.recv() == {"command": "heartbeat_answer"}
|
||
assert server.sessions[client.uid].last_data == {"telemetry": {"battery": 42}, "status": {"ok": True}}
|
||
run(scenario)
|
||
|
||
|
||
@pytest.mark.usefixtures("make_manager")
|
||
def test_watcher_and_tcp_finally_schedule_cleanup_once(client_module, timing, monkeypatch):
|
||
"""Проверяет гонку watcher/EOF: одно закрытие транспорта и один таймер очистки."""
|
||
async def scenario():
|
||
async with timing, RunningServer(client_module) as server:
|
||
client = await server.client("racing")
|
||
session = server.sessions[client.uid]
|
||
calls, removed = [], []
|
||
close_requests = []
|
||
close_started, finish_close = asyncio.Event(), asyncio.Event()
|
||
original_close = session.connection.close
|
||
original_remove = server.router.remove
|
||
original_manager_close = server.sessions.close
|
||
|
||
async def counted_close():
|
||
calls.append(session.uid)
|
||
close_started.set()
|
||
await finish_close.wait()
|
||
await original_close()
|
||
|
||
async def observed_manager_close(uid):
|
||
close_requests.append(uid)
|
||
await original_manager_close(uid)
|
||
|
||
def counted_remove(uid):
|
||
removed.append(uid)
|
||
original_remove(uid)
|
||
|
||
monkeypatch.setattr(session.connection, "close", counted_close)
|
||
monkeypatch.setattr(server.router, "remove", counted_remove)
|
||
monkeypatch.setattr(server.sessions, "close", observed_manager_close)
|
||
poll = await timing.sleeping("session", session.watcher_task)
|
||
timing.clock.advance(session.timeout - timing.clock())
|
||
poll.release()
|
||
try:
|
||
await asyncio.wait_for(close_started.wait(), 1)
|
||
# Hold the watcher's close open while the real handler reaches finally.
|
||
client.writer.write_eof()
|
||
await asyncio.wait_for(server.handlers[0], 1)
|
||
assert close_requests == [session.uid, session.uid]
|
||
finally:
|
||
finish_close.set()
|
||
assert await asyncio.wait_for(client.reader.read(), 1) == b""
|
||
cleanup = await timing.sleeping("manager")
|
||
assert calls == [session.uid]
|
||
await asyncio.sleep(0)
|
||
assert timing.queues["manager"].empty()
|
||
timing.clock.advance(60)
|
||
cleanup.release()
|
||
await asyncio.wait_for(cleanup.task, 1)
|
||
await asyncio.wait_for(asyncio.gather(session.watcher_task, *server.handlers), 1)
|
||
assert removed == [session.uid]
|
||
assert session.uid not in server.sessions.sessions
|
||
run(scenario)
|
||
|
||
|
||
async def assert_send_error_observed(caplog, failure):
|
||
async with asyncio.timeout(1):
|
||
while not any(record.exc_info and record.exc_info[1] is failure for record in caplog.records):
|
||
await asyncio.sleep(0)
|
||
|
||
|
||
def test_background_send_failure_is_observable_and_isolated(client_module, monkeypatch, caplog):
|
||
"""Проверяет видимость фонового ConnectionResetError, закрытие по EOF и работу другого клиента."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
bad = await server.client("broken-send")
|
||
session = server.sessions[bad.uid]
|
||
failure = ConnectionResetError("send regression")
|
||
raised = asyncio.Event()
|
||
|
||
async def failing_send(packet):
|
||
raised.set()
|
||
raise failure
|
||
|
||
monkeypatch.setattr(session.connection, "send", failing_send)
|
||
session.send({"command": "telemetry", "payload": {"battery": 87}})
|
||
await asyncio.wait_for(raised.wait(), 1)
|
||
await assert_send_error_observed(caplog, failure)
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
# A send exception in its own task need not close the reader task.
|
||
# The actual TCP EOF must still close the session normally.
|
||
await server.disconnect(bad)
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
run(scenario)
|
||
|
||
|
||
def test_disconnect_while_background_send_is_waiting(client_module, monkeypatch, caplog):
|
||
"""Проверяет разрыв во время фоновой отправки и наблюдаемость ошибки без повторного close."""
|
||
async def scenario():
|
||
async with RunningServer(client_module) as server:
|
||
healthy = await server.client("healthy")
|
||
client = await server.client("pending-send")
|
||
session = server.sessions[client.uid]
|
||
started, release, finished = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||
original_send = session.connection.send
|
||
original_close = session.connection.close
|
||
closes = []
|
||
failures = []
|
||
|
||
async def pending_send(packet):
|
||
started.set()
|
||
try:
|
||
await release.wait()
|
||
await original_send(packet)
|
||
except ConnectionError as error:
|
||
failures.append(error)
|
||
raise
|
||
finally:
|
||
finished.set()
|
||
|
||
async def counted_close():
|
||
closes.append(session.uid)
|
||
await original_close()
|
||
|
||
monkeypatch.setattr(session.connection, "send", pending_send)
|
||
monkeypatch.setattr(session.connection, "close", counted_close)
|
||
session.send({"command": "telemetry", "payload": {}})
|
||
try:
|
||
await asyncio.wait_for(started.wait(), 1)
|
||
await server.disconnect(client)
|
||
finally:
|
||
release.set()
|
||
await asyncio.wait_for(finished.wait(), 1)
|
||
await asyncio.sleep(0)
|
||
assert closes == [session.uid]
|
||
assert await healthy.heartbeat() == {"command": "heartbeat_answer"}
|
||
for failure in failures:
|
||
await assert_send_error_observed(caplog, failure)
|
||
run(scenario)
|