aws_sqs_client.py

from typing import List, Dict

from boto3 import resource


class SqsClient:
    __slots__ = ['boto_resource', 'max_number_of_messages', 'pool_wait_time_seconds']
    boto_resource: resource
    max_number_of_messages: int
    pool_wait_time_seconds: int

    def __init__(self, boto_resource, max_number_of_messages=10, pool_wait_time_seconds=10):
        self.boto_resource = boto_resource
        self.max_number_of_messages = max_number_of_messages
        self.pool_wait_time_seconds = pool_wait_time_seconds

    def get_messages(self, queue_name: str) -> List[str]:
        return self.boto_resource.Queue(queue_name).receive_messages(
            AttributeNames=["All"],
            MaxNumberOfMessages=self.max_number_of_messages,
            MessageAttributeNames=["string"],
            WaitTimeSeconds=self.pool_wait_time_seconds,
        )

    def send_message(self, message: str, queue_name: str) -> Dict[str, str]:
        queue_resource = self.boto_resource.Queue(queue_name)
        return queue_resource.send_message(MessageBody=message)

test_aws_sqs_client.py

from unittest import TestCase
from unittest.mock import MagicMock

from app.clients.sqs import SqsClient


class SqsClientTest(TestCase):
    def setUp(self):
        self.sample_message = "hello!"
        self.queue_name = "test-queue-name"
        self.max_number_of_messages = 1
        self.pool_wait_time_seconds = 1
        self.list_messages = [self.sample_message]

    def test_when_I_call_get_messages_with_correct_params_then_should_be_success(self):
        boto_resource_mock = MagicMock()

        boto_resource_mock.Queue(self.queue_name).receive_messages = MagicMock(return_value=self.list_messages)

        sqs_client = SqsClient(
            boto_resource=boto_resource_mock,
            max_number_of_messages=self.max_number_of_messages,
            pool_wait_time_seconds=self.pool_wait_time_seconds,
        )

        messages = sqs_client.get_messages(self.queue_name)
        self.assertEqual(messages, self.list_messages)

    def test_when_I_call_send_message_with_correct_params_then_should_be_success(self):
        boto_resource_mock = MagicMock()
        boto_resource_mock.Queue().send_message = MagicMock(return_value={"MessageId": "123456"})

        sqs_client = SqsClient(
            boto_resource=boto_resource_mock,
            max_number_of_messages=self.max_number_of_messages,
            pool_wait_time_seconds=self.pool_wait_time_seconds
        )

        send_message_response = sqs_client.send_message("Message body", self.queue_name)
        self.assertEqual("123456", send_message_response.get("MessageId"))