tests/test_rate_tracker.py

"""
Unit tests for backend/utils/rate_tracker.py
"""
import asyncio
import time
from utils.rate_tracker import rate_tracker, record_call, can_call, get_available_providers, get_status


async def test_rate_record_and_limits():
    print("\n[TEST] Rate Tracker: record and limit enforcement")
    # Use a test provider with small limit
    provider = "test_provider"
    rate_tracker.limits[provider] = 5
    rate_tracker.windows[provider] = type(rate_tracker.windows.get("groq"))(5)

    # Initially should be callable
    allowed = await can_call(provider)
    print(f"✓ Initially can_call: {allowed}")
    assert allowed

    # Record 5 calls
    for i in range(5):
        await record_call(provider)
    
    allowed_after = await can_call(provider)
    print(f"✓ After {5} calls can_call: {allowed_after}")
    assert not allowed_after, "Provider should be exhausted after reaching limit"

    # Get status
    status = await get_status()
    print(f"✓ Status for {provider}: {status[provider]}")
    assert status[provider]["count_last_min"] == 5

    # Sleep past window to allow pruning
    print("Waiting for window to expire (2s)...")
    time.sleep(2)
    # Manually prune by calling count()
    cnt = await rate_tracker.windows[provider].count()
    print(f"✓ Count after short wait: {cnt} (may still be >0)")

    print("\n✅ Rate Tracker basic enforcement test done")
    return True


async def test_available_providers_ordering():
    print("\n[TEST] Rate Tracker: provider ordering by available capacity")
    # Reset windows for predictable state
    for name in ["google", "groq", "cerebras", "openrouter"]:
        rate_tracker.windows[name] = type(rate_tracker.windows.get(name))(rate_tracker.limits[name])

    # Simulate some calls to google and groq
    for _ in range(10):
        await record_call("google")
    for _ in range(2):
        await record_call("groq")

    avail = await get_available_providers()
    print("✓ Availability snapshot:")
    for name, cap in avail:
        print(f"  - {name}: {cap}")

    # Ensure groq has more capacity than google now
    cap_map = {n: c for n, c in avail}
    assert cap_map["groq"] > cap_map["google"]

    print("\n✅ Provider ordering test done")
    return True


async def main():
    await test_rate_record_and_limits()
    await test_available_providers_ordering()


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