"""
Test priority-based rate limiting for dynamic_rate_limiter_v3.

Core tests to validate that priority weights are respected (0.9/0.1) instead of equal splitting (0.5/0.5).
"""

import asyncio
import os
import sys
import time
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, patch

import pytest

sys.path.insert(0, os.path.abspath("../../../.."))

import litellm
from litellm import DualCache, Router
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
    _PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler,
)


class TimeController:
    def __init__(self):
        self._current = datetime.utcnow()

    def now(self) -> datetime:
        return self._current

    def advance(self, seconds: float) -> None:
        self._current += timedelta(seconds=seconds)


@pytest.fixture
def time_controller(monkeypatch):
    controller = TimeController()
    monkeypatch.setattr(time, "time", lambda: controller.now().timestamp())
    return controller


@pytest.mark.asyncio
async def test_priority_weight_allocation():
    """
    Test that priority weights are correctly applied instead of equal splitting.

    With priority_reservation = {"high": 0.9, "low": 0.1}:
    - High priority should get 90% of TPM (900 out of 1000)
    - Low priority should get 10% of TPM (100 out of 1000)

    This validates the core fix where before it would split 50/50.
    """
    # Set up environment for premium feature
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    # Set up priority reservations
    litellm.priority_reservation = {"high": 0.9, "low": 0.1}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "test-model"
    total_tpm = 1000

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": total_tpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Test high priority allocation
    high_priority_user = UserAPIKeyAuth()
    high_priority_user.metadata = {"priority": "high"}

    high_descriptors = handler._create_priority_based_descriptors(
        model=model,
        user_api_key_dict=high_priority_user,
        priority="high",
    )

    assert len(high_descriptors) == 1
    high_descriptor = high_descriptors[0]
    expected_high_tpm = int(total_tpm * 0.9)  # 900
    actual_high_tpm = high_descriptor["rate_limit"]["tokens_per_unit"]

    assert (
        actual_high_tpm == expected_high_tpm
    ), f"High priority should get {expected_high_tpm} TPM (90%), got {actual_high_tpm}"
    assert high_descriptor["value"] == f"{model}:high"

    # Test low priority allocation
    low_priority_user = UserAPIKeyAuth()
    low_priority_user.metadata = {"priority": "low"}

    low_descriptors = handler._create_priority_based_descriptors(
        model=model,
        user_api_key_dict=low_priority_user,
        priority="low",
    )

    assert len(low_descriptors) == 1
    low_descriptor = low_descriptors[0]
    expected_low_tpm = int(total_tpm * 0.1)  # 100
    actual_low_tpm = low_descriptor["rate_limit"]["tokens_per_unit"]

    assert (
        actual_low_tpm == expected_low_tpm
    ), f"Low priority should get {expected_low_tpm} TPM (10%), got {actual_low_tpm}"
    assert low_descriptor["value"] == f"{model}:low"

    # Verify the ratio is 9:1, not 1:1 (equal splitting)
    ratio = actual_high_tpm / actual_low_tpm
    expected_ratio = 9.0
    assert (
        abs(ratio - expected_ratio) < 0.1
    ), f"High:Low ratio should be {expected_ratio}:1, got {ratio}:1"


@pytest.mark.asyncio
async def test_concurrent_priority_requests():
    """
    Test the core issue: 5 concurrent requests with different priorities should get
    proper allocation based on priority weights, not equal splitting.

    This tests the exact scenario mentioned: priorities 0.9 and 0.1 should be 0.9/0.1, not 0.5/0.5.
    """
    # Set up environment for premium feature
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    # Set up the exact scenario from the issue
    litellm.priority_reservation = {"high": 0.9, "low": 0.1}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "test-model"
    total_tpm = 1000

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": total_tpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create 5 concurrent users - 3 high priority, 2 low priority
    high_priority_users = []
    low_priority_users = []

    for i in range(3):  # 3 high priority users
        user = UserAPIKeyAuth()
        user.metadata = {"priority": "high"}
        user.user_id = f"high_user_{i}"
        high_priority_users.append(user)

    for i in range(2):  # 2 low priority users
        user = UserAPIKeyAuth()
        user.metadata = {"priority": "low"}
        user.user_id = f"low_user_{i}"
        low_priority_users.append(user)

    # Test all high priority users get the same allocation (not divided)
    for user in high_priority_users:
        descriptors = handler._create_priority_based_descriptors(
            model=model,
            user_api_key_dict=user,
            priority="high",
        )

        assert len(descriptors) == 1
        descriptor = descriptors[0]
        # Each high priority user should get 900 TPM, not divided by 3
        assert descriptor["rate_limit"]["tokens_per_unit"] == 900, (
            f"High priority user {user.user_id} should get 900 TPM, "
            f"got {descriptor['rate_limit']['tokens_per_unit']}"
        )
        assert descriptor["value"] == f"{model}:high"

    # Test all low priority users get the same allocation (not divided)
    for user in low_priority_users:
        descriptors = handler._create_priority_based_descriptors(
            model=model,
            user_api_key_dict=user,
            priority="low",
        )

        assert len(descriptors) == 1
        descriptor = descriptors[0]
        # Each low priority user should get 100 TPM, not divided by 2
        assert descriptor["rate_limit"]["tokens_per_unit"] == 100, (
            f"Low priority user {user.user_id} should get 100 TPM, "
            f"got {descriptor['rate_limit']['tokens_per_unit']}"
        )
        assert descriptor["value"] == f"{model}:low"


