rabbitmq-test/test.py

import asyncio
import sys
import os

# Add current directory to Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from product.direct_multi_publish import setup_multi_queue_balance, BalancedProducer
from product.direct_publish import setup_direct_exchange, direct_publish
from product.fanout_publish import fanout_publish, setup_fanout_exchange
from product.topic_publish import setup_topic_exchange, topic_publish
from comsumer.direct_consumer import start_all_direct_consumers
from comsumer.fanout_consumer import start_all_fanout_consumers
from comsumer.topic_consumer import start_all_topic_consumers
from comsumer.direct_multi_consumer import start_balanced_consumers


async def run_fanout_test():
    """Run fanout exchange test with producer and consumer"""
    print("=== Running Fanout Exchange Test ===")
    
    # Start consumer in background
    consumer_task = asyncio.create_task(start_all_fanout_consumers())
    
    # Wait for consumer to start
    await asyncio.sleep(1)
    
    # Setup and publish messages
    await setup_fanout_exchange("demo.fanout", "demo.fanout.queue-")
    await fanout_publish(message="hello world", exchange_name="demo.fanout")
    await fanout_publish(message="test message 2", exchange_name="demo.fanout")
    await fanout_publish(message="test message 3", exchange_name="demo.fanout")
    
    # Wait for messages to be processed
    await asyncio.sleep(3)
    
    # Cancel consumer
    consumer_task.cancel()
    print("✅ Fanout test completed successfully!")


async def run_direct_exchange_test():
    """Run direct exchange test with producer and consumer"""
    print("=== Running Direct Exchange Test ===")
    
    # Start consumer in background
    consumer_task = asyncio.create_task(start_all_direct_consumers())
    
    # Wait for consumer to start
    await asyncio.sleep(1)
    
    # Setup exchange and publish messages
    await setup_direct_exchange()
    
    test_messages = [
        ("System crash, unable to start", "error"),  # Route to error queue
        ("Disk space insufficient", "warning"),  # Route to warning queue
        ("User login successful", "info"),  # Route to info queue
        ("Debug info: Database connection successful", "debug")  # Route to info queue
    ]

    for msg, routing_key in test_messages:
        await direct_publish(msg, routing_key)
        await asyncio.sleep(0.5)
    
    # Wait for messages to be processed
    await asyncio.sleep(3)
    
    # Cancel consumer
    consumer_task.cancel()
    print("✅ Direct exchange test completed successfully!")


async def run_topic_exchange_test():
    """Run topic exchange test with producer and consumer"""
    print("=== Running Topic Exchange Test ===")
    
    # Start consumer in background
    consumer_task = asyncio.create_task(start_all_topic_consumers())
    
    # Wait for consumer to start
    await asyncio.sleep(1)
    
    # Setup exchange and publish messages
    await setup_topic_exchange()
    
    test_messages = [
        ("Order creation failed (critical error)", "order.create.critical"),
        ("User login successful", "user.login.success"),
        ("Order payment completed", "order.pay.success"),
        ("System crash (critical error)", "system.crash.critical"),
        ("User login failed", "user.login.failed"),
        ("Normal system log", "system.log.info")  # Won't match any binding key, will be discarded
    ]

    for msg, routing_key in test_messages:
        await topic_publish(msg, routing_key)
        await asyncio.sleep(0.5)
    
    # Wait for messages to be processed
    await asyncio.sleep(3)
    
    # Cancel consumer
    consumer_task.cancel()
    print("✅ Topic exchange test completed successfully!")


async def run_multi_queue_balance_test():
    """Run multi-queue load balancing test with producer and consumer"""
    print("=== Running Multi-Queue Load Balancing Test ===")
    
    queue_count = 3
    
    # Start consumer in background
    consumer_task = asyncio.create_task(start_balanced_consumers(queue_count=queue_count))
    
    # Wait for consumer to start
    await asyncio.sleep(1)
    
    # Setup and publish messages
    await setup_multi_queue_balance(queue_count=queue_count)
    
    producer = BalancedProducer(queue_count=queue_count)
    await producer.connect()
    
    for i in range(10):
        await producer.publish(f"Task {i + 1}: Multi-queue load balancing test")
        await asyncio.sleep(0.3)
    
    await producer.close()
    
    # Wait for messages to be processed
    await asyncio.sleep(3)
    
    # Cancel consumer
    consumer_task.cancel()
    print("✅ Multi-queue load balancing test completed successfully!")


async def main():
    """Main function to run all tests"""
    print("🚀 Starting RabbitMQ Tests...")
    
    try:
        # Run all tests
        await run_fanout_test()
        await asyncio.sleep(1)
        
        await run_direct_exchange_test()
        await asyncio.sleep(1)
        
        await run_topic_exchange_test()
        await asyncio.sleep(1)
        
        await run_multi_queue_balance_test()
        
        print("\n🎉 All tests completed successfully!")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())