Testing consumer directly

Author: harishkannarao

app/config.py

@@ -37,6 +37,7 @@ class Settings(BaseSettings):
 
     # rabbit mq queue/exchange settings end
     app_rabbit_inbound_queue: str = "in-bound-queue"
+    app_rabbit_inbound_concurrency: int = 10
     app_rabbit_inbound_exchange: str = "in-bound-ex"
     app_rabbit_inbound_routing_key: str = "in-bound-rk"
 

app/consumer/__init__.py

@@ -0,0 +1,29 @@
+import asyncio
+from _asyncio import Task
+
+import structlog
+from aio_pika.abc import AbstractIncomingMessage
+from app.rabbit_mq.rabbit_mq_client import get_connection
+from app.config import settings
+
+
+async def process_inbound_message_task(message: AbstractIncomingMessage):
+    logger = structlog.get_logger()
+    async with message.process():  # Automatically ACKs if no exception occurs
+        payload_string: str = message.body.decode()
+        logger.info(
+            f"Processed inbound message {payload_string}", payload=payload_string
+        )
+
+
+async def start_inbound_consumer() -> Task[str]:
+    logger = structlog.get_logger()
+    # Create a dedicated channel for this consumer group
+    channel = await get_connection().channel()
+    # MAX CONCURRENCY SETTING: Limit unacknowledged messages on this channel
+    await channel.set_qos(prefetch_count=settings.app_rabbit_inbound_concurrency)
+    queue = await channel.get_queue(settings.app_rabbit_inbound_queue)
+    # Fire-and-forget the consumer loop into the background loop
+    task: Task[str] = asyncio.create_task(queue.consume(process_inbound_message_task))
+    logger.info("Started inbound message consumer")
+    return task

app/consumer/inbound_consumer.py

@@ -1,3 +1,6 @@
+import asyncio
+from asyncio import Task
+
 import structlog
 import uvicorn
 from fastapi import FastAPI
@@ -34,6 +37,7 @@
 from app.routers.dependency import router as dependency_router
 from app.routers.external_faq import router as external_faq_router
 from app.routers.sample_form import router as sample_form_router
+from app.consumer.inbound_consumer import start_inbound_consumer
 
 # Initialize logging at application startup
 setup_logging(json_logs=settings.app_json_logs, db_logs=settings.app_db_log_sql)
@@ -70,12 +74,21 @@ async def lifespan(_app: FastAPI):
         apply_db_migrations()
     if settings.app_db_enabled:
         await database.connect()
+    consumer_tasks: list[Task[str]] = []
     if settings.app_rabbit_mq_connect:
         await open_rabbit_connection()
         if not settings.app_rabbit_mq_passive:
             await configure_rabbitmq()
+        consumer_tasks.append(await start_inbound_consumer())
     yield
     if settings.app_rabbit_mq_connect:
+        # Cancel any long-running active consumer loop tasks
+        for task in consumer_tasks:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
         await close_rabbit_connection()
     if engine is not None:
         engine.dispose()

app/main.py

@@ -0,0 +1,27 @@
+import json
+
+import aio_pika
+import structlog
+from fastapi.encoders import jsonable_encoder
+
+from app.rabbit_mq.rabbit_mq_client import get_connection
+from app.model.response.sample import Sample
+from app.config import settings
+
+
+async def publish_to_inbound(samples: list[Sample]):
+    logger = structlog.get_logger()
+    async with get_connection().channel() as channel:
+        exchange = await channel.get_exchange(settings.app_rabbit_inbound_exchange)
+        payload_string: str = json.dumps(jsonable_encoder(samples))
+        message = aio_pika.Message(
+            body=payload_string.encode(),
+            delivery_mode=aio_pika.DeliveryMode.PERSISTENT,  # Survives broker restart
+        )
+        await exchange.publish(
+            message, routing_key=settings.app_rabbit_inbound_routing_key
+        )
+        logger.info(
+            f"Published to inbound queue {payload_string}", payload=payload_string
+        )
+        return

app/producer/__init__.py

@@ -2,7 +2,7 @@
 
 import aio_pika
 import structlog