@pytest.mark.asyncio
async def test_100_concurrent_priority_requests(time_controller):
    """
    Stress test: 100 concurrent requests with mixed priorities over 10 seconds.

    This validates that the priority system works correctly under high load:
    - 70 high priority requests (should get 900 TPM each)
    - 30 low priority requests (should get 100 TPM each)
    - Spread across 10 seconds to simulate real-world load
    """
    # Set up environment for premium feature
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    # Set up priority reservations
    litellm.priority_reservation = {"high": 0.9, "low": 0.1}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(
        internal_usage_cache=dual_cache, time_provider=time_controller.now
    )

    model = "stress-test-model"
    total_tpm = 1000

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": total_tpm,
                    "rpm": 500,  # Also test RPM limits
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create 100 users: 70 high priority, 30 low priority
    all_users = []

    # 70 high priority users
    for i in range(70):
        user = UserAPIKeyAuth()
        user.metadata = {"priority": "high"}
        user.user_id = f"high_stress_user_{i}"
        all_users.append((user, "high", 900, 450))  # expected TPM, expected RPM

    # 30 low priority users
    for i in range(30):
        user = UserAPIKeyAuth()
        user.metadata = {"priority": "low"}
        user.user_id = f"low_stress_user_{i}"
        all_users.append((user, "low", 100, 50))  # expected TPM, expected RPM

    async def test_user_descriptors(user_data):
        """Test descriptor creation for a single user."""
        user, priority, expected_tpm, expected_rpm = user_data

        descriptors = handler._create_priority_based_descriptors(
            model=model,
            user_api_key_dict=user,
            priority=priority,
        )

        assert (
            len(descriptors) == 1
        ), f"User {user.user_id} should have exactly 1 descriptor"
        descriptor = descriptors[0]

        # Validate TPM allocation
        actual_tpm = descriptor["rate_limit"]["tokens_per_unit"]
        assert (
            actual_tpm == expected_tpm
        ), f"User {user.user_id} ({priority}) should get {expected_tpm} TPM, got {actual_tpm}"

        # Validate RPM allocation
        actual_rpm = descriptor["rate_limit"]["requests_per_unit"]
        assert (
            actual_rpm == expected_rpm
        ), f"User {user.user_id} ({priority}) should get {expected_rpm} RPM, got {actual_rpm}"

        # Validate descriptor key
        assert descriptor["value"] == f"{model}:{priority}"
        assert descriptor["key"] == "priority_model"

        return {
            "user_id": user.user_id,
            "priority": priority,
            "tpm": actual_tpm,
            "rpm": actual_rpm,
            "success": True,
        }

    # Run all 100 requests concurrently to simulate high load
    start_time = time.time()

    # Split into batches to simulate requests over 10 seconds
    batch_size = 10  # 10 requests per batch
    batches = [
        all_users[i : i + batch_size] for i in range(0, len(all_users), batch_size)
    ]

    all_results = []

    for batch_idx, batch in enumerate(batches):
        # Process each batch concurrently
        batch_tasks = [test_user_descriptors(user_data) for user_data in batch]
        batch_results = await asyncio.gather(*batch_tasks, return_exceptions=True)
        all_results.extend(batch_results)

        # Add small delay between batches to spread over ~10 seconds
        if batch_idx < len(batches) - 1:  # Don't sleep after last batch
            await asyncio.sleep(0)
            time_controller.advance(1.0)  # simulate 1s passing between batches

    end_time = time.time()
    total_duration = end_time - start_time

    # Validate that the test ran over approximately 10 seconds
    assert (
        total_duration >= 9.0
    ), f"Test should take ~10 seconds, took {total_duration:.2f}s"
    assert total_duration <= 15.0, f"Test took too long: {total_duration:.2f}s"

    # Validate all requests were successful
    successful_results = [
        r for r in all_results if isinstance(r, dict) and r.get("success")
    ]
    assert (
        len(successful_results) == 100
    ), f"Expected 100 successful results, got {len(successful_results)}"

    # Validate priority distribution
    high_priority_results = [r for r in successful_results if r["priority"] == "high"]
    low_priority_results = [r for r in successful_results if r["priority"] == "low"]

    assert (
        len(high_priority_results) == 70
    ), f"Expected 70 high priority results, got {len(high_priority_results)}"
    assert (
        len(low_priority_results) == 30
    ), f"Expected 30 low priority results, got {len(low_priority_results)}"

    # Validate all high priority users got correct allocation
    for result in high_priority_results:
        assert (
            result["tpm"] == 900
        ), f"High priority user {result['user_id']} got {result['tpm']} TPM, expected 900"
        assert (
            result["rpm"] == 450
        ), f"High priority user {result['user_id']} got {result['rpm']} RPM, expected 450"

    # Validate all low priority users got correct allocation
    for result in low_priority_results:
        assert (
            result["tpm"] == 100
        ), f"Low priority user {result['user_id']} got {result['tpm']} TPM, expected 100"
        assert (
            result["rpm"] == 50
        ), f"Low priority user {result['user_id']} got {result['rpm']} RPM, expected 50"

    print(f"✅ Successfully processed 100 concurrent requests in {total_duration:.2f}s")
    print(f"   - 70 high priority users: 900 TPM, 450 RPM each")
    print(f"   - 30 low priority users: 100 TPM, 50 RPM each")
    print(f"   - Priority ratio maintained: 9:1 (TPM) and 9:1 (RPM)")


