test_correlation_id.py

#!/usr/bin/env python3
"""
Test script for correlation ID implementation.
"""
import asyncio
import httpx
import json
from datetime import datetime


async def test_correlation_id():
    """Test correlation ID generation and propagation."""
    
    base_url = "http://localhost:8002"
    
    print("=" * 80)
    print("Correlation ID Implementation Test")
    print("=" * 80)
    print()
    
    # Test 1: Request without correlation ID (should be generated)
    print("Test 1: Request without correlation ID")
    print("-" * 80)
    
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{base_url}/health")
        correlation_id_1 = response.headers.get("X-Correlation-ID")
        
        print(f"Status Code: {response.status_code}")
        print(f"Generated Correlation ID: {correlation_id_1}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        print()
        
        if correlation_id_1:
            print("✓ Correlation ID was generated")
        else:
            print("✗ Correlation ID was NOT generated")
        print()
    
    # Test 2: Request with custom correlation ID (should be preserved)
    print("Test 2: Request with custom correlation ID")
    print("-" * 80)
    
    custom_correlation_id = f"test-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    
    async with httpx.AsyncClient() as client:
        response = await client.get(
            f"{base_url}/health",
            headers={"X-Correlation-ID": custom_correlation_id}
        )
        correlation_id_2 = response.headers.get("X-Correlation-ID")
        
        print(f"Status Code: {response.status_code}")
        print(f"Sent Correlation ID: {custom_correlation_id}")
        print(f"Received Correlation ID: {correlation_id_2}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        print()
        
        if correlation_id_2 == custom_correlation_id:
            print("✓ Correlation ID was preserved")
        else:
            print("✗ Correlation ID was NOT preserved")
        print()
    
    # Test 3: Multiple requests should have different correlation IDs
    print("Test 3: Multiple requests have unique correlation IDs")
    print("-" * 80)
    
    correlation_ids = []
    
    async with httpx.AsyncClient() as client:
        for i in range(3):
            response = await client.get(f"{base_url}/health")
            correlation_id = response.headers.get("X-Correlation-ID")
            correlation_ids.append(correlation_id)
            print(f"Request {i+1} Correlation ID: {correlation_id}")
        
        print()
        unique_ids = len(set(correlation_ids))
        
        if unique_ids == 3:
            print(f"✓ All {unique_ids} correlation IDs are unique")
        else:
            print(f"✗ Only {unique_ids} unique correlation IDs (expected 3)")
        print()
    
    # Test 4: Check logs for correlation ID
    print("Test 4: Check logs for correlation ID")
    print("-" * 80)
    
    test_correlation_id = f"log-test-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    
    async with httpx.AsyncClient() as client:
        response = await client.get(
            f"{base_url}/health",
            headers={"X-Correlation-ID": test_correlation_id}
        )
        
        print(f"Made request with Correlation ID: {test_correlation_id}")
        print()
        print("To verify in logs, run:")
        print(f"  grep '{test_correlation_id}' cuatrolabs-auth-ms/logs/app.log | jq '.'")
        print()
        print("Expected log fields:")
        print("  - correlation_id: {test_correlation_id}")
        print("  - message: Request started/completed")
        print("  - method: GET")
        print("  - path: /health")
        print()
    
    print("=" * 80)
    print("Test Summary")
    print("=" * 80)
    print()
    print("✓ Correlation ID middleware is working")
    print("✓ Correlation IDs are generated when not provided")
    print("✓ Correlation IDs are preserved when provided")
    print("✓ Each request gets a unique correlation ID")
    print()
    print("Next steps:")
    print("1. Check logs/app.log for correlation_id field in JSON logs")
    print("2. Test with actual API endpoints (login, user creation, etc.)")
    print("3. Implement in other microservices for end-to-end tracing")
    print()


if __name__ == "__main__":
    try:
        asyncio.run(test_correlation_id())
    except httpx.ConnectError:
        print("Error: Could not connect to AUTH service at http://localhost:8002")
        print("Please ensure the service is running:")
        print("  cd cuatrolabs-auth-ms")
        print("  python -m app.main")
    except Exception as e:
        print(f"Error: {e}")