python-sdk/tests/test_rate_limiter.py

#!/usr/bin/env python3
"""Test rate limiting and error recovery."""

import sys
import asyncio
from pathlib import Path

import pytest
from rich.table import Table
from rich.console import Console

# Add hanzo src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "pkg" / "hanzo" / "src"))

from hanzo.rate_limiter import (
    RateLimiter,
    ErrorRecovery,
    RateLimitConfig,
    smart_limiter,
)


async def test_rate_limiter():
    """Test rate limiter functionality."""
    console = Console()

    console.print("\n[bold cyan]Testing Rate Limiter[/bold cyan]\n")

    # Create a strict rate limiter for testing
    config = RateLimitConfig(requests_per_minute=5, requests_per_hour=100, burst_size=2)
    limiter = RateLimiter(config)

    # Test rapid requests
    console.print("[bold]Testing rate limiting (5 requests/minute):[/bold]")

    for i in range(8):
        allowed, wait = await limiter.check_rate_limit()

        if allowed:
            await limiter.acquire()
            console.print(f"✅ Request {i + 1} allowed")
        else:
            console.print(f"⏳ Request {i + 1} blocked, wait {wait:.1f}s")

        # Small delay between requests
        await asyncio.sleep(0.1)

    # Test status
    console.print("\n[bold]Rate limiter status:[/bold]")
    status = limiter.get_status()

    table = Table(title="Rate Limiter Status")
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="white")

    for key, value in status.items():
        table.add_row(key.replace("_", " ").title(), str(value))

    console.print(table)

    # Test error recovery
    console.print("\n[bold]Testing error recovery:[/bold]")

    recovery = ErrorRecovery(limiter)

    # Simulate a function that fails initially
    attempt_count = 0

    async def flaky_function():
        nonlocal attempt_count
        attempt_count += 1

        if attempt_count < 3:
            console.print(f"  Attempt {attempt_count} failed")
            raise Exception("Simulated error")

        console.print(f"  Attempt {attempt_count} succeeded!")
        return "Success"

    try:
        result = await recovery.with_retry(flaky_function, max_retries=5)
        console.print(f"[green]Result: {result}[/green]")
    except Exception as e:
        console.print(f"[red]Failed after retries: {e}[/red]")

    # Test smart limiter
    console.print("\n[bold]Testing smart rate limiter:[/bold]")

    # Simulate API calls
    async def mock_api_call(api_type: str):
        console.print(f"  Calling {api_type} API...")
        await asyncio.sleep(0.1)
        return f"Response from {api_type}"

    # Test different API types
    for api_type in ["openai", "anthropic", "local", "free"]:
        try:
            result = await smart_limiter.execute_with_limit(api_type, mock_api_call, api_type)
            console.print(f"  ✅ {api_type}: {result}")
        except Exception as e:
            console.print(f"  ❌ {api_type}: {e}")

    # Show all limiter statuses
    console.print("\n[bold]All API limiter statuses:[/bold]")
    all_status = smart_limiter.get_all_status()

    for api_type, status in all_status.items():
        if status["total_requests"] > 0:
            console.print(f"\n{api_type}:")
            console.print(f"  Requests: {status['total_requests']}")
            console.print(f"  Errors: {status['total_errors']}")
            console.print(f"  Last minute: {status['requests_last_minute']}")

    console.print("\n[green]✅ Rate limiting tests completed![/green]")

    return True


if __name__ == "__main__":
    success = asyncio.run(test_rate_limiter())
    sys.exit(0 if success else 1)