@pytest.mark.asyncio
async def test_concurrent_pre_call_hooks_stress():
    """
    Stress test: 50 concurrent pre-call hooks with saturation-aware priority enforcement.

    Tests priority-based rate limiting in strict mode (>80% saturation).
    Mocks high saturation to force strict mode where priorities are enforced.
    Premium users (80% allocation) should have >90% success rate.
    Standard users (20% allocation) should have ~70% success rate with 30% random limiting.
    """
    # Set up environment for premium feature
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    litellm.priority_reservation = {"premium": 0.8, "standard": 0.2}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "pre-call-stress-model"
    total_tpm = 2000

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": total_tpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Mock the v3 limiter to simulate different scenarios
    successful_requests = []
    rate_limited_requests = []

    # Mock saturation check to return high saturation (forces strict mode)
    async def mock_get_cache(key, litellm_parent_otel_span=None, local_only=False):
        """Mock cache to simulate high saturation."""
        # Return high usage to trigger strict mode (>80% saturation)
        if ":requests" in key or ":tokens" in key:
            return 1800  # 1800/2000 = 90% saturation
        return None

    async def mock_should_rate_limit(
        descriptors, parent_otel_span=None, read_only=False
    ):
        """Mock rate limiter that handles saturation-aware descriptors."""
        descriptor = descriptors[0]
        descriptor_key = descriptor["key"]
        descriptor_value = descriptor["value"]

        # Handle model-wide tracking (for both generous and strict mode tracking)
        if descriptor_key == "model_saturation_check":
            # Always allow model-wide tracking (doesn't enforce in our mock)
            return {
                "overall_code": "OK",
                "statuses": [
                    {
                        "code": "OK",
                        "descriptor_key": descriptor_value,
                        "rate_limit_type": "tokens_per_unit",
                        "limit_remaining": 10000,
                    }
                ],
            }

        # Handle priority-specific enforcement in strict mode
        elif descriptor_key == "priority_model":
            # Extract priority from value like "pre-call-stress-model:premium"
            priority = descriptor_value.split(":")[-1]

            if priority == "premium":
                # Allow all premium requests
                return {
                    "overall_code": "OK",
                    "statuses": [
                        {
                            "code": "OK",
                            "descriptor_key": descriptor_value,
                            "rate_limit_type": "tokens_per_unit",
                            "limit_remaining": 1000,
                        }
                    ],
                }
            else:
                # Rate limit some standard requests (simulate load)
                import random

                if random.random() < 0.3:  # 30% of standard requests get rate limited
                    return {
                        "overall_code": "OVER_LIMIT",
                        "statuses": [
                            {
                                "code": "OVER_LIMIT",
                                "descriptor_key": descriptor_value,
                                "rate_limit_type": "tokens_per_unit",
                                "limit_remaining": 0,
                            }
                        ],
                    }
                else:
                    return {
                        "overall_code": "OK",
                        "statuses": [
                            {
                                "code": "OK",
                                "descriptor_key": descriptor_value,
                                "rate_limit_type": "tokens_per_unit",
                                "limit_remaining": 100,
                            }
                        ],
                    }

        # Default: allow
        return {
            "overall_code": "OK",
            "statuses": [
                {
                    "code": "OK",
                    "descriptor_key": descriptor_value,
                    "rate_limit_type": "tokens_per_unit",
                    "limit_remaining": 1000,
                }
            ],
        }

    # Create 50 users: 30 premium, 20 standard
    users = []

    for i in range(30):
        user = UserAPIKeyAuth()
        user.metadata = {"priority": "premium"}
        user.user_id = f"premium_hook_user_{i}"
        users.append((user, "premium"))

    for i in range(20):
        user = UserAPIKeyAuth()
        user.metadata = {"priority": "standard"}
        user.user_id = f"standard_hook_user_{i}"
        users.append((user, "standard"))

    async def make_request(user_data):
        """Make a pre-call hook request."""
        user, priority = user_data

        try:
            result = await handler.async_pre_call_hook(
                user_api_key_dict=user,
                cache=DualCache(),
                data={"model": model},
                call_type="completion",
            )

            # If no exception, request was allowed
            successful_requests.append(
                {"user_id": user.user_id, "priority": priority, "result": "allowed"}
            )
            return {
                "status": "success",
                "user_id": user.user_id,
                "priority": priority,
            }

        except Exception as e:
            # Request was rate limited
            rate_limited_requests.append(
                {"user_id": user.user_id, "priority": priority, "error": str(e)}
            )
            return {
                "status": "rate_limited",
                "user_id": user.user_id,
                "priority": priority,
            }

    # Run all 50 requests concurrently with patches applied to the entire batch
    start_time = time.time()
    with (
        patch.object(
            handler.v3_limiter, "should_rate_limit", side_effect=mock_should_rate_limit
        ),
        patch.object(
            handler.internal_usage_cache, "async_get_cache", side_effect=mock_get_cache
        ),
    ):
        tasks = [make_request(user_data) for user_data in users]
        results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.time()

    # Analyze results
    successful_count = len(
        [r for r in results if isinstance(r, dict) and r["status"] == "success"]
    )
    rate_limited_count = len(
        [r for r in results if isinstance(r, dict) and r["status"] == "rate_limited"]
    )

    # Validate that premium users were mostly successful (priority worked)
    premium_results = [
        r for r in results if isinstance(r, dict) and r["priority"] == "premium"
    ]
    premium_success = len([r for r in premium_results if r["status"] == "success"])

    standard_results = [
        r for r in results if isinstance(r, dict) and r["priority"] == "standard"
    ]
    standard_success = len([r for r in standard_results if r["status"] == "success"])

    # Premium users should have higher success rate due to priority
    premium_success_rate = (
        premium_success / len(premium_results) if premium_results else 0
    )
    standard_success_rate = (
        standard_success / len(standard_results) if standard_results else 0
    )

    assert (
        premium_success_rate >= 0.9
    ), f"Premium success rate should be >= 90%, got {premium_success_rate:.2%}"
    assert (
        standard_success_rate >= 0.5
    ), f"Standard success rate should be >= 50% (with 30% random limiting, allows for variance), got {standard_success_rate:.2%}"

    # Allow for the case where both are 100% due to timing/mocking issues
    # The test is inherently flaky due to random behavior
    if premium_success_rate < 1.0 or standard_success_rate < 1.0:
        assert (
            premium_success_rate >= standard_success_rate
        ), "Premium should have >= success rate than standard"

    total_duration = end_time - start_time

    print(f"✅ Processed 50 concurrent pre-call hooks in {total_duration:.2f}s")
    print(
        f"   - Premium users: {premium_success}/{len(premium_results)} success ({premium_success_rate:.1%})"
    )
    print(
        f"   - Standard users: {standard_success}/{len(standard_results)} success ({standard_success_rate:.1%})"
    )
    print(f"   - Total successful: {successful_count}/50 ({successful_count/50:.1%})")
    print(f"   - Priority system working: Premium > Standard success rates")


# These tests make actual async_pre_call_hook calls to simulate real traffic