-from aio_pika.abc import AbstractRobustConnection, AbstractChannel
+from aio_pika.abc import AbstractRobustConnection
 
 from app.config import settings
 
@@ -65,9 +65,3 @@ def get_connection() -> AbstractRobustConnection:
             "Ensure 'set_connection' was called during app startup."
         )
     return _connection
-
-
-async def get_channel() -> AbstractChannel:
-    """Utility helper to quickly grab an isolated channel off the global connection."""
-    conn = get_connection()
-    return await conn.channel()

app/producer/inbound_producer.py

@@ -2,13 +2,13 @@
 import structlog
 
 from app.config import settings
-from app.rabbit_mq.rabbit_mq_client import get_channel
+from app.rabbit_mq.rabbit_mq_client import get_connection
 
 
 async def configure_rabbitmq():
     logger = structlog.get_logger()
 
-    async with await get_channel() as channel:
+    async with await get_connection().channel() as channel:
         logger.info("Configuring queues and exchanges")
 
         inbound_queue = await channel.declare_queue(

app/rabbit_mq/rabbit_mq_client.py

@@ -0,0 +1,87 @@
+import datetime
+import json
+import uuid
+from decimal import Decimal
+from typing import MutableMapping, Any
+
+import aio_pika
+import pytest
+import structlog
+from assertpy import assert_that
+from fastapi.encoders import jsonable_encoder
+from fastapi.testclient import TestClient
+from tenacity import Retrying, stop_after_delay, wait_fixed
+
+
+from app.model.response.sample import Sample
+
+
+@pytest.mark.asyncio
+async def test_inbound_consumer_handles_message_successfully(
+    test_client: TestClient,
+    captured_logs: list[MutableMapping[str, Any]],
+):
+    sample1: Sample = Sample(
+        id=uuid.uuid4(),
+        username=f"user-{uuid.uuid4()}",
+        bool_field=True,
+        float_field=2.0,
+        decimal_field=Decimal("3.1"),
+        created_datetime=datetime.datetime.now(datetime.UTC),
+        updated_datetime=datetime.datetime.now(datetime.UTC),
+        version=1,
+    )
+    sample2: Sample = Sample(
+        id=uuid.uuid4(),
+        username=f"user-{uuid.uuid4()}",
+        bool_field=None,
+        float_field=None,
+        decimal_field=None,
+        created_datetime=datetime.datetime.now(datetime.UTC),
+        updated_datetime=datetime.datetime.now(datetime.UTC),
+        version=1,
+    )
+    samples: list[Sample] = [sample1, sample2]
+
+    await publish_to_inbound(samples)
+
+    for attempt in Retrying(
+        stop=stop_after_delay(5), wait=wait_fixed(0.5), reraise=True
+    ):
+        with attempt:
+            assert_that(len(captured_logs)).is_greater_than(0)
+            inbound_consumer_logs = list(
+                filter(
+                    lambda entry: str(entry["event"]).startswith(
+                        "Processed inbound message"
+                    ),
+                    captured_logs,
+                )
+            )
+            assert_that(inbound_consumer_logs).is_length(1)
+
+
+async def publish_to_inbound(samples: list[Sample]):
+    logger = structlog.get_logger()
+    from app.config import settings
+
+    rabbit_mq_url = (
+        f"amqp://{settings.app_rabbit_mq_username}:{settings.app_rabbit_mq_password}"
+        f"@{settings.app_rabbit_mq_host}:{settings.app_rabbit_mq_port}/{settings.app_rabbit_mq_vhost}"
+    )
+
+    async with await aio_pika.connect_robust(rabbit_mq_url) as connection:
+        async with connection.channel() as channel:
+            exchange = await channel.get_exchange(settings.app_rabbit_inbound_exchange)
+            payload_string: str = json.dumps(jsonable_encoder(samples))
+            message = aio_pika.Message(
+                body=payload_string.encode(),
+                delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+            )
+            await exchange.publish(
+                message, routing_key=settings.app_rabbit_inbound_routing_key
+            )
+            logger.info(
+                f"Published to inbound queue {payload_string}", payload=payload_string
+            )
+            return