feat(sqs): improve validation for queue_url

Author: gmcrocetti

aws_lambda_powertools/utilities/batch/sqs.py

@@ -15,17 +15,26 @@
 
 class PartialSQSProcessor(BasePartialProcessor):
     def __init__(self):
-        self._client = boto3.client("sqs")
+        self.client = boto3.client("sqs")
         self.success_messages: List = []
         self.fail_messages: List = []
 
         super().__init__()
 
     def get_queue_url(self):
+        """
+        Format QueueUrl from first records entry
+        """
+        if not getattr(self, "records", None):
+            return
+
         *_, account_id, queue_name = self.records[0]["eventSourceARN"].split(":")
-        return f"{self._client._endpoint.host}/{account_id}/{queue_name}"
+        return f"{self.client._endpoint.host}/{account_id}/{queue_name}"
 
     def get_entries_to_clean(self):
+        """
+        Format messages to use in batch deletion
+        """
         return [{"Id": msg["messageId"], "ReceiptHandle": msg["receiptHandle"]} for msg in self.success_messages]
 
     def _process_record(self, record):
@@ -37,21 +46,22 @@ def _process_record(self, record):
 
     def _prepare(self):
         """
+        Remove results from previous executions.
         """
         self.success_messages.clear()
         self.fail_messages.clear()
 
     def _clean(self):
         """
+        Delete messages from Queue in case of partial failure.
         """
-        # skip only failures or only successes
         if not (self.fail_messages and self.success_messages):
             return
 
         queue_url = self.get_queue_url()
         entries_to_remove = self.get_entries_to_clean()
 
-        return self._client.delete_message_batch(QueueUrl=queue_url, Entries=entries_to_remove)
+        return self.client.delete_message_batch(QueueUrl=queue_url, Entries=entries_to_remove)
 
 
 @lambda_handler_decorator

tests/functional/test_utilities_batch.py

@@ -0,0 +1,148 @@
+from typing import Callable
+import pytest
+
+from botocore.stub import Stubber
+
+from aws_lambda_powertools.utilities.batch import PartialSQSProcessor, partial_sqs_processor
+
+
+@pytest.fixture
+def sqs_event_factory() -> Callable:
+    def factory(body: str):
+        return {
+            "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
+            "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a",
+            "body": body,
+            "attributes": {},
+            "messageAttributes": {},
+            "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+            "eventSource": "aws:sqs",
+            "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+            "awsRegion": "us-east-1",
+        }
+
+    return factory
+
+
+@pytest.fixture
+def record_handler() -> Callable:
+    def handler(record):
+        body = record["body"]
+        if "fail" in body:
+            raise Exception("Failed to process record.")
+        return body
+
+    return handler
+
+
+def test_partial_sqs_processor_context_with_failure(sqs_event_factory, record_handler):
+    """ Test processor with one record failing """
+    processor = PartialSQSProcessor()
+
+    fail_record = sqs_event_factory("fail")
+    success_record = sqs_event_factory("success")
+
+    records = [fail_record, success_record]
+
+    response = {"Successful": [{"Id": fail_record["messageId"]},], "Failed": []}
+
+    with Stubber(processor.client) as stubber:
+        stubber.add_response("delete_message_batch", response)
+
+        with processor(records, record_handler) as ctx:
+            result = ctx.process()
+
+        stubber.assert_no_pending_responses()
+
+    assert result == [
+        ("fail", ("Failed to process record.",), fail_record),
+        ("success", success_record["body"], success_record),
+    ]
+
+
+def test_partial_sqs_processor_context_only_success(sqs_event_factory, record_handler):
+    """ Test processor without failure """
+    processor = PartialSQSProcessor()
+
+    first_record = sqs_event_factory("success")
+    second_record = sqs_event_factory("success")
+
+    records = [first_record, second_record]
+
+    with processor(records, record_handler) as ctx:
+        result = ctx.process()
+
+    assert result == [
+        ("success", first_record["body"], first_record),
+        ("success", second_record["body"], second_record),
+    ]
+
+
+def test_partial_sqs_processor_context_multiple_calls(sqs_event_factory, record_handler):
+    """ Test processor without failure """
+    processor = PartialSQSProcessor()
+
+    first_record = sqs_event_factory("success")
+    second_record = sqs_event_factory("success")
+
+    records = [first_record, second_record]
+
+    with processor(records, record_handler) as ctx:
+        ctx.process()
+
+    with processor([first_record], record_handler) as ctx:
+        ctx.process()
+
+    assert processor.success_messages == [first_record]
+
+
+def test_partial_sqs_processor_middleware_with_default(sqs_event_factory, record_handler):
+    """ Test middleware with default partial processor """
+    processor = PartialSQSProcessor()
+
+    @partial_sqs_processor(record_handler=record_handler, processor=processor)
+    def lambda_handler(event, context):
+        return True
+
+    fail_record = sqs_event_factory("fail")
+
+    event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("success")]}
+    response = {"Successful": [{"Id": fail_record["messageId"]},], "Failed": []}
+
+    with Stubber(processor.client) as stubber:
+        stubber.add_response("delete_message_batch", response)
+
+        result = lambda_handler(event, {})
+
+        stubber.assert_no_pending_responses()
+
+    assert result is True
+
+
+def test_partial_sqs_processor_middleware_with_custom(capsys, sqs_event_factory, record_handler):
+    """ Test middle with custom partial processor """
+    class CustomProcessor(PartialSQSProcessor):
+        def failure_handler(self, record, exception):
+            print("Oh no ! It's a failure.")
+            return super().failure_handler(record, exception)
+
+    processor = CustomProcessor()
+
+    @partial_sqs_processor(record_handler=record_handler, processor=processor)
+    def lambda_handler(event, context):
+        return True
+
+    fail_record = sqs_event_factory("fail")
+
+    event = {"Records": [sqs_event_factory("fail"), sqs_event_factory("success")]}
+    response = {"Successful": [{"Id": fail_record["messageId"]},], "Failed": []}
+
+    with Stubber(processor.client) as stubber:
+        stubber.add_response("delete_message_batch", response)
+
+        result = lambda_handler(event, {})
+
+        stubber.assert_no_pending_responses()
+
+    assert result is True
+    assert capsys.readouterr().out == "Oh no ! It's a failure.\n"