@pytest.mark.asyncio
async def test_fake_calls_case_1_no_rate_limiting_at_capacity():
    """
    Test Case 1: Saturation-Aware Rate Limiting at 50% Threshold

    System: 100 RPM capacity, saturation_threshold=50%
    Key A: priority_reservation=0.75 (75 RPM reserved)
    Key B: priority_reservation=0.25 (25 RPM reserved)
    Traffic A: 1 request
    Traffic B: 100 requests

    Expected behavior:
    - Key A: 1 request succeeds (low traffic)
    - Key B: ~25-26 requests succeed (capped at reservation when saturation >= 50%)

    Once saturation hits 50%, strict mode enforces priority-based limits.
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    # Set up priority reservations
    litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "fake-call-test-1"
    total_rpm = 100

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create users
    key_a_user = UserAPIKeyAuth()
    key_a_user.metadata = {"priority": "key_a"}
    key_a_user.user_id = "key_a_user"

    key_b_user = UserAPIKeyAuth()
    key_b_user.metadata = {"priority": "key_b"}
    key_b_user.user_id = "key_b_user"

    # Track results
    successful_requests = {"key_a": 0, "key_b": 0}
    rate_limited_requests = {"key_a": 0, "key_b": 0}

    async def make_request(user, priority_name, request_id):
        """Make a single request and track the result."""
        try:
            result = await handler.async_pre_call_hook(
                user_api_key_dict=user,
                cache=dual_cache,
                data={"model": model},
                call_type="completion",
            )

            if result is None:
                successful_requests[priority_name] += 1
                return {"status": "success", "priority": priority_name}
            else:
                rate_limited_requests[priority_name] += 1
                return {"status": "rate_limited", "priority": priority_name}

        except Exception as e:
            rate_limited_requests[priority_name] += 1
            return {
                "status": "rate_limited",
                "priority": priority_name,
                "error": str(e),
            }

    # Send 1 request from key_a, 100 from key_b
    tasks = []

    for i in range(1):
        tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}"))

    for i in range(100):
        tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}"))

    start_time = time.time()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.time()

    # Analyze results
    total_successful = successful_requests["key_a"] + successful_requests["key_b"]
    total_rate_limited = rate_limited_requests["key_a"] + rate_limited_requests["key_b"]

    print(f"Test Case 1 - Saturation-Aware Rate Limiting:")
    print(f"   - Duration: {end_time - start_time:.2f}s")
    print(f"   - Key A: {successful_requests['key_a']}/1 successful (reserved 75 RPM)")
    print(
        f"   - Key B: {successful_requests['key_b']}/100 successful (reserved 25 RPM)"
    )
    print(f"   - Total successful: {total_successful}/101")
    print(f"   - Total rate limited: {total_rate_limited}/101")

    # Key A should get its 1 request
    assert (
        successful_requests["key_a"] == 1
    ), f"Key A should get 1 request, got {successful_requests['key_a']}"

    # Key B can send until saturation hits 50% (which is ~50 total requests)
    # After that, strict mode enforces its 25 RPM reservation
    # Due to race conditions in concurrent execution, allow 45-52 successful requests
    assert (
        45 <= successful_requests["key_b"] <= 52
    ), f"Key B should get ~49 requests (45-52), got {successful_requests['key_b']}"

    # Verify approximately half of key_b requests were rate limited
    assert (
        rate_limited_requests["key_b"] >= 45
    ), f"Key B should have ≥45 rate limited requests, got {rate_limited_requests['key_b']}"


@pytest.mark.asyncio
async def test_fake_calls_case_2_priority_queue_during_saturation():
    """
    Test Case 2: Priority Queue Behavior During Saturation

    System: 100 RPM capacity
    Key A: priority_reservation=0.75 (75 RPM reserved)
    Key B: priority_reservation=0.25 (25 RPM reserved)
    Traffic A: 200 RPM
    Traffic B: 200 RPM
    Expected A: 75 RPM (75% of capacity)
    Expected B: 25 RPM (25% of capacity)

    When total traffic exceeds capacity, rate limiting enforces priority reservations.
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "fake-call-test-2"
    total_rpm = 100

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create users
    key_a_user = UserAPIKeyAuth()
    key_a_user.metadata = {"priority": "key_a"}
    key_a_user.user_id = "key_a_user"

    key_b_user = UserAPIKeyAuth()
    key_b_user.metadata = {"priority": "key_b"}
    key_b_user.user_id = "key_b_user"

    # Track results
    successful_requests = {"key_a": 0, "key_b": 0}
    rate_limited_requests = {"key_a": 0, "key_b": 0}

    async def make_request(user, priority_name, request_id):
        """Make a single request and track the result."""
        try:
            result = await handler.async_pre_call_hook(
                user_api_key_dict=user,
                cache=dual_cache,
                data={"model": model},
                call_type="completion",
            )

            if result is None:
                successful_requests[priority_name] += 1
                return {"status": "success", "priority": priority_name}
            else:
                rate_limited_requests[priority_name] += 1
                return {"status": "rate_limited", "priority": priority_name}

        except Exception as e:
            rate_limited_requests[priority_name] += 1
            return {
                "status": "rate_limited",
                "priority": priority_name,
                "error": str(e),
            }

    # Send 200 requests from each priority (over capacity)
    tasks = []

    for i in range(200):
        tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}"))

    for i in range(200):
        tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}"))

    start_time = time.time()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.time()

    # Analyze results
    total_successful = successful_requests["key_a"] + successful_requests["key_b"]

    key_a_success_rate = successful_requests["key_a"] / 200
    key_b_success_rate = successful_requests["key_b"] / 200

    print(f"Test Case 2 - Priority Queue Behavior During Saturation:")
    print(f"   - Duration: {end_time - start_time:.2f}s")
    print(
        f"   - Key A: {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})"
    )
    print(
        f"   - Key B: {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})"
    )
    print(f"   - Total successful: {total_successful}/400")

    # Key A should get significantly more requests than Key B (75:25 ratio)
    assert (
        key_a_success_rate > key_b_success_rate
    ), f"Key A should have higher success rate: {key_a_success_rate:.1%} vs {key_b_success_rate:.1%}"

    # Check ratio is approximately 3:1 (75:25)
    if total_successful > 0:
        key_a_share = successful_requests["key_a"] / total_successful
        expected_key_a_share = 0.75

        print(
            f"   - Key A got {key_a_share:.1%} of successful requests (expected ~75%)"
        )

        # Allow tolerance for timing effects
        assert (
            abs(key_a_share - expected_key_a_share) < 0.2
        ), f"Key A share should be ~75%, got {key_a_share:.1%}"


