diff --git a/cloudrift/messaging/sqs.py b/cloudrift/messaging/sqs.py index 0bec142..4c62a1c 100644 --- a/cloudrift/messaging/sqs.py +++ b/cloudrift/messaging/sqs.py @@ -134,6 +134,7 @@ async def _ensure(self): async def close(self) -> None: client_cm, self._client_cm = self._client_cm, None self._client = None + self._pending.clear() if client_cm is not None: await client_cm.__aexit__(None, None, None) @@ -237,6 +238,7 @@ async def receive( for m in response.get("Messages", []): attrs = m.get("Attributes", {}) receive_count = attrs.get("ApproximateReceiveCount") + self._pending[m["ReceiptHandle"]] = m["Body"] messages.append( Message( id=m["MessageId"], @@ -263,6 +265,9 @@ async def nack(self, receipt_handle: str) -> None: ) except ClientError as e: self._raise(e) + finally: + # the handle goes stale on redelivery; redelivery stores a new one + self._pending.pop(receipt_handle, None) async def delete(self, receipt_handle: str) -> None: client = await self._ensure() @@ -344,6 +349,7 @@ async def purge(self) -> None: client = await self._ensure() try: await client.purge_queue(QueueUrl=self.queue_url) + self._pending.clear() except ClientError as e: self._raise(e) diff --git a/tests/test_messaging.py b/tests/test_messaging.py index 64e1df5..3a5cc40 100644 --- a/tests/test_messaging.py +++ b/tests/test_messaging.py @@ -9,6 +9,7 @@ REGION = "us-east-1" QUEUE_NAME = "test-queue" +DLQ_NAME = "test-queue-dlq" FIFO_QUEUE_NAME = "test-queue.fifo" @@ -116,6 +117,110 @@ async def test_group_id_on_standard_queue_raises(sqs_backend): await sqs_backend.send({"x": 1}, group_id="g1") +# --- dead_letter / get_queue_depth --- + + +async def test_dead_letter_moves_message_to_dlq(sqs_backend, sqs_client): + await sqs_backend.send({"poison": True, "id": 7}) + [m] = await sqs_backend.receive(max_messages=1) + + await sqs_backend.dead_letter(m.receipt_handle, reason="schema mismatch") + + # gone from the source queue (and from the pending map) + assert await sqs_backend.receive(max_messages=10) == [] + assert m.receipt_handle not in sqs_backend._pending + + # present on the DLQ with the original body and the reason attribute + resp = sqs_client.receive_message( + QueueUrl=sqs_backend._dlq_test_url, + MaxNumberOfMessages=1, + MessageAttributeNames=["All"], + ) + dlq_messages = resp.get("Messages", []) + assert len(dlq_messages) == 1 + assert json.loads(dlq_messages[0]["Body"]) == {"poison": True, "id": 7} + reason_attr = dlq_messages[0]["MessageAttributes"]["DeadLetterReason"] + assert reason_attr["StringValue"] == "schema mismatch" + sqs_client.delete_message( + QueueUrl=sqs_backend._dlq_test_url, + ReceiptHandle=dlq_messages[0]["ReceiptHandle"], + ) + + +async def test_dead_letter_unknown_handle_raises(sqs_backend): + from cloudrift.core.exceptions import MessagingError + + with pytest.raises(MessagingError, match="No pending message"): + await sqs_backend.dead_letter("bogus-handle", reason="x") + + +async def test_dead_letter_without_dlq_raises(moto_server, sqs_client): + """A queue with no RedrivePolicy and no explicit dlq_url cannot dead-letter.""" + from cloudrift.core.exceptions import MessagingError + + queue_url = sqs_client.create_queue(QueueName="no-dlq-queue")["QueueUrl"] + backend = get_queue( + "sqs", + queue_url=queue_url, + aws_access_key_id="test", + aws_secret_access_key="test", + region=REGION, + endpoint_url=moto_server, + ) + try: + await backend.send({"x": 1}) + [m] = await backend.receive(max_messages=1) + with pytest.raises(MessagingError, match="No dead-letter queue configured"): + await backend.dead_letter(m.receipt_handle, reason="x") + finally: + await backend.close() + sqs_client.delete_queue(QueueUrl=queue_url) + + +async def test_dead_letter_with_explicit_dlq_url(moto_server, sqs_client): + """dlq_url= passed at construction wins over RedrivePolicy resolution.""" + dlq_url = sqs_client.create_queue(QueueName="explicit-dlq")["QueueUrl"] + queue_url = sqs_client.create_queue(QueueName="explicit-src")["QueueUrl"] + backend = get_queue( + "sqs", + queue_url=queue_url, + dlq_url=dlq_url, + aws_access_key_id="test", + aws_secret_access_key="test", + region=REGION, + endpoint_url=moto_server, + ) + try: + await backend.send({"n": 1}) + [m] = await backend.receive(max_messages=1) + await backend.dead_letter(m.receipt_handle, reason="explicit") + resp = sqs_client.receive_message(QueueUrl=dlq_url, MaxNumberOfMessages=1) + assert json.loads(resp["Messages"][0]["Body"]) == {"n": 1} + finally: + await backend.close() + sqs_client.delete_queue(QueueUrl=queue_url) + sqs_client.delete_queue(QueueUrl=dlq_url) + + +async def test_nack_drops_pending_entry(sqs_backend): + await sqs_backend.send({"retry": 1}) + [m] = await sqs_backend.receive(max_messages=1) + assert m.receipt_handle in sqs_backend._pending + await sqs_backend.nack(m.receipt_handle) + assert m.receipt_handle not in sqs_backend._pending + # redelivery stores a fresh handle, and dead_letter works on it + [again] = await sqs_backend.receive(max_messages=1) + await sqs_backend.dead_letter(again.receipt_handle, reason="after nack") + + +async def test_get_queue_depth(sqs_backend): + assert await sqs_backend.get_queue_depth() == 0 + await sqs_backend.send({"a": 1}) + await sqs_backend.send({"b": 2}) + assert await sqs_backend.get_queue_depth() == 2 + await sqs_backend.purge() + + # --- FIFO tests --- diff --git a/tests/test_messaging_azure.py b/tests/test_messaging_azure.py index 4be5a2c..fe7e4d3 100644 --- a/tests/test_messaging_azure.py +++ b/tests/test_messaging_azure.py @@ -198,6 +198,48 @@ async def test_nack_unknown_handle_raises(): await backend.nack("missing") +async def test_dead_letter_calls_dead_letter_message_and_releases(): + backend = _make_backend(session_enabled=True) + client = MagicMock() + receiver = AsyncMock() + raw = _make_received_message(lock_token="tok-dl") + receiver.receive_messages.return_value = [raw] + client.get_queue_receiver.return_value = receiver + _patch_client(backend, client) + + [m] = await backend.receive() + await backend.dead_letter(m.receipt_handle, reason="schema mismatch") + + receiver.dead_letter_message.assert_awaited_once_with( + raw, reason="schema mismatch", error_description="schema mismatch" + ) + assert backend._pending == {} + assert backend._receiver_tokens == {} + receiver.__aexit__.assert_awaited() + + +async def test_dead_letter_unknown_handle_raises(): + backend = _make_backend() + with pytest.raises(MessagingError, match="No pending message"): + await backend.dead_letter("missing", reason="x") + + +async def test_get_queue_depth_uses_admin_client(): + backend = _make_backend() + props = MagicMock() + props.active_message_count = 5 + admin = AsyncMock() + admin.get_queue_runtime_properties.return_value = props + admin.__aenter__.return_value = admin + with patch( + "azure.servicebus.aio.management.ServiceBusAdministrationClient" + ) as admin_cls: + admin_cls.from_connection_string.return_value = admin + depth = await backend.get_queue_depth() + assert depth == 5 + admin.get_queue_runtime_properties.assert_awaited_once_with("test-queue") + + async def test_session_enabled_threads_through_factories(): with patch("azure.identity.aio.ManagedIdentityCredential"): b = AzureServiceBusBackend.from_managed_identity(