tests/unit/test_utilities_batch.py

@@ -0,0 +1,126 @@
+import pytest
+
+from aws_lambda_powertools.utilities.batch import PartialSQSProcessor
+
+
+@pytest.fixture(scope="module")
+def sqs_event():
+    return {
+        "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
+        "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a",
+        "body": "",
+        "attributes": {},
+        "messageAttributes": {},
+        "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+        "eventSource": "aws:sqs",
+        "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+        "awsRegion": "us-east-1",
+    }
+
+
+def test_partial_sqs_get_queue_url_with_records(mocker, sqs_event):
+    expected_url = "https://queue.amazonaws.com/123456789012/my-queue"
+
+    records_mock = mocker.patch.object(PartialSQSProcessor, "records", create=True, new_callable=mocker.PropertyMock)
+    records_mock.return_value = [sqs_event]
+
+    result = PartialSQSProcessor().get_queue_url()
+    assert result == expected_url
+
+
+def test_partial_sqs_get_queue_url_without_records():
+    assert PartialSQSProcessor().get_queue_url() == None
+
+
+def test_partial_sqs_get_entries_to_clean_with_success(mocker, sqs_event):
+    expected_entries = [{"Id": sqs_event["messageId"], "ReceiptHandle": sqs_event["receiptHandle"]}]
+
+    success_messages_mock = mocker.patch.object(
+        PartialSQSProcessor, "success_messages", new_callable=mocker.PropertyMock
+    )
+    success_messages_mock.return_value = [sqs_event]
+
+    result = PartialSQSProcessor().get_entries_to_clean()
+
+    assert result == expected_entries
+
+
+def test_partial_sqs_get_entries_to_clean_without_success(mocker):
+    expected_entries = []
+
+    success_messages_mock = mocker.patch.object(
+        PartialSQSProcessor, "success_messages", new_callable=mocker.PropertyMock
+    )
+    success_messages_mock.return_value = []
+
+    result = PartialSQSProcessor().get_entries_to_clean()
+
+    assert result == expected_entries
+
+
+def test_partial_sqs_process_record_success(mocker):
+    expected_value = mocker.sentinel.expected_value
+
+    success_result = mocker.sentinel.success_result
+    record = mocker.sentinel.record
+
+    handler_mock = mocker.patch.object(PartialSQSProcessor, "handler", create=True, return_value=success_result)
+    success_handler_mock = mocker.patch.object(PartialSQSProcessor, "success_handler", return_value=expected_value)
+
+    result = PartialSQSProcessor()._process_record(record)
+
+    handler_mock.assert_called_once_with(record)
+    success_handler_mock.assert_called_once_with(record, success_result)
+
+    assert result == expected_value
+
+
+def test_partial_sqs_process_record_failure(mocker):
+    expected_value = mocker.sentinel.expected_value
+
+    failure_result = Exception()
+    record = mocker.sentinel.record
+
+    handler_mock = mocker.patch.object(PartialSQSProcessor, "handler", create=True, side_effect=failure_result)
+    failure_handler_mock = mocker.patch.object(PartialSQSProcessor, "failure_handler", return_value=expected_value)
+
+    result = PartialSQSProcessor()._process_record(record)
+
+    handler_mock.assert_called_once_with(record)
+    failure_handler_mock.assert_called_once_with(record, failure_result)
+
+    assert result == expected_value
+
+
+def test_partial_sqs_prepare(mocker):
+    processor = PartialSQSProcessor()
+
+    success_messages_mock = mocker.patch.object(processor, "success_messages", spec=list)
+    failed_messages_mock = mocker.patch.object(processor, "fail_messages", spec=list)
+
+    processor._prepare()
+
+    success_messages_mock.clear.assert_called_once()
+    failed_messages_mock.clear.assert_called_once()
+
+
+def test_partial_sqs_clean(monkeypatch, mocker):
+    processor = PartialSQSProcessor()
+    records = [mocker.sentinel.record]
+
+    monkeypatch.setattr(processor, "fail_messages", records)
+    monkeypatch.setattr(processor, "success_messages", records)
+
+    queue_url_mock = mocker.patch.object(PartialSQSProcessor, "get_queue_url")
+    entries_to_clean_mock = mocker.patch.object(PartialSQSProcessor, "get_entries_to_clean")
+
+    queue_url_mock.return_value = mocker.sentinel.queue_url
+    entries_to_clean_mock.return_value = mocker.sentinel.entries_to_clean
+
+    client_mock = mocker.patch.object(processor, "client", autospec=True)
+
+    processor._clean()
+
+    client_mock.delete_message_batch.assert_called_once_with(
+        QueueUrl=mocker.sentinel.queue_url, Entries=mocker.sentinel.entries_to_clean
+    )