@pytest.mark.asyncio
async def test_fake_calls_case_3_spillover_capacity_default_keys():
    """
    Test Case 3: Spillover Capacity for Default Keys

    System: 100 RPM capacity
    Key A: priority_reservation=0.75 (75 RPM reserved)
    Key B: nothing set (default)
    Key C: nothing set (default)
    Key D: nothing set (default)
    Traffic A: 150 RPM
    Traffic B: 150 RPM
    Traffic C: 150 RPM
    Traffic D: 150 RPM
    Expected A: 75 RPM (75% reserved)
    Expected B: ~8.3 RPM (remaining 25 RPM / 3 default keys)
    Expected C: ~8.3 RPM
    Expected D: ~8.3 RPM

    Tests spillover behavior where default keys share remaining capacity.
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    litellm.priority_reservation = {"key_a": 0.75}
    litellm.priority_reservation_settings.default_priority = 0.25

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "fake-call-test-3"
    total_rpm = 100

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create users
    key_a_user = UserAPIKeyAuth()
    key_a_user.metadata = {"priority": "key_a"}
    key_a_user.user_id = "key_a_user"

    key_b_user = UserAPIKeyAuth()
    key_b_user.metadata = {}
    key_b_user.user_id = "key_b_user"

    key_c_user = UserAPIKeyAuth()
    key_c_user.metadata = {}
    key_c_user.user_id = "key_c_user"

    key_d_user = UserAPIKeyAuth()
    key_d_user.metadata = {}
    key_d_user.user_id = "key_d_user"

    # Track results
    successful_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0}
    rate_limited_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0}

    async def make_request(user, key_name, request_id):
        """Make a single request and track the result."""
        try:
            result = await handler.async_pre_call_hook(
                user_api_key_dict=user,
                cache=dual_cache,
                data={"model": model},
                call_type="completion",
            )

            if result is None:
                successful_requests[key_name] += 1
                return {"status": "success", "key": key_name}
            else:
                rate_limited_requests[key_name] += 1
                return {"status": "rate_limited", "key": key_name}

        except Exception as e:
            rate_limited_requests[key_name] += 1
            return {"status": "rate_limited", "key": key_name, "error": str(e)}

    # Send 150 requests from each key (600 total, 6x over capacity)
    tasks = []

    for i in range(150):
        tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}"))

    for i in range(150):
        tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}"))

    for i in range(150):
        tasks.append(make_request(key_c_user, "key_c", f"key_c_{i}"))

    for i in range(150):
        tasks.append(make_request(key_d_user, "key_d", f"key_d_{i}"))

    start_time = time.time()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.time()

    # Analyze results
    total_successful = sum(successful_requests.values())

    print(f"Test Case 3 - Spillover Capacity for Default Keys:")
    print(f"   - Duration: {end_time - start_time:.2f}s")
    print(f"   - Key A: {successful_requests['key_a']}/150 successful")
    print(f"   - Key B: {successful_requests['key_b']}/150 successful (default)")
    print(f"   - Key C: {successful_requests['key_c']}/150 successful (default)")
    print(f"   - Key D: {successful_requests['key_d']}/150 successful (default)")
    print(f"   - Total successful: {total_successful}/600")

    # Key A should get the most requests (75% of capacity)
    assert (
        successful_requests["key_a"] > successful_requests["key_b"]
    ), "Key A should get more than Key B"
    assert (
        successful_requests["key_a"] > successful_requests["key_c"]
    ), "Key A should get more than Key C"
    assert (
        successful_requests["key_a"] > successful_requests["key_d"]
    ), "Key A should get more than Key D"

    # Default keys should get similar amounts (spillover capacity)
    avg_default = (
        successful_requests["key_b"]
        + successful_requests["key_c"]
        + successful_requests["key_d"]
    ) / 3
    print(f"   - Average default key success: {avg_default:.1f}")


@pytest.mark.asyncio
async def test_fake_calls_case_4_over_allocated_with_normalization():
    """
    Test Case 4: Over-Allocated Priority reservations with Normalization

    System: 100 RPM capacity
    Key A: priority_reservation=0.60 (60% requested)
    Key B: priority_reservation=0.80 (80% requested)
    Total: 140% (over-allocated, should normalize to 43%/57%)
    Traffic A: 200 RPM
    Traffic B: 200 RPM

    With saturation-aware rate limiting:
    - Initially, requests are allowed through in generous mode (under 80% saturation)
    - Once saturated, strict priority-based limits kick in with normalized weights
    - Due to concurrent burst, total successful may exceed 100 RPM in the test window
    - This test verifies normalization works and total capacity is reasonably bounded
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    litellm.priority_reservation = {"key_a": 0.60, "key_b": 0.80}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "fake-call-test-4"
    total_rpm = 100

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create users
    key_a_user = UserAPIKeyAuth()
    key_a_user.metadata = {"priority": "key_a"}
    key_a_user.user_id = "key_a_user"

    key_b_user = UserAPIKeyAuth()
    key_b_user.metadata = {"priority": "key_b"}
    key_b_user.user_id = "key_b_user"

    # Track results
    successful_requests = {"key_a": 0, "key_b": 0}
    rate_limited_requests = {"key_a": 0, "key_b": 0}

    async def make_request(user, priority_name, request_id):
        """Make a single request and track the result."""
        try:
            result = await handler.async_pre_call_hook(
                user_api_key_dict=user,
                cache=dual_cache,
                data={"model": model},
                call_type="completion",
            )

            if result is None:
                successful_requests[priority_name] += 1
                return {"status": "success", "priority": priority_name}
            else:
                rate_limited_requests[priority_name] += 1
                return {"status": "rate_limited", "priority": priority_name}

        except Exception as e:
            rate_limited_requests[priority_name] += 1
            return {
                "status": "rate_limited",
                "priority": priority_name,
                "error": str(e),
            }

    # Send 200 requests from each key (400 total, 4x over capacity)
    tasks = []

    for i in range(200):
        tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}"))

    for i in range(200):
        tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}"))

    start_time = time.time()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.time()

    # Analyze results
    total_successful = successful_requests["key_a"] + successful_requests["key_b"]

    key_a_success_rate = successful_requests["key_a"] / 200
    key_b_success_rate = successful_requests["key_b"] / 200

    print(f"Test Case 4 - Over-Allocated Priority Reservations with Normalization:")
    print(f"   - Duration: {end_time - start_time:.2f}s")
    print(
        f"   - Key A (0.60): {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})"
    )
    print(
        f"   - Key B (0.80): {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})"
    )
    print(f"   - Total successful: {total_successful}/400")

    # With saturation-aware behavior:
    # 1. Verify total capacity is reasonably bounded (not all 400 requests succeed)
    assert (
        total_successful < 300
    ), f"Total requests should be bounded by saturation detection, got {total_successful}/400"

    # 2. Verify significant rate limiting occurred (at least 50% blocked)
    assert (
        total_successful < 200
    ), f"At least 50% of requests should be rate limited, got {total_successful}/400 successful"

    # 3. Verify both keys got some requests through (normalization is working)
    assert successful_requests["key_a"] > 0, "Key A should get some requests"
    assert successful_requests["key_b"] > 0, "Key B should get some requests"

    print(
        f"   - Normalization test PASSED: Both priorities got requests, "
        f"total bounded to {total_successful} (under 200)"
    )


