apm-component 1.0.0__py3-none-any.whl

apm_component/broker_rabbitmq_aio_pika_async.py

@@ -0,0 +1,697 @@
+"""
+RabbitMQ Broker Component - aio_pika module (async)
+Covers: async connection, errors, queue/exchange management, publish, consume,
+        ack/nack, message properties, bulk publish, dead letter, serialization
+"""
+
+import asyncio
+import json
+import time
+import aio_pika
+from apm_component.config import get_rabbitmq_config
+
+
+_PREFIX = "ct_aiopika_"
+
+
+# ---------------------------------------------------------------------------
+# Connection helpers
+# ---------------------------------------------------------------------------
+def _get_config():
+    return get_rabbitmq_config()
+
+
+def _get_url():
+    cfg = _get_config()
+    return f"amqp://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}{cfg['vhost']}"
+
+
+def _queue_name(suffix):
+    return f"{_PREFIX}{suffix}"
+
+
+def _exchange_name(suffix):
+    return f"{_PREFIX}ex_{suffix}"
+
+
+# ---------------------------------------------------------------------------
+# Connection
+# ---------------------------------------------------------------------------
+async def _connect():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        return {"status": "connected", "is_closed": conn.is_closed}
+    finally:
+        await conn.close()
+
+
+def connect():
+    return asyncio.run(_connect())
+
+
+async def _connect_basic():
+    conn = await aio_pika.connect(_get_url())
+    try:
+        return {"status": "connected_basic", "is_closed": conn.is_closed}
+    finally:
+        await conn.close()
+
+
+def connect_basic():
+    return asyncio.run(_connect_basic())
+
+
+# ---------------------------------------------------------------------------
+# Connection error scenarios
+# ---------------------------------------------------------------------------
+async def _connect_error_wrong_host():
+    cfg = _get_config()
+    url = f"amqp://{cfg['user']}:{cfg['password']}@192.0.2.1:5672/"
+    try:
+        conn = await asyncio.wait_for(aio_pika.connect(url), timeout=3)
+        await conn.close()
+        return {"status": "unexpected_success"}
+    except Exception as e:
+        return {"status": "error", "msg": str(e)}
+
+
+def connect_error_wrong_host():
+    return asyncio.run(_connect_error_wrong_host())
+
+
+async def _connect_error_wrong_port():
+    cfg = _get_config()
+    url = f"amqp://{cfg['user']}:{cfg['password']}@{cfg['host']}:55672/"
+    try:
+        conn = await asyncio.wait_for(aio_pika.connect(url), timeout=3)
+        await conn.close()
+        return {"status": "unexpected_success"}
+    except Exception as e:
+        return {"status": "error", "msg": str(e)}
+
+
+def connect_error_wrong_port():
+    return asyncio.run(_connect_error_wrong_port())
+
+
+async def _connect_error_wrong_credentials():
+    cfg = _get_config()
+    url = f"amqp://bad_user:bad_pass@{cfg['host']}:{cfg['port']}/"
+    try:
+        conn = await asyncio.wait_for(aio_pika.connect(url), timeout=5)
+        await conn.close()
+        return {"status": "unexpected_success"}
+    except Exception as e:
+        return {"status": "error", "msg": str(e)}
+
+
+def connect_error_wrong_credentials():
+    return asyncio.run(_connect_error_wrong_credentials())
+
+
+# ---------------------------------------------------------------------------
+# Queue management
+# ---------------------------------------------------------------------------
+async def _declare_queue():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("basic"), durable=True)
+        return {
+            "status": "queue_declared",
+            "queue": q.name,
+            "message_count": q.declaration_result.message_count,
+            "consumer_count": q.declaration_result.consumer_count,
+        }
+    finally:
+        await conn.close()
+
+
+def declare_queue():
+    return asyncio.run(_declare_queue())
+
+
+async def _purge_queue():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("basic"), durable=True)
+        purged = await q.purge()
+        return {"status": "queue_purged", "messages_purged": purged}
+    finally:
+        await conn.close()
+
+
+def purge_queue():
+    return asyncio.run(_purge_queue())
+
+
+async def _delete_queue():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("delete_test"))
+        await q.delete()
+        return {"status": "queue_deleted"}
+    finally:
+        await conn.close()
+
+
+def delete_queue():
+    return asyncio.run(_delete_queue())
+
+
+# ---------------------------------------------------------------------------
+# Exchange management
+# ---------------------------------------------------------------------------
+async def _declare_exchange():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        ex = await channel.declare_exchange(
+            _exchange_name("direct"), aio_pika.ExchangeType.DIRECT, durable=True
+        )
+        return {"status": "exchange_declared", "exchange": ex.name}
+    finally:
+        await conn.close()
+
+
+def declare_exchange():
+    return asyncio.run(_declare_exchange())
+
+
+async def _bind_queue_to_exchange():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        ex = await channel.declare_exchange(
+            _exchange_name("direct"), aio_pika.ExchangeType.DIRECT, durable=True
+        )
+        q = await channel.declare_queue(_queue_name("bound"), durable=True)
+        await q.bind(ex, routing_key="test_key")
+        return {"status": "queue_bound", "queue": q.name, "exchange": ex.name, "routing_key": "test_key"}
+    finally:
+        await conn.close()
+
+
+def bind_queue_to_exchange():
+    return asyncio.run(_bind_queue_to_exchange())
+
+
+async def _delete_exchange():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        ex = await channel.declare_exchange(
+            _exchange_name("del_test"), aio_pika.ExchangeType.FANOUT
+        )
+        await ex.delete()
+        return {"status": "exchange_deleted"}
+    finally:
+        await conn.close()
+
+
+def delete_exchange():
+    return asyncio.run(_delete_exchange())
+
+
+# ---------------------------------------------------------------------------
+# Publishing
+# ---------------------------------------------------------------------------
+async def _publish_basic():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("pub_basic"), durable=True)
+        await q.purge()
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"hello aio_pika", delivery_mode=aio_pika.DeliveryMode.PERSISTENT),
+            routing_key=q.name,
+        )
+        return {"status": "published"}
+    finally:
+        await conn.close()
+
+
+def publish_basic():
+    return asyncio.run(_publish_basic())
+
+
+async def _publish_to_exchange():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        ex = await channel.declare_exchange(
+            _exchange_name("pub"), aio_pika.ExchangeType.DIRECT, durable=True
+        )
+        q = await channel.declare_queue(_queue_name("pub_ex"), durable=True)
+        await q.bind(ex, routing_key="pub_key")
+        await q.purge()
+        await ex.publish(
+            aio_pika.Message(body=b"exchange message", delivery_mode=aio_pika.DeliveryMode.PERSISTENT),
+            routing_key="pub_key",
+        )
+        return {"status": "published_to_exchange"}
+    finally:
+        await conn.close()
+
+
+def publish_to_exchange():
+    return asyncio.run(_publish_to_exchange())
+
+
+async def _publish_with_properties():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("pub_props"), durable=True)
+        await q.purge()
+        msg = aio_pika.Message(
+            body=json.dumps({"key": "value"}).encode(),
+            delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+            content_type="application/json",
+            headers={"x-custom": "test_header"},
+            priority=5,
+            expiration=60,
+        )
+        await channel.default_exchange.publish(msg, routing_key=q.name)
+        return {"status": "published_with_properties"}
+    finally:
+        await conn.close()
+
+
+def publish_with_properties():
+    return asyncio.run(_publish_with_properties())
+
+
+async def _publish_bulk():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("pub_bulk"), durable=True)
+        await q.purge()
+        count = 10
+        for i in range(count):
+            await channel.default_exchange.publish(
+                aio_pika.Message(body=f"bulk_msg_{i}".encode(), delivery_mode=aio_pika.DeliveryMode.PERSISTENT),
+                routing_key=q.name,
+            )
+        return {"status": "bulk_published", "count": count}
+    finally:
+        await conn.close()
+
+
+def publish_bulk():
+    return asyncio.run(_publish_bulk())
+
+
+async def _publish_json():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("pub_json"), durable=True)
+        await q.purge()
+        payload = {"name": "test", "values": [1, 2, 3], "nested": {"a": True}}
+        msg = aio_pika.Message(
+            body=json.dumps(payload).encode(),
+            content_type="application/json",
+            delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+        )
+        await channel.default_exchange.publish(msg, routing_key=q.name)
+        return {"status": "json_published"}
+    finally:
+        await conn.close()
+
+
+def publish_json():
+    return asyncio.run(_publish_json())
+
+
+# ---------------------------------------------------------------------------
+# Consuming
+# ---------------------------------------------------------------------------
+async def _consume_basic_get():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("pub_basic"), durable=True)
+        await q.purge()
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"consume_me"),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.ack()
+            return {"status": "consumed", "body": msg.body.decode()}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def consume_basic_get():
+    return asyncio.run(_consume_basic_get())
+
+
+async def _consume_with_ack():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("ack_test"), durable=True)
+        await q.purge()
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"ack_me"),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.ack()
+            return {"status": "consumed_acked", "body": msg.body.decode()}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def consume_with_ack():
+    return asyncio.run(_consume_with_ack())
+
+
+async def _consume_with_nack():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("nack_test"), durable=True)
+        await q.purge()
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"nack_me"),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.nack(requeue=False)
+            return {"status": "consumed_nacked", "body": msg.body.decode()}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def consume_with_nack():
+    return asyncio.run(_consume_with_nack())
+
+
+async def _consume_with_reject():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("reject_test"), durable=True)
+        await q.purge()
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"reject_me"),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.reject(requeue=False)
+            return {"status": "consumed_rejected", "body": msg.body.decode()}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def consume_with_reject():
+    return asyncio.run(_consume_with_reject())
+
+
+async def _consume_with_qos():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        await channel.set_qos(prefetch_count=2)
+        q = await channel.declare_queue(_queue_name("qos_test"), durable=True)
+        await q.purge()
+        for i in range(5):
+            await channel.default_exchange.publish(
+                aio_pika.Message(body=f"qos_{i}".encode()),
+                routing_key=q.name,
+            )
+        messages = []
+        for _ in range(2):
+            msg = await q.get(timeout=5)
+            if msg:
+                messages.append(msg.body.decode())
+                await msg.ack()
+        return {"status": "qos_consumed", "messages": messages, "count": len(messages)}
+    finally:
+        await conn.close()
+
+
+def consume_with_qos():
+    return asyncio.run(_consume_with_qos())
+
+
+async def _consume_json():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("pub_json"), durable=True)
+        await q.purge()
+        payload = {"type": "json_test", "data": [1, 2, 3]}
+        await channel.default_exchange.publish(
+            aio_pika.Message(
+                body=json.dumps(payload).encode(),
+                content_type="application/json",
+            ),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.ack()
+            parsed = json.loads(msg.body.decode())
+            return {"status": "json_consumed", "payload": parsed}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def consume_json():
+    return asyncio.run(_consume_json())
+
+
+# ---------------------------------------------------------------------------
+# Publisher confirms (enabled by default in aio_pika)
+# ---------------------------------------------------------------------------
+async def _publish_with_confirm():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("confirm_test"), durable=True)
+        await q.purge()
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"confirmed_msg", delivery_mode=aio_pika.DeliveryMode.PERSISTENT),
+            routing_key=q.name,
+        )
+        return {"status": "publish_confirmed"}
+    finally:
+        await conn.close()
+
+
+def publish_with_confirm():
+    return asyncio.run(_publish_with_confirm())
+
+
+# ---------------------------------------------------------------------------
+# Dead letter queue
+# ---------------------------------------------------------------------------
+async def _dead_letter_queue():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        dlx = await channel.declare_exchange(
+            _exchange_name("dlx"), aio_pika.ExchangeType.DIRECT, durable=True
+        )
+        dlq = await channel.declare_queue(_queue_name("dlq"), durable=True)
+        await dlq.bind(dlx, routing_key="dead")
+
+        main_q_name = _queue_name("dl_main")
+        try:
+            temp_q = await channel.declare_queue(main_q_name, passive=True)
+            await temp_q.delete()
+        except Exception:
+            pass
+
+        channel = await conn.channel()
+        main_q = await channel.declare_queue(
+            main_q_name,
+            durable=True,
+            arguments={
+                "x-dead-letter-exchange": _exchange_name("dlx"),
+                "x-dead-letter-routing-key": "dead",
+            },
+        )
+        await dlq.purge()
+
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"will_be_dead_lettered"),
+            routing_key=main_q.name,
+        )
+        msg = await main_q.get(timeout=5)
+        if msg:
+            await msg.nack(requeue=False)
+
+        await asyncio.sleep(0.5)
+
+        dl_msg = await dlq.get(timeout=5)
+        if dl_msg:
+            await dl_msg.ack()
+            return {"status": "dead_letter_ok", "body": dl_msg.body.decode()}
+        return {"status": "dead_letter_no_message"}
+    finally:
+        await conn.close()
+
+
+def dead_letter_queue():
+    return asyncio.run(_dead_letter_queue())
+
+
+# ---------------------------------------------------------------------------
+# Error scenarios
+# ---------------------------------------------------------------------------
+async def _error_consume_nonexistent_queue():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        try:
+            await channel.declare_queue("nonexistent_q_xyz_12345", passive=True)
+            return {"status": "unexpected_success"}
+        except Exception as e:
+            return {"status": "error", "msg": str(e)}
+    finally:
+        await conn.close()
+
+
+def error_consume_nonexistent_queue():
+    return asyncio.run(_error_consume_nonexistent_queue())
+
+
+# ---------------------------------------------------------------------------
+# Unique / advanced
+# ---------------------------------------------------------------------------
+async def _delayed_publish():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("delayed"), durable=True)
+        await q.purge()
+        await asyncio.sleep(1)
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=b"delayed_msg"),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.ack()
+            return {"status": "delayed_publish_ok", "body": msg.body.decode()}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def delayed_publish():
+    return asyncio.run(_delayed_publish())
+
+
+async def _publish_binary():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        q = await channel.declare_queue(_queue_name("binary"), durable=True)
+        await q.purge()
+        binary_data = bytes(range(256))
+        await channel.default_exchange.publish(
+            aio_pika.Message(body=binary_data, content_type="application/octet-stream"),
+            routing_key=q.name,
+        )
+        msg = await q.get(timeout=5)
+        if msg:
+            await msg.ack()
+            return {"status": "binary_ok", "length": len(msg.body)}
+        return {"status": "no_message"}
+    finally:
+        await conn.close()
+
+
+def publish_binary():
+    return asyncio.run(_publish_binary())
+
+
+# ---------------------------------------------------------------------------
+# Cleanup helper
+# ---------------------------------------------------------------------------
+async def _cleanup():
+    conn = await aio_pika.connect_robust(_get_url())
+    try:
+        channel = await conn.channel()
+        queues = [
+            "basic", "pub_basic", "pub_ex", "pub_props", "pub_bulk", "pub_json",
+            "ack_test", "nack_test", "reject_test", "qos_test", "confirm_test",
+            "dl_main", "dlq", "delete_test", "bound", "delayed", "binary",
+        ]
+        for q_suffix in queues:
+            try:
+                q = await channel.declare_queue(_queue_name(q_suffix), passive=True)
+                await q.delete()
+            except Exception:
+                channel = await conn.channel()
+
+        exchanges = ["direct", "pub", "dlx", "del_test"]
+        for ex_suffix in exchanges:
+            try:
+                ex = await channel.declare_exchange(
+                    _exchange_name(ex_suffix), aio_pika.ExchangeType.DIRECT, passive=True
+                )
+                await ex.delete()
+            except Exception:
+                channel = await conn.channel()
+        return {"status": "cleanup_done"}
+    finally:
+        await conn.close()
+
+
+def cleanup():
+    return asyncio.run(_cleanup())
+
+
+# ---------------------------------------------------------------------------
+# Global route-to-function mapping
+# ---------------------------------------------------------------------------
+ROUTE_MAP = {
+    "/broker/rabbitmq/aio_pika/connect": connect,
+    "/broker/rabbitmq/aio_pika/connect_basic": connect_basic,
+    "/broker/rabbitmq/aio_pika/connect_error_wrong_host": connect_error_wrong_host,
+    "/broker/rabbitmq/aio_pika/connect_error_wrong_port": connect_error_wrong_port,
+    "/broker/rabbitmq/aio_pika/connect_error_wrong_credentials": connect_error_wrong_credentials,
+    "/broker/rabbitmq/aio_pika/declare_queue": declare_queue,
+    "/broker/rabbitmq/aio_pika/purge_queue": purge_queue,
+    "/broker/rabbitmq/aio_pika/delete_queue": delete_queue,
+    "/broker/rabbitmq/aio_pika/declare_exchange": declare_exchange,
+    "/broker/rabbitmq/aio_pika/bind_queue_to_exchange": bind_queue_to_exchange,
+    "/broker/rabbitmq/aio_pika/delete_exchange": delete_exchange,
+    "/broker/rabbitmq/aio_pika/publish_basic": publish_basic,
+    "/broker/rabbitmq/aio_pika/publish_to_exchange": publish_to_exchange,
+    "/broker/rabbitmq/aio_pika/publish_with_properties": publish_with_properties,
+    "/broker/rabbitmq/aio_pika/publish_bulk": publish_bulk,
+    "/broker/rabbitmq/aio_pika/publish_json": publish_json,
+    "/broker/rabbitmq/aio_pika/consume_basic_get": consume_basic_get,
+    "/broker/rabbitmq/aio_pika/consume_with_ack": consume_with_ack,
+    "/broker/rabbitmq/aio_pika/consume_with_nack": consume_with_nack,
+    "/broker/rabbitmq/aio_pika/consume_with_reject": consume_with_reject,
+    "/broker/rabbitmq/aio_pika/consume_with_qos": consume_with_qos,
+    "/broker/rabbitmq/aio_pika/consume_json": consume_json,
+    "/broker/rabbitmq/aio_pika/publish_with_confirm": publish_with_confirm,
+    "/broker/rabbitmq/aio_pika/dead_letter_queue": dead_letter_queue,
+    "/broker/rabbitmq/aio_pika/error_consume_nonexistent_queue": error_consume_nonexistent_queue,
+    "/broker/rabbitmq/aio_pika/delayed_publish": delayed_publish,
+    "/broker/rabbitmq/aio_pika/publish_binary": publish_binary,
+    "/broker/rabbitmq/aio_pika/cleanup": cleanup,
+}