diff --git a/Justfile b/Justfile index 4f82e5b..1e9100c 100644 --- a/Justfile +++ b/Justfile @@ -18,8 +18,10 @@ test-fast *args: test *args: #!/bin/bash + set -euo pipefail trap 'echo; docker compose logs && docker compose down --remove-orphans' EXIT docker compose up -d + python3 scripts/wait_for_stomp_brokers.py uv run pytest {{args}} run-artemis: diff --git a/packages/stompman/test_stompman/test_integration.py b/packages/stompman/test_stompman/test_integration.py index ef1978e..1558bff 100644 --- a/packages/stompman/test_stompman/test_integration.py +++ b/packages/stompman/test_stompman/test_integration.py @@ -1,9 +1,7 @@ import asyncio from collections.abc import AsyncGenerator, Callable from contextlib import asynccontextmanager -from datetime import timedelta from itertools import starmap -from typing import Final from uuid import uuid4 import pytest @@ -20,7 +18,24 @@ parse_header, ) -DESTINATION: Final = "DLQ" + +async def wait_for_reconnect(client: stompman.Client, initial_reconnection_count: int) -> None: + def is_reconnected() -> bool: + return ( + client._connection_manager._reconnection_count > initial_reconnection_count + and client._connection_manager._active_connection_state is not None + ) + + while not is_reconnected(): # noqa: ASYNC110 + await asyncio.sleep(0.05) + + +async def force_reconnect(client: stompman.Client) -> None: + connection_state = await client._connection_manager._get_active_connection_state() + initial_reconnection_count = client._connection_manager._reconnection_count + client._connection_manager._clear_active_connection_state(stompman.ConnectionLostError(reason="test reconnect")) + await connection_state.connection.close() + await asyncio.wait_for(wait_for_reconnect(client, initial_reconnection_count), timeout=5) @asynccontextmanager @@ -37,35 +52,28 @@ async def test_consumption_survives_forced_reconnects( received: list[bytes] = [] received_event = asyncio.Event() - async def handle_message(frame: stompman.MessageFrame) -> None: # noqa: RUF029 + async def handle_message(frame: stompman.AckableMessageFrame) -> None: received.append(frame.body) + await frame.ack() received_event.set() + destination = "DLQ" + async with ( - stompman.Client( - servers=[connection_parameters], - connection_confirmation_timeout=10, - no_message_restart_interval=timedelta(milliseconds=300), - ) as consumer, + stompman.Client(servers=[connection_parameters], connection_confirmation_timeout=10) as consumer, stompman.Client(servers=[connection_parameters], connection_confirmation_timeout=10) as producer, ): async def consume_after_reconnects() -> None: for index in range(iterations): - initial_reconnection_count = consumer._connection_manager._reconnection_count - await asyncio.sleep(0.6) - assert consumer._connection_manager._reconnection_count > initial_reconnection_count, ( - f"iteration {index}: expected forced reconnect to fire" - ) + await force_reconnect(consumer) payload = f"msg-{index}-{uuid4()}".encode() received_event.clear() - await producer.send(body=payload, destination=DESTINATION) + await producer.send(body=payload, destination=destination) await asyncio.wait_for(received_event.wait(), timeout=5) assert payload in received, f"iteration {index}: {payload!r} not delivered" - subscription = await consumer.subscribe( - destination=DESTINATION, handler=handle_message, on_suppressed_exception=print - ) + subscription = await consumer.subscribe_with_manual_ack(destination=destination, handler=handle_message) try: await consume_after_reconnects() finally: @@ -79,16 +87,18 @@ async def consume_after_reconnects() -> None: @pytest.mark.anyio async def test_ok(connection_parameters: stompman.ConnectionParameters) -> None: - async def produce() -> None: + async def produce(destination: str) -> None: + await subscribed_event.wait() + for message in messages[200:]: - await producer.send(body=message, destination=DESTINATION, headers={"hello": "from outside transaction"}) + await producer.send(body=message, destination=destination, headers={"hello": "from outside transaction"}) async with producer.begin() as transaction: for message in messages[:200]: - await transaction.send(body=message, destination=DESTINATION, headers={"hello": "from transaction"}) + await transaction.send(body=message, destination=destination, headers={"hello": "from transaction"}) - async def consume() -> None: - received_messages = [] + async def consume(destination: str) -> None: + received_messages: list[bytes] = [] event = asyncio.Event() async def handle_message(frame: stompman.MessageFrame) -> None: # noqa: RUF029 @@ -97,22 +107,25 @@ async def handle_message(frame: stompman.MessageFrame) -> None: # noqa: RUF029 event.set() subscription = await consumer.subscribe( - destination=DESTINATION, handler=handle_message, on_suppressed_exception=print + destination=destination, handler=handle_message, on_suppressed_exception=print ) - await asyncio.wait_for(event.wait(), timeout=5) + subscribed_event.set() + await asyncio.wait_for(event.wait(), timeout=15) await subscription.unsubscribe() assert sorted(received_messages) == sorted(messages) messages = [str(uuid4()).encode() for _ in range(1000)] + destination = "DLQ" + subscribed_event = asyncio.Event() async with ( create_client(connection_parameters) as consumer, create_client(connection_parameters) as producer, asyncio.TaskGroup() as task_group, ): - task_group.create_task(consume()) - task_group.create_task(produce()) + task_group.create_task(consume(destination)) + task_group.create_task(produce(destination)) def generate_frames( diff --git a/pyproject.toml b/pyproject.toml index 3359107..e205134 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,6 @@ dev = [ "pytest==9.0.3", "pytest-cov==7.1.0", "pytest-timeout==2.4.0", - "pytest-rerunfailures==16.3", "ruff==0.15.15", "uvloop==0.22.1", ] @@ -62,7 +61,7 @@ ignore = [ extend-per-file-ignores = { "*/test_*/*" = ["S101", "SLF001", "ARG", "PLR6301"] } [tool.pytest.ini_options] -addopts = "--cov -s -vv --reruns 6 --only-rerun FailedAllConnectAttemptsError" +addopts = "--cov -s -vv" [tool.coverage.report] skip_covered = true diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/wait_for_stomp_brokers.py b/scripts/wait_for_stomp_brokers.py new file mode 100644 index 0000000..527654b --- /dev/null +++ b/scripts/wait_for_stomp_brokers.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import socket +import sys +import time +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class Broker: + name: str + host: str + port: int + login: str + passcode: str + + +class BrokerProtocolError(Exception): + pass + + +class BrokerWaitTimeoutError(Exception): + pass + + +BROKERS = (Broker(name="ActiveMQ Artemis", host="127.0.0.1", port=9000, login="admin", passcode=":=123"),) +CONNECT_TIMEOUT_SECONDS = 2.0 +READ_TIMEOUT_SECONDS = 2.0 +OVERALL_TIMEOUT_SECONDS = 90.0 +RETRY_INTERVAL_SECONDS = 1.0 + + +def _build_connect_frame(broker: Broker) -> bytes: + return ( + "CONNECT\n" + "accept-version:1.2\n" + f"host:{broker.host}\n" + f"login:{broker.login}\n" + f"passcode:{broker.passcode}\n" + "heart-beat:0,0\n" + "\n" + "\0" + ).encode() + + +def _read_stomp_frame(sock: socket.socket) -> bytes: + response = b"" + while b"\0" not in response: + chunk = sock.recv(4096) + if not chunk: + break + response += chunk + return response + + +def _probe_broker(broker: Broker) -> None: + with socket.create_connection((broker.host, broker.port), timeout=CONNECT_TIMEOUT_SECONDS) as sock: + sock.settimeout(READ_TIMEOUT_SECONDS) + sock.sendall(_build_connect_frame(broker)) + response = _read_stomp_frame(sock) + if response.lstrip(b"\n").startswith(b"CONNECTED\n"): + sock.sendall(b"DISCONNECT\n\n\0") + return + + decoded_response = response.decode(errors="replace") + message = f"{broker.name} did not accept a STOMP connection: {decoded_response!r}" + raise BrokerProtocolError(message) + + +def _wait_for_broker(broker: Broker, deadline: float) -> None: + last_error: Exception | None = None + while time.monotonic() < deadline: + try: + _probe_broker(broker) + except (BrokerProtocolError, OSError) as error: + last_error = error + else: + sys.stdout.write(f"{broker.name} is ready on {broker.host}:{broker.port}\n") + sys.stdout.flush() + return + time.sleep(RETRY_INTERVAL_SECONDS) + + message = f"Timed out waiting for {broker.name} on {broker.host}:{broker.port}: {last_error}" + raise BrokerWaitTimeoutError(message) + + +def main() -> int: + deadline = time.monotonic() + OVERALL_TIMEOUT_SECONDS + for broker in BROKERS: + _wait_for_broker(broker, deadline) + return 0 + + +if __name__ == "__main__": + sys.exit(main())