@pytest.mark.asyncio
async def test_fake_calls_case_5_default_value_priority_reservation():
    """
    Test Case 5: Default value for priority reservation

    System: 100 RPM capacity
    Key A: priority_reservation=0.50 (50 RPM)
    Key B: priority_reservation=0.20 (20 RPM)
    Key C: priority_reservation=0.05 (5 RPM)
    Key D: nothing set (uses default_priority=0.05, 5 RPM)
    Traffic A: 150 RPM
    Traffic B: 150 RPM
    Traffic C: 150 RPM
    Traffic D: 150 RPM
    Expected A: 55 RPM (normalized)
    Expected B: 25 RPM (normalized)
    Expected C: 10 RPM (normalized)
    Expected D: 10 RPM (normalized)

    Tests complex scenario with explicit priorities and default priority.
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    litellm.priority_reservation = {"key_a": 0.50, "key_b": 0.20, "key_c": 0.05}
    litellm.priority_reservation_settings.default_priority = 0.05

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "fake-call-test-5"
    total_rpm = 100

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create users
    key_a_user = UserAPIKeyAuth()
    key_a_user.metadata = {"priority": "key_a"}
    key_a_user.user_id = "key_a_user"

    key_b_user = UserAPIKeyAuth()
    key_b_user.metadata = {"priority": "key_b"}
    key_b_user.user_id = "key_b_user"

    key_c_user = UserAPIKeyAuth()
    key_c_user.metadata = {"priority": "key_c"}
    key_c_user.user_id = "key_c_user"

    key_d_user = UserAPIKeyAuth()
    key_d_user.metadata = {}
    key_d_user.user_id = "key_d_user"

    # Track results
    successful_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0}
    rate_limited_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0}

    async def make_request(user, key_name, request_id):
        """Make a single request and track the result."""
        try:
            result = await handler.async_pre_call_hook(
                user_api_key_dict=user,
                cache=dual_cache,
                data={"model": model},
                call_type="completion",
            )

            if result is None:
                successful_requests[key_name] += 1
                return {"status": "success", "key": key_name}
            else:
                rate_limited_requests[key_name] += 1
                return {"status": "rate_limited", "key": key_name}

        except Exception as e:
            rate_limited_requests[key_name] += 1
            return {"status": "rate_limited", "key": key_name, "error": str(e)}

    # Send 150 requests from each key (600 total, 6x over capacity)
    tasks = []

    for i in range(150):
        tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}"))

    for i in range(150):
        tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}"))

    for i in range(150):
        tasks.append(make_request(key_c_user, "key_c", f"key_c_{i}"))

    for i in range(150):
        tasks.append(make_request(key_d_user, "key_d", f"key_d_{i}"))

    start_time = time.time()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    end_time = time.time()

    # Analyze results
    total_successful = sum(successful_requests.values())

    print(f"Test Case 5 - Default value for priority reservation:")
    print(f"   - Duration: {end_time - start_time:.2f}s")
    print(f"   - Key A (0.50): {successful_requests['key_a']}/150 successful")
    print(f"   - Key B (0.20): {successful_requests['key_b']}/150 successful")
    print(f"   - Key C (0.05): {successful_requests['key_c']}/150 successful")
    print(f"   - Key D (default 0.05): {successful_requests['key_d']}/150 successful")
    print(f"   - Total successful: {total_successful}/600")

    # Verify priority ordering: A > B > C ≈ D
    assert (
        successful_requests["key_a"] > successful_requests["key_b"]
    ), "Key A should get more than Key B"
    assert (
        successful_requests["key_b"] > successful_requests["key_c"]
    ), "Key B should get more than Key C"

    # Key C and Key D should get similar amounts (both have 0.05 priority)
    key_c_vs_d_ratio = successful_requests["key_c"] / max(
        successful_requests["key_d"], 1
    )
    print(f"   - Key C vs Key D ratio: {key_c_vs_d_ratio:.2f} (expected ~1.0)")

    if total_successful > 0:
        key_a_share = successful_requests["key_a"] / total_successful
        print(
            f"   - Key A got {key_a_share:.1%} of successful requests (expected ~55-62%)"
        )


@pytest.mark.asyncio
async def test_default_priority_shared_pool():
    """
    Test that keys without explicit priority share ONE default pool, not get individual allocations.

    With default_priority=0.25:
    - Key A, B, C (no priority) should share ONE 25 RPM pool
    - NOT get 25 RPM each (which would be 75 RPM total)
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    litellm.priority_reservation = {"prod": 0.75}
    litellm.priority_reservation_settings.default_priority = 0.25

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "test-default-pool"
    total_rpm = 100

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Create 3 users without explicit priority
    user_a = UserAPIKeyAuth()
    user_a.metadata = {}
    user_a.user_id = "user_a"

    user_b = UserAPIKeyAuth()
    user_b.metadata = {}
    user_b.user_id = "user_b"

    user_c = UserAPIKeyAuth()
    user_c.metadata = {}
    user_c.user_id = "user_c"

    # Get descriptors for each
    desc_a = handler._create_priority_based_descriptors(
        model=model, user_api_key_dict=user_a, priority=None
    )
    desc_b = handler._create_priority_based_descriptors(
        model=model, user_api_key_dict=user_b, priority=None
    )
    desc_c = handler._create_priority_based_descriptors(
        model=model, user_api_key_dict=user_c, priority=None
    )

    # All should use the SAME shared pool key
    assert desc_a[0]["value"] == f"{model}:default_pool"
    assert desc_b[0]["value"] == f"{model}:default_pool"
    assert desc_c[0]["value"] == f"{model}:default_pool"

    # All should have same limit (25 RPM SHARED, not 25 RPM each)
    assert desc_a[0]["rate_limit"]["requests_per_unit"] == 25
    assert desc_b[0]["rate_limit"]["requests_per_unit"] == 25
    assert desc_c[0]["rate_limit"]["requests_per_unit"] == 25

    # Verify explicit priority uses different pool
    user_prod = UserAPIKeyAuth()
    user_prod.metadata = {"priority": "prod"}
    desc_prod = handler._create_priority_based_descriptors(
        model=model, user_api_key_dict=user_prod, priority="prod"
    )

    assert desc_prod[0]["value"] == f"{model}:prod"
    assert desc_prod[0]["rate_limit"]["requests_per_unit"] == 75
    assert desc_prod[0]["value"] != desc_a[0]["value"]  # Different pools

    print("✅ Default priority test passed:")
    print(f"   - 3 keys without priority share ONE pool: {desc_a[0]['value']}")
    print(f"   - Shared pool limit: {desc_a[0]['rate_limit']['requests_per_unit']} RPM")
    print(f"   - Explicit priority 'prod' uses separate pool: {desc_prod[0]['value']}")


