ferruzzi commented on code in PR #56371: URL: https://github.com/apache/airflow/pull/56371#discussion_r2430914490
########## providers/amazon/tests/unit/amazon/aws/hooks/test_sqs.py: ########## @@ -24,18 +24,130 @@ from airflow.providers.amazon.aws.hooks.sqs import SqsHook +QUEUE_NAME = "test-queue" QUEUE_URL = "https://sqs.region.amazonaws.com/123456789/test-queue"; MESSAGE_BODY = "test message" +DELAY = 5 +DEDUPE = "banana" +MSG_ATTRIBUTES = { + "Author": { + "StringValue": "test-user", + "DataType": "String", + }, + "Priority": { + "StringValue": "1", + "DataType": "Number", + }, +} + MESSAGE_ID_KEY = "MessageId" +SEND_MESSAGE_DEFAULTS = { + "DelaySeconds": 0, + "MessageAttributes": {}, +} + +@mock_aws class TestSqsHook: - @mock_aws - def test_get_conn(self): + @pytest.fixture(autouse=True) + def setup_test_queue(self): + """Create a test queue before each test.""" hook = SqsHook(aws_conn_id="aws_default") + self.queue_url = hook.create_queue(queue_name=QUEUE_NAME) + yield + + @pytest.fixture + def hook(self): + return SqsHook(aws_conn_id="aws_default") + + def test_get_conn(self, hook): assert hook.get_conn() is not None + def test_create_queue(self, hook): + """Test that create_queue creates a queue and returns the queue URL.""" + queue_name = "test-create-queue" + queue_url = hook.create_queue(queue_name=queue_name) + + assert isinstance(queue_url, str) + assert queue_name in queue_url + + def test_create_queue_with_attributes(self, hook): + """Test creating a queue with custom attributes.""" + queue_name = "test-queue-with-attributes" + attributes = { + "DelaySeconds": "5", + "MaximumMessageSize": "262144", + } + + queue_url = hook.create_queue(queue_name=queue_name, attributes=attributes) + + assert isinstance(queue_url, str) + assert queue_name in queue_url + + # Verify attributes were actually set on the queue + queue_attrs = hook.get_conn().get_queue_attributes( + QueueUrl=queue_url, AttributeNames=["DelaySeconds", "MaximumMessageSize"] + ) + assert queue_attrs["Attributes"]["DelaySeconds"] == "5" + assert queue_attrs["Attributes"]["MaximumMessageSize"] == "262144" + + def test_send_message(self, hook): + """Test sending a message to a queue.""" + response = hook.send_message(queue_url=self.queue_url, message_body=MESSAGE_BODY) + + assert isinstance(response, dict) + assert MESSAGE_ID_KEY in response + assert "MD5OfMessageBody" in response + + def test_send_message_with_attributes(self, hook): + """Test sending a message with message attributes.""" + message_attributes = { + "Author": { + "StringValue": "test-user", + "DataType": "String", + }, + "Priority": { + "StringValue": "1", + "DataType": "Number", + }, + } Review Comment: Same as above, you already have a `MSG_ATTRIBUTES` global -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
