"""Unit tests for adaptive rate control."""

import pytest
import asyncio
import time
from azure_discovery.utils.rate_control import AdaptiveRateController, RateLimitConfig


class TestAdaptiveRateController:
    """Tests for adaptive rate controller."""
    
    @pytest.mark.asyncio
    async def test_acquire_enforces_rate_limit(self):
        """Test that acquire() enforces the rate limit."""
        config = RateLimitConfig(initial_rps=10.0)
        controller = AdaptiveRateController(config)
        
        start = time.time()
        
        # Make 5 requests
        for _ in range(5):
            await controller.acquire()
        
        elapsed = time.time() - start
        
        # Should take at least 0.4 seconds (5 requests at 10 RPS)
        assert elapsed >= 0.4
    
    @pytest.mark.asyncio
    async def test_record_success_increases_rate(self):
        """Test that sustained success increases rate."""
        config = RateLimitConfig(
            initial_rps=10.0,
            max_rps=50.0,
            increase_factor=2.0,
        )
        controller = AdaptiveRateController(config)
        
        initial_rps = controller.get_current_rps()
        
        # Record 10 successes (threshold for increase)
        for _ in range(10):
            await controller.record_success()
        
        new_rps = controller.get_current_rps()
        assert new_rps > initial_rps
        assert new_rps == initial_rps * 2.0
    
    @pytest.mark.asyncio
    async def test_record_throttle_decreases_rate(self):
        """Test that throttling decreases rate."""
        config = RateLimitConfig(
            initial_rps=10.0,
            min_rps=1.0,
            decrease_factor=0.5,
        )
        controller = AdaptiveRateController(config)
        
        initial_rps = controller.get_current_rps()
        
        await controller.record_throttle()
        
        new_rps = controller.get_current_rps()
        assert new_rps < initial_rps
        assert new_rps == initial_rps * 0.5
    
    @pytest.mark.asyncio
    async def test_respects_min_max_bounds(self):
        """Test that rate stays within configured bounds."""
        config = RateLimitConfig(
            initial_rps=10.0,
            min_rps=5.0,
            max_rps=20.0,
            increase_factor=3.0,
            decrease_factor=0.1,
        )
        controller = AdaptiveRateController(config)
        
        # Try to increase beyond max
        for _ in range(10):
            await controller.record_success()
        
        assert controller.get_current_rps() <= config.max_rps
        
        # Try to decrease below min
        for _ in range(10):
            await controller.record_throttle()
        
        assert controller.get_current_rps() >= config.min_rps
    
    def test_get_backoff_delay(self):
        """Test exponential backoff calculation."""
        config = RateLimitConfig(backoff_base=2.0)
        controller = AdaptiveRateController(config)
        
        assert controller.get_backoff_delay(0) == 1.0  # 2^0
        assert controller.get_backoff_delay(1) == 2.0  # 2^1
        assert controller.get_backoff_delay(2) == 4.0  # 2^2
        assert controller.get_backoff_delay(3) == 8.0  # 2^3
        
        # Should cap at 60 seconds
        assert controller.get_backoff_delay(10) == 60.0
    
    @pytest.mark.asyncio
    async def test_reset(self):
        """Test controller reset."""
        config = RateLimitConfig(initial_rps=10.0)
        controller = AdaptiveRateController(config)
        
        # Modify state
        for _ in range(10):
            await controller.record_success()
        
        modified_rps = controller.get_current_rps()
        assert modified_rps != config.initial_rps
        
        # Reset
        await controller.reset()
        
        reset_rps = controller.get_current_rps()
        assert reset_rps == config.initial_rps