@pytest.mark.asyncio
async def test_async_log_success_event_increments_by_actual_tokens():
    """
    Test that async_log_success_event increments token counters by actual token usage.

    This validates the fix for Bug 1: Token count was incrementing by 1 instead of actual usage.
    The async_log_success_event should increment both model_saturation_check and priority_model
    counters by the actual completion_tokens (when rate_limit_type=output).
    """
    from unittest.mock import MagicMock

    from litellm.types.utils import ModelResponse, Usage

    os.environ["LITELLM_LICENSE"] = "test-license-key"
    litellm.priority_reservation = {"dev": 0.1, "prod": 0.9}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "test-token-increment"
    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": 1000,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Track what gets incremented
    increment_calls = []

    async def mock_increment(pipeline_operations, parent_otel_span=None):
        for op in pipeline_operations:
            increment_calls.append(
                {
                    "key": op["key"],
                    "increment_value": op["increment_value"],
                }
            )

    handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment

    # Create mock response with 50 completion tokens
    mock_response = MagicMock(spec=ModelResponse)
    mock_response.usage = MagicMock(spec=Usage)
    mock_response.usage.prompt_tokens = 10
    mock_response.usage.completion_tokens = 50
    mock_response.usage.total_tokens = 60

    # Create kwargs with priority in user_api_key_auth_metadata
    kwargs = {
        "standard_logging_object": {
            "metadata": {
                "user_api_key_auth_metadata": {"priority": "dev"},
            },
            "model_group": model,
        },
        "litellm_params": {
            "metadata": {"model_group": model},
        },
    }

    with patch(
        "litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs",
        return_value=model,
    ):
        await handler.async_log_success_event(
            kwargs=kwargs,
            response_obj=mock_response,
            start_time=None,
            end_time=None,
        )

    # Verify increments happened with actual token count (60 total tokens)
    assert (
        len(increment_calls) == 2
    ), f"Expected 2 increment calls, got {len(increment_calls)}"

    # Both should increment by 50 (total_tokens, since rate_limit_type defaults to 'total')
    for call in increment_calls:
        assert (
            call["increment_value"] == 60
        ), f"Expected increment of 60 tokens, got {call['increment_value']} for key {call['key']}"

    # Verify correct keys were used
    keys = [call["key"] for call in increment_calls]
    assert any(
        "model_saturation_check" in k for k in keys
    ), "Should increment model_saturation_check"
    assert any(
        "priority_model" in k and "dev" in k for k in keys
    ), "Should increment priority_model with 'dev' priority"


@pytest.mark.asyncio
async def test_saturation_check_cache_ttl_configuration():
    """
    Test that saturation_check_cache_ttl controls how long saturation values are cached locally.

    This validates the configurable TTL for multi-node consistency:
    - When saturation_check_cache_ttl is set, local cache should expire after that duration
    - After expiration, fresh values should be fetched from Redis
    - This prevents nodes from having stale saturation data in multi-node deployments
    """
    os.environ["LITELLM_LICENSE"] = "test-license-key"

    # Set a short TTL for testing (5 seconds)
    original_ttl = litellm.priority_reservation_settings.saturation_check_cache_ttl
    litellm.priority_reservation_settings.saturation_check_cache_ttl = 5

    try:
        dual_cache = DualCache()
        handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

        model = "test-saturation-ttl"
        llm_router = Router(
            model_list=[
                {
                    "model_name": model,
                    "litellm_params": {
                        "model": "gpt-3.5-turbo",
                        "api_key": "test-key",
                        "api_base": "test-base",
                        "rpm": 100,
                        "tpm": 1000,
                    },
                }
            ]
        )
        handler.update_variables(llm_router=llm_router)

        # Verify the TTL getter returns configured value
        assert (
            handler._get_saturation_check_cache_ttl() == 5
        ), "TTL should be configurable via priority_reservation_settings"

        # Track async_get_cache calls to verify TTL is passed
        get_cache_calls = []
        original_get_cache = handler.internal_usage_cache.async_get_cache

        async def mock_get_cache(
            key, litellm_parent_otel_span=None, local_only=False, **kwargs
        ):
            get_cache_calls.append(
                {
                    "key": key,
                    "ttl": kwargs.get("ttl"),
                    "local_only": local_only,
                }
            )
            return None  # Simulate cache miss

        handler.internal_usage_cache.async_get_cache = mock_get_cache

        # Call _get_saturation_value_from_cache
        counter_key = handler.v3_limiter.create_rate_limit_keys(
            key="model_saturation_check",
            value=model,
            rate_limit_type="requests",
        )

        await handler._get_saturation_value_from_cache(counter_key=counter_key)

        # Verify async_get_cache was called with the configured TTL
        assert len(get_cache_calls) == 1, "Expected 1 cache call"
        assert (
            get_cache_calls[0]["ttl"] == 5
        ), f"Expected TTL of 5 seconds, got {get_cache_calls[0]['ttl']}"
        assert (
            get_cache_calls[0]["local_only"] is False
        ), "Should check Redis (local_only=False) for multi-node consistency"

        # Test with different TTL value
        get_cache_calls.clear()
        litellm.priority_reservation_settings.saturation_check_cache_ttl = 30

        await handler._get_saturation_value_from_cache(counter_key=counter_key)

        assert (
            get_cache_calls[0]["ttl"] == 30
        ), f"TTL should update to 30 seconds, got {get_cache_calls[0]['ttl']}"

        print("Saturation check cache TTL test passed:")
        print(
            "   - TTL is configurable via priority_reservation_settings.saturation_check_cache_ttl"
        )
        print(
            "   - TTL is passed to async_get_cache for local cache expiration control"
        )
        print(
            "   - local_only=False ensures Redis is checked for multi-node consistency"
        )

    finally:
        # Restore original TTL
        litellm.priority_reservation_settings.saturation_check_cache_ttl = original_ttl


