Repository: run-llama/llama_deploy

Issue:
Make the topic configurable for AWSMessageQueue



Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6e8fc5ea..072f75b8 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -36,6 +36,7 @@ repos:
             "llama_index_core",
             "types-Deprecated",
             "types-PyYAML",
+            "types-botocore",
             "types-aiobotocore",
             "types-protobuf==4.24.0.4",
             "types-redis",
@@ -48,7 +49,7 @@ repos:
             --ignore-missing-imports,
             --python-version=3.11,
           ]
-        exclude: ^(examples/|e2e_tests/)
+        exclude: ^(examples/|e2e_tests/|tests/message_queues/test_aws.py)
 
   - repo: https://github.com/adamchainz/blacken-docs
     rev: 1.16.0
diff --git a/llama_deploy/__init__.py b/llama_deploy/__init__.py
index a7ca6af8..7f5b7c7b 100644
--- a/llama_deploy/__init__.py
+++ b/llama_deploy/__init__.py
@@ -24,7 +24,7 @@
 root_logger.addHandler(console_handler)
 
 root_logger.setLevel(logging.INFO)
-root_logger.propagate = False
+root_logger.propagate = True
 
 
 __all__ = [
diff --git a/llama_deploy/apiserver/deployment.py b/llama_deploy/apiserver/deployment.py
index 7f72bb34..987b14f7 100644
--- a/llama_deploy/apiserver/deployment.py
+++ b/llama_deploy/apiserver/deployment.py
@@ -283,7 +283,7 @@ def _load_message_queue_client(
             cfg = SimpleMessageQueueConfig()
 
         if cfg.type == "aws":
-            return AWSMessageQueue(**cfg.model_dump())
+            return AWSMessageQueue(cfg)  # type: ignore
         elif cfg.type == "kafka":
             return KafkaMessageQueue(cfg)  # type: ignore
         elif cfg.type == "rabbitmq":
diff --git a/llama_deploy/deploy/deploy.py b/llama_deploy/deploy/deploy.py
index 1a47f03a..e6432780 100644
--- a/llama_deploy/deploy/deploy.py
+++ b/llama_deploy/deploy/deploy.py
@@ -64,7 +64,7 @@ def _get_message_queue_client(config: BaseSettings) -> BaseMessageQueue:
     if isinstance(config, SimpleMessageQueueConfig):
         return SimpleMessageQueue(config)  # type: ignore
     elif isinstance(config, AWSMessageQueueConfig):
-        return AWSMessageQueue(**config.model_dump())
+        return AWSMessageQueue(config)  # type: ignore
     elif isinstance(config, KafkaMessageQueueConfig):
         return KafkaMessageQueue(config)  # type: ignore
     elif isinstance(config, RabbitMQMessageQueueConfig):
diff --git a/llama_deploy/message_queues/aws.py b/llama_deploy/message_queues/aws.py
index f8ae9a15..da09594b 100644
--- a/llama_deploy/message_queues/aws.py
+++ b/llama_deploy/message_queues/aws.py
@@ -1,22 +1,22 @@
 """AWS SNS and SQS Message Queue."""
 
 import json
+from asyncio import CancelledError
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
+from typing import TYPE_CHECKING, Any, Literal
 
-from pydantic import BaseModel, Field, PrivateAttr, SecretStr
+from pydantic import BaseModel, Field, SecretStr
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from llama_deploy.message_consumers.base import (
     BaseMessageQueueConsumer,
     StartConsumingCallable,
 )
-from llama_deploy.message_queues.base import BaseMessageQueue
+from llama_deploy.message_queues.base import AbstractMessageQueue
 from llama_deploy.messages.base import QueueMessage
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from aiobotocore.session import AioSession, ClientCreatorContext
-    from botocore.config import Config
 
 logger = getLogger(__name__)
 
@@ -48,12 +48,12 @@ class AWSMessageQueueConfig(BaseSettings):
     model_config = SettingsConfigDict()
 
     type: Literal["aws"] = Field(default="aws", exclude=True)
-    aws_region: str
-    aws_access_key_id: Optional[str] = Field(default=None, exclude=True)
-    aws_secret_access_key: Optional[str] = Field(default=None, exclude=True)
+    aws_region: str = "us-east-1"
+    aws_access_key_id: SecretStr | None = Field(default=None, exclude=True)
+    aws_secret_access_key: SecretStr | None = Field(default=None, exclude=True)
 
 
-class AWSMessageQueue(BaseMessageQueue):
+class AWSMessageQueue(AbstractMessageQueue):
     """AWS SQS integration with aiobotocore client.
 
     This class creates and interacts with SNS topics and SQS queues. It includes methods
@@ -77,17 +77,7 @@ class AWSMessageQueue(BaseMessageQueue):
         aws_secret_access_key (Optional[SecretStr]): AWS secret access key (fallback to environment variables).
     """
 
-    aws_region: str
-    aws_access_key_id: Optional[SecretStr] = None
-    aws_secret_access_key: Optional[SecretStr] = None
-    topics: List["Topic"] = []
-    queues: List["Queue"] = []
-    subscriptions: List["Subscription"] = []
-
-    _aio_session: Optional["AioSession"] = PrivateAttr(None)
-    _retry_config: "Config" = PrivateAttr()
-
-    def __init__(self, **kwargs: Any) -> None:
+    def __init__(self, config: AWSMessageQueueConfig | None = None) -> None:
         """Initialize AWSMessageQueue with config."""
         try:
             from botocore.config import Config
@@ -95,17 +85,14 @@ def __init__(self, **kwargs: Any) -> None:
             raise ValueError(
                 "Missing `botocore` package. Please install by running `pip install llama-deploy[aws]`."
             )
-
-        super().__init__(**kwargs)
-        self._aio_session = None
-
+        self._config = config or AWSMessageQueueConfig()
+        self._topics: list["Topic"] = []
+        self._queues: list["Queue"] = []
+        self._subscriptions: list["Subscription"] = []
+        self._aio_session: "AioSession | None" = None
         # AWS retry configuration with exponential backoff
         self._retry_config = Config(retries={"max_attempts": 5, "mode": "adaptive"})
-
-        if not self.aws_region:
-            raise ValueError("AWS region must be provided.")
-
-        if not self.aws_access_key_id and not self.aws_secret_access_key:
+        if not self._config.aws_access_key_id or not self._config.aws_secret_access_key:
             logger.info(
                 "Using default AWS credential provider chain (IAM Role or environment variables)."
             )
@@ -125,17 +112,17 @@ def _get_client(self, service_name: str) -> "ClientCreatorContext":
         """Returns a client context manager."""
         session = self._get_aio_session()
         client_kwargs = {
-            "region_name": self.aws_region,
+            "region_name": self._config.aws_region,
             "config": self._retry_config,
             "service_name": service_name,
         }
 
         # IAM Role or environment variable fallback
-        if self.aws_access_key_id and self.aws_secret_access_key:
+        if self._config.aws_access_key_id and self._config.aws_secret_access_key:
             client_kwargs.update(
                 {
-                    "aws_access_key_id": self.aws_access_key_id.get_secret_value(),
-                    "aws_secret_access_key": self.aws_secret_access_key.get_secret_value(),
+                    "aws_access_key_id": self._config.aws_access_key_id.get_secret_value(),
+                    "aws_secret_access_key": self._config.aws_secret_access_key.get_secret_value(),
                 }
             )
 
@@ -149,6 +136,7 @@ async def get_topic_by_name(self, topic_name: str) -> "Topic":
             try:
                 # First, check if the topic exists
                 response = await client.list_topics()  # type: ignore
+                print(response)
                 for topic in response.get("Topics", []):
                     if f"{topic_name}.fifo" in topic["TopicArn"]:
                         logger.info(f"SNS topic {topic_name} already exists.")
@@ -180,8 +168,8 @@ async def _create_sns_topic(self, topic_name: str) -> "Topic":
                 raise
 
         topic = Topic(arn=response["TopicArn"], name=topic_name)
-        if topic.arn not in [t.arn for t in self.topics]:
-            self.topics.append(topic)
+        if topic.arn not in [t.arn for t in self._topics]:
+            self._topics.append(topic)
         return topic
 
     async def _create_sqs_queue(self, queue_name: str) -> "Queue":
@@ -211,8 +199,8 @@ async def _create_sqs_queue(self, queue_name: str) -> "Queue":
                 raise
 
         queue = Queue(arn=queue_arn, url=queue_url, name=queue_name)
-        if queue.arn not in [q.arn for q in self.queues]:
-            self.queues.append(queue)
+        if queue.arn not in [q.arn for q in self._queues]:
+            self._queues.append(queue)
         return queue
 
     async def _update_queue_policy(self, queue: "Queue", topic: "Topic") -> None:
@@ -253,8 +241,8 @@ async def _subscribe_queue_to_topic(
                 )
                 raise
         subscription = Subscription(arn=response["SubscriptionArn"])
-        if subscription.arn not in [s.arn for s in self.subscriptions]:
-            self.subscriptions.append(subscription)
+        if subscription.arn not in [s.arn for s in self._subscriptions]:
+            self._subscriptions.append(subscription)
 
         # Update the SQS queue policy to allow SNS topic to send messages
         await self._update_queue_policy(queue, topic)
@@ -266,7 +254,7 @@ async def _publish(self, message: QueueMessage, topic: str) -> Any:
         from botocore.exceptions import ClientError
 
         message_body = json.dumps(message.model_dump())
-        _topic = await self.get_topic_by_name(message.type)
+        _topic = await self.get_topic_by_name(topic)
 
         try:
             async with self._get_client("sns") as client:
@@ -277,13 +265,15 @@ async def _publish(self, message: QueueMessage, topic: str) -> Any:
                     MessageGroupId=message.id_,  # Assigning message id as the group id for simplicity
                     MessageDeduplicationId=message.id_,
                 )
-                logger.info(f"Published {message.type} message {message.id_}")
+                logger.info(
+                    f"Published message {message.id_} of type {message.type} on topic {topic}"
+                )
                 return response
         except ClientError as e:
             logger.error(f"Could not publish message to SQS queue: {e}")
             raise
 
-    async def cleanup(self, *args: Any, **kwargs: Dict[str, Any]) -> None:
+    async def cleanup(self, *args: Any, **kwargs: dict[str, Any]) -> None:
         """Perform cleanup of queues and topics."""
         from botocore.exceptions import ClientError
 
@@ -292,7 +282,7 @@ async def cleanup(self, *args: Any, **kwargs: Dict[str, Any]) -> None:
             self._get_client("sns") as sns_client,
         ):
             # Delete all SQS queues
