Fix flaky integration tests

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:

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(

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

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())