@pytest.mark.asyncio
async def test_async_log_success_event_uses_team_priority_from_auth_metadata():
    """
    Test that async_log_success_event correctly retrieves priority from user_api_key_auth_metadata.

    This validates the fix where priority is retrieved from standard_logging_metadata.user_api_key_auth_metadata
    instead of just standard_logging_metadata.priority. This is important for team-based priority inheritance.
    """
    from unittest.mock import MagicMock

    from litellm.types.utils import ModelResponse, Usage

    os.environ["LITELLM_LICENSE"] = "test-license-key"
    litellm.priority_reservation = {"team_priority": 0.8, "default": 0.2}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "test-team-priority"
    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": 1000,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    # Track incremented keys to verify priority is used correctly
    incremented_keys = []

    async def mock_increment(pipeline_operations, parent_otel_span=None):
        for op in pipeline_operations:
            incremented_keys.append(op["key"])

    handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment

    # Create mock response
    mock_response = MagicMock(spec=ModelResponse)
    mock_response.usage = MagicMock(spec=Usage)
    mock_response.usage.prompt_tokens = 10
    mock_response.usage.completion_tokens = 20
    mock_response.usage.total_tokens = 30

    # Simulate team metadata inheritance: priority is in user_api_key_auth_metadata
    # This is how the proxy passes team metadata to the callback
    kwargs = {
        "standard_logging_object": {
            "metadata": {
                # Priority NOT at top level (this would fail before the fix)
                # Priority IS in user_api_key_auth_metadata (team inheritance)
                "user_api_key_auth_metadata": {"priority": "team_priority"},
            },
            "model_group": model,
        },
        "litellm_params": {
            "metadata": {"model_group": model},
        },
    }

    with patch(
        "litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs",
        return_value=model,
    ):
        await handler.async_log_success_event(
            kwargs=kwargs,
            response_obj=mock_response,
            start_time=None,
            end_time=None,
        )

    # Verify the priority_model key uses 'team_priority' (not 'default_pool')
    priority_keys = [k for k in incremented_keys if "priority_model" in k]
    assert (
        len(priority_keys) == 1
    ), f"Expected 1 priority_model key, got {len(priority_keys)}"

    # The key should contain 'team_priority', not 'default_pool'
    assert "team_priority" in priority_keys[0], (
        f"Expected priority key to use 'team_priority' from user_api_key_auth_metadata, "
        f"got key: {priority_keys[0]}"
    )
    assert (
        "default_pool" not in priority_keys[0]
    ), f"Priority key should NOT use 'default_pool', should use team's priority. Got: {priority_keys[0]}"


@pytest.mark.asyncio
async def test_priority_429_includes_model_name_and_configured_limits():
    """
    The priority-based 429 should tell operators which model was hit and what
    the model's configured TPM/RPM are, so they can decide whether to tune the
    priority allocation or the model limits.

    Regression test for the previous message that read:
        "Priority-based rate limit exceeded. Priority: prod,
         Rate limit type: tokens, Remaining: -664145,
         Model saturation: 86.3%"
    -- with no indication of which model was hit.
    """
    from fastapi import HTTPException

    os.environ["LITELLM_LICENSE"] = "test-license-key"
    litellm.priority_reservation = {"prod": 0.5}

    dual_cache = DualCache()
    handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)

    model = "gpt-4o-test"
    total_tpm = 1_000_000
    total_rpm = 10_000

    llm_router = Router(
        model_list=[
            {
                "model_name": model,
                "litellm_params": {
                    "model": "gpt-3.5-turbo",
                    "api_key": "test-key",
                    "api_base": "test-base",
                    "tpm": total_tpm,
                    "rpm": total_rpm,
                },
            }
        ]
    )
    handler.update_variables(llm_router=llm_router)

    user = UserAPIKeyAuth()
    user.metadata = {"priority": "prod"}
    user.user_id = "prod_user"

    model_group_info = handler.llm_router.get_model_group_info(model_group=model)

    # Force the atomic check+increment to return OVER_LIMIT for the
    # priority_model descriptor. saturation=0.95 keeps us above the
    # default saturation threshold so priority limits are enforced.
    over_limit_response = {
        "overall_code": "OVER_LIMIT",
        "statuses": [
            {
                "code": "OVER_LIMIT",
                "descriptor_key": "priority_model",
                "rate_limit_type": "tokens",
                "limit_remaining": -664145,
                "current_limit": int(total_tpm * 0.5),
            }
        ],
    }

    with patch.object(
        handler.v3_limiter,
        "atomic_check_and_increment_by_n",
        new=AsyncMock(return_value=over_limit_response),
    ):
        with pytest.raises(HTTPException) as exc_info:
            await handler._check_rate_limits(
                model=model,
                model_group_info=model_group_info,
                user_api_key_dict=user,
                priority="prod",
                saturation=0.95,
                data={"model": model},
            )

    assert exc_info.value.status_code == 429
    detail = exc_info.value.detail
    assert isinstance(detail, dict)
    error_msg = detail["error"]

    # New fields added by this change -- the whole point of the fix.
    assert f"Model: {model}" in error_msg, error_msg
    assert f"Model TPM: {total_tpm}" in error_msg, error_msg
    assert f"Model RPM: {total_rpm}" in error_msg, error_msg

    # Existing fields must still be present (no regression).
    assert "Priority-based rate limit exceeded" in error_msg, error_msg
    assert "Priority: prod" in error_msg, error_msg
    assert "Rate limit type: tokens" in error_msg, error_msg
    assert "Model saturation:" in error_msg, error_msg