-            for queue in self.queues:
+            for queue in self._queues:
                 try:
                     await sqs_client.delete_queue(QueueUrl=queue.url)  # type: ignore
                     logger.info(f"Deleted SQS queue {queue.name}")
@@ -300,7 +290,7 @@ async def cleanup(self, *args: Any, **kwargs: Dict[str, Any]) -> None:
                     logger.error(f"Could not delete SQS queue {queue.name}: {e}")
 
             # Delete all SNS topics
-            for topic in self.topics:
+            for topic in self._topics:
                 try:
                     await sns_client.delete_topic(TopicArn=topic.arn)  # type: ignore
                     logger.info(f"Deleted SNS topic {topic.name}")
@@ -313,41 +303,46 @@ async def register_consumer(
         """Register a new consumer."""
         from botocore.exceptions import ClientError
 
-        _topic = await self._create_sns_topic(topic_name=consumer.message_type)
-        queue = await self._create_sqs_queue(queue_name=consumer.message_type)
+        topic = topic or consumer.message_type
+        _topic = await self._create_sns_topic(topic_name=topic)
+        queue = await self._create_sqs_queue(queue_name=topic)
         await self._subscribe_queue_to_topic(queue=queue, topic=_topic)
-        logger.info(f"Registered consumer {consumer.id_}: {consumer.message_type}")
+        logger.info(f"Registered consumer {consumer.id_} for topic {topic}")
 
         async def start_consuming_callable() -> None:
             """Start consuming messages."""
             while True:
-                async with self._get_client("sqs") as client:
-                    try:
-                        response = await client.receive_message(  # type: ignore
-                            QueueUrl=queue.url,
-                            WaitTimeSeconds=2,
-                        )
-                        messages = response.get("Messages", [])
-                        for msg in messages:
-                            receipt_handle = msg["ReceiptHandle"]
-                            message_body = json.loads(msg["Body"])
-                            queue_message_data = json.loads(message_body["Message"])
-                            queue_message = QueueMessage.model_validate(
-                                queue_message_data
+                try:
+                    async with self._get_client("sqs") as client:
+                        try:
+                            response = await client.receive_message(  # type: ignore
+                                QueueUrl=queue.url,
+                                WaitTimeSeconds=2,
                             )
-                            await consumer.process_message(queue_message)
-                            await client.delete_message(  # type: ignore
-                                QueueUrl=queue.url, ReceiptHandle=receipt_handle
+                            messages = response.get("Messages", [])
+                            for msg in messages:
+                                receipt_handle = msg["ReceiptHandle"]
+                                message_body = json.loads(msg["Body"])
+                                queue_message_data = json.loads(message_body["Message"])
+                                queue_message = QueueMessage.model_validate(
+                                    queue_message_data
+                                )
+                                await consumer.process_message(queue_message)
+                                await client.delete_message(  # type: ignore
+                                    QueueUrl=queue.url, ReceiptHandle=receipt_handle
+                                )
+                        except ClientError as e:
+                            logger.error(
+                                f"Error receiving messages from SQS queue: {e}"
                             )
-                    except ClientError as e:
-                        logger.error(f"Error receiving messages from SQS queue: {e}")
+                except CancelledError:
+                    return
 
         return start_consuming_callable
 
     def as_config(self) -> BaseModel:
         """Return the current configuration as an AWSMessageQueueConfig object."""
-        config = AWSMessageQueueConfig(aws_region=self.aws_region)
-        return config
+        return self._config
 
     async def deregister_consumer(self, consumer: BaseMessageQueueConsumer) -> Any:
         """Deregister a consumer.


Test patch:
diff --git a/.github/workflows/e2e_test.yml b/.github/workflows/e2e_test.yml
index 21124f21..fd5983fe 100644
--- a/.github/workflows/e2e_test.yml
+++ b/.github/workflows/e2e_test.yml
@@ -33,4 +33,7 @@ jobs:
         shell: bash
         run: poetry install --with dev
       - name: Run All E2E Tests
+        env:
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
         run: poetry run pytest e2e_tests/
diff --git a/e2e_tests/message_queues/aws/__init__.py b/e2e_tests/message_queues/aws/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/e2e_tests/message_queues/aws/conftest.py b/e2e_tests/message_queues/aws/conftest.py
new file mode 100644
index 00000000..3559b625
--- /dev/null
+++ b/e2e_tests/message_queues/aws/conftest.py
@@ -0,0 +1,90 @@
+import asyncio
+import multiprocessing
+import time
+
+import pytest
+
+from llama_deploy import (
+    ControlPlaneConfig,
+    WorkflowServiceConfig,
+    deploy_core,
+    deploy_workflow,
+)
+from llama_deploy.message_queues import AWSMessageQueue, AWSMessageQueueConfig
+
+from .workflow import BasicWorkflow
+
+
+@pytest.fixture
+def mq():
+    return AWSMessageQueue()
+
+
+def run_workflow_one():
+    asyncio.run(
+        deploy_workflow(
+            BasicWorkflow(timeout=10, name="Workflow one"),
+            WorkflowServiceConfig(
+                host="127.0.0.1",
+                port=8003,
+                service_name="basic",
+            ),
+            ControlPlaneConfig(topic_namespace="core_one", port=8001),
+        )
+    )
+
+
+def run_workflow_two():
+    asyncio.run(
+        deploy_workflow(
+            BasicWorkflow(timeout=10, name="Workflow two"),
+            WorkflowServiceConfig(
+                host="127.0.0.1",
+                port=8004,
+                service_name="basic",
+            ),
+            ControlPlaneConfig(topic_namespace="core_two", port=8002),
+        )
+    )
+
+
+def run_core_one():
+    asyncio.run(
+        deploy_core(
+            ControlPlaneConfig(topic_namespace="core_one", port=8001),
+            AWSMessageQueueConfig(),
+        )
+    )
+
+
+def run_core_two():
+    asyncio.run(
+        deploy_core(
+            ControlPlaneConfig(topic_namespace="core_two", port=8002),
+            AWSMessageQueueConfig(),
+        )
+    )
+
+
+@pytest.fixture
+def control_planes(kafka_service):
+    p1 = multiprocessing.Process(target=run_core_one)
+    p1.start()
+
+    p2 = multiprocessing.Process(target=run_core_two)
+    p2.start()
+
+    time.sleep(3)
+
+    p3 = multiprocessing.Process(target=run_workflow_one)
+    p3.start()
+
+    p4 = multiprocessing.Process(target=run_workflow_two)
+    p4.start()
+
+    yield
+
+    p1.terminate()
+    p2.terminate()
+    p3.terminate()
+    p4.terminate()
diff --git a/e2e_tests/message_queues/aws/test_message_queue.py b/e2e_tests/message_queues/aws/test_message_queue.py
new file mode 100644
index 00000000..35406690
--- /dev/null
+++ b/e2e_tests/message_queues/aws/test_message_queue.py
@@ -0,0 +1,50 @@
+import asyncio
+
+import pytest
+
+from llama_deploy import Client
+from llama_deploy.message_consumers.callable import CallableMessageConsumer
+from llama_deploy.messages import QueueMessage
+
+
+@pytest.mark.asyncio
+async def test_roundtrip(mq):
+    received_messages = []
+
+    # register a consumer
+    def message_handler(message: QueueMessage) -> None:
+        received_messages.append(message)
+
+    test_consumer = CallableMessageConsumer(
+        message_type="test_message", handler=message_handler
+    )
+    start_consuming_callable = await mq.register_consumer(test_consumer, topic="test")
+    t = asyncio.create_task(start_consuming_callable())
+    await asyncio.sleep(1)
+
+    # produce a message
+    test_message = QueueMessage(type="test_message", data={"message": "this is a test"})
+    await mq.publish(test_message, topic="test")
+    await asyncio.sleep(1)
+
+    t.cancel()
+    await t
+
+    assert len(received_messages) == 1
+    assert test_message in received_messages
+
+
+@pytest.mark.asyncio
+async def _test_multiple_control_planes(control_planes):
+    c1 = Client(control_plane_url="http://localhost:8001")
+    c2 = Client(control_plane_url="http://localhost:8002")
+
+    session = await c1.core.sessions.create()
+    r1 = await session.run("basic", arg="Hello One!")
+    await c1.core.sessions.delete(session.id)
+    assert r1 == "Workflow one received Hello One!"
+
+    session = await c2.core.sessions.create()
+    r2 = await session.run("basic", arg="Hello Two!")
+    await c2.core.sessions.delete(session.id)
+    assert r2 == "Workflow two received Hello Two!"
diff --git a/e2e_tests/message_queues/aws/workflow.py b/e2e_tests/message_queues/aws/workflow.py
new file mode 100644
index 00000000..cf3d863d
--- /dev/null
+++ b/e2e_tests/message_queues/aws/workflow.py
@@ -0,0 +1,12 @@
+from llama_index.core.workflow import Context, StartEvent, StopEvent, Workflow, step
+
+
+class BasicWorkflow(Workflow):
+    def __init__(self, *args, **kwargs):
+        self._name = kwargs.pop("name")
+        super().__init__(*args, **kwargs)
+
+    @step()
+    async def run_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        received = ev.get("arg")
+        return StopEvent(result=f"{self._name} received {received}")
diff --git a/tests/apiserver/test_deployment.py b/tests/apiserver/test_deployment.py
index 84b571be..7bd57079 100644
--- a/tests/apiserver/test_deployment.py
+++ b/tests/apiserver/test_deployment.py
@@ -91,7 +91,7 @@ def test_deployment__load_message_queues(mocked_deployment: Deployment) -> None:
         mocked_config = mock.MagicMock(type="aws")
         mocked_config.model_dump.return_value = {"foo": "aws"}
         mocked_deployment._load_message_queue_client(mocked_config)
-        m.assert_called_with(**{"foo": "aws"})
+        m.assert_called_with(mocked_config)
 
     with mock.patch("llama_deploy.apiserver.deployment.KafkaMessageQueue") as m:
         mocked_config = mock.MagicMock(type="kafka")
diff --git a/tests/message_queues/test_aws.py b/tests/message_queues/test_aws.py
new file mode 100644
index 00000000..bb988060
--- /dev/null
+++ b/tests/message_queues/test_aws.py
@@ -0,0 +1,431 @@
+import logging
+from typing import Any
+from unittest import mock
+
+import pytest
+import pytest_asyncio
+from aiobotocore.session import AioSession
+from botocore.exceptions import ClientError
+from pydantic import SecretStr
+
+from llama_deploy.message_queues.aws import AWSMessageQueue, AWSMessageQueueConfig
+
+
+@pytest_asyncio.fixture
+async def aws_queue(monkeypatch: Any) -> Any:
+    mq = AWSMessageQueue()
+    yield mq
+    await mq.cleanup()
+
+
+@pytest.fixture
+def botocore_not_installed(monkeypatch: Any) -> Any:
+    """Mock the import mechanism to raise ImportError even when package is installed"""
+    # Store the original import
+    original_import = __import__
+
+    def mock_import(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name.startswith("botocore"):
+            raise ImportError(f"No module named '{name}'")
+        return original_import(name, *args, **kwargs)
+
+    # Replace the import mechanism with our mock
+    monkeypatch.setattr("builtins.__import__", mock_import)
+
+
+@pytest.fixture
+def aiobotocore_not_installed(monkeypatch: Any) -> Any:
+    """Mock the import mechanism to raise ImportError even when package is installed"""
+    # Store the original import
+    original_import = __import__
+
+    def mock_import(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name.startswith("aiobotocore"):
+            raise ImportError(f"No module named '{name}'")
+        return original_import(name, *args, **kwargs)
+
+    # Replace the import mechanism with our mock
+    monkeypatch.setattr("builtins.__import__", mock_import)
+
+
+def test_ctor(caplog) -> None:
+    caplog.set_level(logging.INFO, logger="llama_deploy.message_queues.aws")
+    AWSMessageQueue(AWSMessageQueueConfig(aws_access_key_id=SecretStr("")))
+    assert "Using default AWS credential provider chain" in caplog.text
+
+
+def test_missing_deps(botocore_not_installed: Any) -> None:
+    with pytest.raises(ValueError, match="Missing `botocore` package"):
+        AWSMessageQueue()
+
+
+def test__get_aio_session_missing_deps(aiobotocore_not_installed: Any) -> None:
+    with pytest.raises(ValueError, match="Missing `aiobotocore`"):
+        mq = AWSMessageQueue()
+        mq._get_aio_session()
+
+
+def test__get_aio_session(aws_queue: AWSMessageQueue) -> None:
+    mq = AWSMessageQueue()
+    assert type(mq._get_aio_session()) is AioSession
+
+
+@pytest.mark.asyncio
+async def test__get_client() -> None:
+    mq = AWSMessageQueue(
+        AWSMessageQueueConfig(
+            aws_access_key_id=SecretStr("secret"),
+            aws_secret_access_key=SecretStr("moar_secret"),
+        )
+    )
+    async with mq._get_client("sns") as client:
+        assert client
+
+
+@pytest.mark.asyncio
+async def test_get_topic_by_name(aws_queue: AWSMessageQueue) -> None:
+    mock_response = {
+        "Topics": [{"TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic.fifo"}]
+    }
+
+    class MockSNSClient:
+        def __init__(self, should_raise: bool = False):
+            self.should_raise = should_raise
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        async def list_topics(self):
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="list_topics",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return mock_response
+
+    aws_queue._get_client = lambda x: MockSNSClient()  # type: ignore
+    topic = await aws_queue.get_topic_by_name("test-topic")
+    assert topic.arn == "arn:aws:sns:us-east-1:123456789012:test-topic.fifo"
+    assert topic.name == "test-topic"
+
+    mock_response["Topics"] = []
+    with pytest.raises(ValueError, match="Could not find topic test-topic"):
+        await aws_queue.get_topic_by_name("test-topic")
+
+    # Test ClientError exception
+    aws_queue._get_client = lambda x: MockSNSClient(should_raise=True)  # type: ignore
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue.get_topic_by_name("test-topic")
+
+
+@pytest.mark.asyncio
+async def test_create_sns_topic(aws_queue: AWSMessageQueue) -> None:
+    from botocore.exceptions import ClientError
+
+    # Mock responses
+    existing_topic_response = {
+        "Topics": [
+            {"TopicArn": "arn:aws:sns:us-east-1:123456789012:existing-topic.fifo"}
+        ]
+    }
+
+    new_topic_response = {
+        "TopicArn": "arn:aws:sns:us-east-1:123456789012:new-topic.fifo"
+    }
+
+    class MockSNSClient:
+        def __init__(self, should_raise: bool = False):
+            self.should_raise = should_raise
+            self.create_topic_called = False
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        async def delete_topic(self, *args, **kwargs):
+            pass
+
+        async def list_topics(self):
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="list_topics",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return existing_topic_response
+
+        async def create_topic(self, Name: str, Attributes: dict):
+            self.create_topic_called = True
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="create_topic",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return new_topic_response
+
+    # Test finding existing topic
+    mock_client = MockSNSClient()
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+
+    topic = await aws_queue._create_sns_topic("existing-topic")
+    assert topic.arn == "arn:aws:sns:us-east-1:123456789012:existing-topic.fifo"
+    assert topic.name == "existing-topic"
+    assert not mock_client.create_topic_called
+
+    # Test creating new topic
+    existing_topic_response["Topics"] = []  # Clear existing topics
+    topic = await aws_queue._create_sns_topic("new-topic")
+    assert topic.arn == "arn:aws:sns:us-east-1:123456789012:new-topic.fifo"
+    assert topic.name == "new-topic"
+    assert mock_client.create_topic_called
+    assert topic in aws_queue._topics
+
+    # Test ClientError during list_topics
+    mock_client = MockSNSClient(should_raise=True)
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue._create_sns_topic("error-topic")
+
+    # Test ClientError during create_topic
+    mock_client = MockSNSClient(should_raise=True)
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+    existing_topic_response["Topics"] = []  # Force create_topic to be called
+
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue._create_sns_topic("error-topic")
+
+
+@pytest.mark.asyncio
+async def test_create_sqs_queue(aws_queue: AWSMessageQueue) -> None:
+    from botocore.exceptions import ClientError
+
+    # Mock responses
+    existing_queue_response = {
+        "QueueUrls": [
+            "https://sqs.us-east-1.amazonaws.com/123456789012/existing-queue.fifo"
+        ]
+    }
+
+    new_queue_response = {
+        "QueueUrl": "https://sqs.us-east-1.amazonaws.com/123456789012/new-queue.fifo"
+    }
+
+    queue_attributes_response = {
+        "Attributes": {"QueueArn": "arn:aws:sqs:us-east-1:123456789012:test-queue.fifo"}
+    }
+
+    class MockSQSClient:
+        def __init__(self, should_raise: bool = False):
+            self.should_raise = should_raise
+            self.create_queue_called = False
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        async def delete_queue(self, *args, **kwargs):
+            pass
+
+        async def list_queues(self, QueueNamePrefix: str):
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="list_queues",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return existing_queue_response
+
+        async def create_queue(self, QueueName: str, Attributes: dict):
+            self.create_queue_called = True
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="create_queue",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return new_queue_response
+
+        async def get_queue_attributes(self, QueueUrl: str, AttributeNames: list):
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="get_queue_attributes",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return queue_attributes_response
+
+    # Test finding existing queue
+    mock_client = MockSQSClient()
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+
+    queue = await aws_queue._create_sqs_queue("existing-queue")
+    assert (
+        queue.url
+        == "https://sqs.us-east-1.amazonaws.com/123456789012/existing-queue.fifo"
+    )
+    assert queue.arn == "arn:aws:sqs:us-east-1:123456789012:test-queue.fifo"
+    assert queue.name == "existing-queue"
+    assert not mock_client.create_queue_called
+    assert queue in aws_queue._queues
+
+    # Test creating new queue
+    existing_queue_response["QueueUrls"] = []  # Clear existing queues
+    queue = await aws_queue._create_sqs_queue("new-queue")
+    assert (
+        queue.url == "https://sqs.us-east-1.amazonaws.com/123456789012/new-queue.fifo"
+    )
+    assert queue.arn == "arn:aws:sqs:us-east-1:123456789012:test-queue.fifo"
+    assert queue.name == "new-queue"
+    assert mock_client.create_queue_called
+
+    # Test ClientError during list_queues
+    mock_client = MockSQSClient(should_raise=True)
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue._create_sqs_queue("error-queue")
+
+    # Test ClientError during create_queue
+    mock_client = MockSQSClient(should_raise=True)
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+    existing_queue_response["QueueUrls"] = []  # Force create_queue to be called
+
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue._create_sqs_queue("error-queue")
+
+    # Test ClientError during get_queue_attributes
+    mock_client = MockSQSClient(should_raise=True)
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue._create_sqs_queue("error-queue")
+
+
+@pytest.mark.asyncio
+async def test_update_queue_policy(aws_queue: AWSMessageQueue) -> None:
+    class MockSQSClient:
+        def __init__(self, should_raise: bool = False):
+            self.should_raise = should_raise
+            self.set_queue_attributes_called = False
+            self.last_policy = None
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        async def set_queue_attributes(self, QueueUrl: str, Attributes: dict):
+            self.set_queue_attributes_called = True
+            self.last_policy = Attributes.get("Policy")
+
+    # Test successful policy update
+    mock_client = MockSQSClient()
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+
+    await aws_queue._update_queue_policy(
+        mock.MagicMock(arn="foo_arn"), mock.MagicMock(arn="bar_arn")
+    )  # type: ignore
+    assert mock_client.set_queue_attributes_called
+
+
+@pytest.mark.asyncio
+async def test_subscribe_queue_to_topic(aws_queue: AWSMessageQueue) -> None:
+    from botocore.exceptions import ClientError
+
+    class MockSNSClient:
+        def __init__(self, should_raise: bool = False):
+            self.should_raise = should_raise
+            self.subscribe_called = False
+
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc_val, exc_tb):
+            pass
+
+        async def subscribe(self, TopicArn: str, Protocol: str, Endpoint: str):
+            self.subscribe_called = True
+            if self.should_raise:
+                raise ClientError(
+                    operation_name="subscribe",
+                    error_response={
+                        "Error": {
+                            "Code": "InternalError",
+                            "Message": "AWS Internal Error",
+                        }
+                    },
+                )
+            return {
+                "SubscriptionArn": "arn:aws:sns:us-east-1:123456789012:test-topic:subscription-id"
+            }
+
+    # Mock for _update_queue_policy
+    update_policy_called = False
+
+    async def mock_update_queue_policy(queue, topic):  # type: ignore
+        nonlocal update_policy_called
+        update_policy_called = True
+
+    # Test successful subscription
+    mock_client = MockSNSClient()
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+    aws_queue._update_queue_policy = mock_update_queue_policy  # type: ignore
+
+    subscription = await aws_queue._subscribe_queue_to_topic(
+        mock.MagicMock(arn="foo_arn"), mock.MagicMock(arn="bar_arn")
+    )
+
+    assert mock_client.subscribe_called
+    assert update_policy_called
+    assert (
+        subscription.arn
+        == "arn:aws:sns:us-east-1:123456789012:test-topic:subscription-id"
+    )
+    assert subscription in aws_queue._subscriptions
+
+    # Test ClientError during subscribe
+    mock_client = MockSNSClient(should_raise=True)
+    aws_queue._get_client = lambda x: mock_client  # type: ignore
+    update_policy_called = False
+
+    with pytest.raises(ClientError, match="AWS Internal Error"):
+        await aws_queue._subscribe_queue_to_topic(
+            mock.MagicMock(arn="foo_arn"), mock.MagicMock(arn="bar_arn")
+        )
+    # Policy update shouldn't be called if subscription fails
+    assert not update_policy_called