152 lines
4.5 KiB
Python
152 lines
4.5 KiB
Python
"""Unit tests for the RabbitMQ consumer base hook crash guard (ticket 007)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from aio_pika import Message
|
|
|
|
from contract_check.core.db.enums import FailureClass
|
|
from contract_check.core.mq.consumer import Consumer
|
|
from contract_check.core.mq.messages import PipelineMessage
|
|
from contract_check.core.mq.topology import H_ATTEMPT
|
|
|
|
pytestmark = pytest.mark.unit
|
|
|
|
|
|
class _FakeChannel:
|
|
def __init__(self) -> None:
|
|
self.published: list[tuple[Message, str]] = []
|
|
self.acks: list[Message] = []
|
|
|
|
async def get_exchange(self, name: str, ensure: bool = True) -> _FakeExchange:
|
|
return _FakeExchange(self)
|
|
|
|
@property
|
|
def default_exchange(self) -> _FakeExchange:
|
|
return _FakeExchange(self)
|
|
|
|
|
|
class _FakeExchange:
|
|
def __init__(self, channel: _FakeChannel) -> None:
|
|
self._channel = channel
|
|
|
|
async def publish(self, message: Message, routing_key: str) -> None:
|
|
self._channel.published.append((message, routing_key))
|
|
|
|
|
|
class _FakeIncomingMessage:
|
|
def __init__(self, body: bytes, headers: dict[str, Any] | None = None) -> None:
|
|
self.body = body
|
|
self.headers = headers or {}
|
|
self.correlation_id = str(uuid.uuid4())
|
|
self.content_type = "application/json"
|
|
self._ack = False
|
|
|
|
async def ack(self) -> None:
|
|
self._ack = True
|
|
|
|
|
|
class _TestMessage(PipelineMessage):
|
|
pass
|
|
|
|
|
|
class _CrashHookConsumer(Consumer[_TestMessage]):
|
|
queue = "test.q"
|
|
routing_key = "test"
|
|
message_model = _TestMessage
|
|
retry_exchange = "retry.x"
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
"amqp://test",
|
|
origin="test",
|
|
prefetch=1,
|
|
max_attempts=3,
|
|
retry_base_ms=1000,
|
|
)
|
|
self.channel = _FakeChannel()
|
|
self._channel = self.channel
|
|
self.failure_calls: list[tuple[Any, ...]] = []
|
|
self.dlq_calls: list[tuple[Any, ...]] = []
|
|
self.raise_on_failure = True
|
|
self.raise_on_dlq = True
|
|
|
|
async def handle(self, payload: _TestMessage) -> None:
|
|
raise RuntimeError("handler failed")
|
|
|
|
def classify(self, exc: BaseException) -> FailureClass:
|
|
return "infra"
|
|
|
|
async def on_failure(
|
|
self,
|
|
payload: _TestMessage,
|
|
failure_class: FailureClass,
|
|
attempt: int,
|
|
error: str,
|
|
) -> None:
|
|
self.failure_calls.append((payload, failure_class, attempt, error))
|
|
if self.raise_on_failure:
|
|
raise RuntimeError("on_failure crashed")
|
|
|
|
async def on_dlq(self, payload: _TestMessage, failure_class: FailureClass, error: str) -> None:
|
|
self.dlq_calls.append((payload, failure_class, error))
|
|
if self.raise_on_dlq:
|
|
raise RuntimeError("on_dlq crashed")
|
|
|
|
|
|
async def test_on_failure_crash_republishes_to_retry_with_unchanged_attempt() -> None:
|
|
consumer = _CrashHookConsumer()
|
|
consumer.raise_on_failure = True
|
|
consumer.raise_on_dlq = False
|
|
payload = _TestMessage(
|
|
correlation_id=uuid.uuid4(),
|
|
document_id=uuid.uuid4(),
|
|
user_id=uuid.uuid4(),
|
|
attempt=2,
|
|
)
|
|
msg = _FakeIncomingMessage(
|
|
payload.model_dump_json().encode("utf-8"),
|
|
headers={H_ATTEMPT: 2},
|
|
)
|
|
|
|
await consumer._on_message(msg)
|
|
|
|
assert len(consumer.channel.published) == 1
|
|
published, routing_key = consumer.channel.published[0]
|
|
assert routing_key == "retry.test"
|
|
assert published.headers["x-attempt"] == 2 # unchanged
|
|
assert msg._ack is True
|
|
|
|
|
|
async def test_on_dlq_crash_republishes_to_retry_with_unchanged_attempt() -> None:
|
|
from contract_check.core.errors import TerminalError
|
|
|
|
class _TerminalConsumer(_CrashHookConsumer):
|
|
async def handle(self, payload: _TestMessage) -> None:
|
|
raise TerminalError("terminal")
|
|
|
|
consumer = _TerminalConsumer()
|
|
consumer.raise_on_failure = False
|
|
consumer.raise_on_dlq = True
|
|
payload = _TestMessage(
|
|
correlation_id=uuid.uuid4(),
|
|
document_id=uuid.uuid4(),
|
|
user_id=uuid.uuid4(),
|
|
attempt=1,
|
|
)
|
|
msg = _FakeIncomingMessage(
|
|
payload.model_dump_json().encode("utf-8"),
|
|
headers={H_ATTEMPT: 1},
|
|
)
|
|
|
|
await consumer._on_message(msg)
|
|
|
|
assert len(consumer.channel.published) == 1
|
|
published, routing_key = consumer.channel.published[0]
|
|
assert routing_key == "retry.test"
|
|
assert published.headers["x-attempt"] == 1 # unchanged
|
|
assert msg._ack is True
|
|
assert len(consumer.dlq_calls) == 1
|