litellm/tests/unit/proxy/auth/test_auth_checks.py

# What is this?
## Tests if 'get_end_user_object' works as expected

import sys, os, asyncio, time, random, uuid
import hashlib
import json
import traceback
from dotenv import load_dotenv

load_dotenv()

from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from copy import deepcopy
from types import SimpleNamespace
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch

import pytest, litellm
import httpx
from fastapi import Request, status
from prisma import Prisma
from litellm._service_logger import ServiceTypes
from litellm.proxy._types import LiteLLM_OrganizationTable, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import get_org_object, get_user_object
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.caching.caching import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy._types import (
    LiteLLM_EndUserTable,
    LiteLLM_BudgetTable,
    LiteLLM_UserTable,
    LiteLLM_TeamTable,
    Litellm_EntityType,
    LitellmUserRoles,
    ProxyException,
)
from litellm.proxy.utils import PrismaClient
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.auth.auth_checks import (
    can_team_access_model,
    delete_cache_key_objects,
    is_model_cost_zero,
    virtual_key_soft_budget_check,
    _team_soft_budget_check,
    common_checks,
)
from litellm.proxy.utils import InternalUsageCache, ProxyLogging
from litellm.proxy.utils import CallInfo
from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER


async def _run_team_trace_id_check(
    *,
    route: str,
    method: str = "POST",
    request_body: dict[str, object] | None = None,
    headers: Mapping[str, str] | None = None,
    team_metadata: Mapping[str, object] | None = None,
    team_object: LiteLLM_TeamTable | None = None,
    registered_pass_through_endpoint: bool = False,
) -> bool:
    request_payload: Final[dict[str, object]] = {} if request_body is None else request_body
    raw_headers: Final = tuple(
        (name.lower().encode("latin-1"), value.encode("latin-1"))
        for name, value in (() if headers is None else headers.items())
    )
    request: Final = Request(
        {
            "type": "http",
            "method": method,
            "path": route,
            "headers": raw_headers,
            "query_string": b"",
        }
    )
    if registered_pass_through_endpoint:
        endpoint: Final = SimpleNamespace()
        setattr(endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, True)
        request.scope["endpoint"] = endpoint
    team: Final = (
        team_object
        if team_object is not None
        else LiteLLM_TeamTable(
            team_id="trace-id-team",
            metadata=None if team_metadata is None else dict(team_metadata),
        )
    )
    valid_token: Final = UserAPIKeyAuth(token="test-token", team_id=team.team_id)
    user: Final = LiteLLM_UserTable(user_id="trace-id-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
    return await common_checks(
        request_body=request_payload,
        team_object=team,
        user_object=user,
        end_user_object=None,
        global_proxy_spend=None,
        general_settings={},
        route=route,
        llm_router=None,
        proxy_logging_obj=MagicMock(spec=ProxyLogging),
        valid_token=valid_token,
        request=request,
        skip_budget_checks=True,
    )


@pytest.mark.parametrize(
    "route",
    ["/chat/completions", "/v1/messages", "/mcp/", "/a2a/my-agent"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_llm_mcp_and_agent_requests(route: str) -> None:
    request_body: Final[dict[str, object]] = {"metadata": {"source": "unit-test"}}
    original_body: Final = deepcopy(request_body)
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route=route,
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"
    assert exc_info.value.message == (
        "Team 'trace-id-team' requires a trace ID on every LLM, MCP and agent request. "
        "Send an x-litellm-trace-id header. "
        "LLM requests can also send a W3C traceparent header or set metadata.trace_id in the body."
    )
    assert request_body == original_body


@pytest.mark.parametrize(
    ("headers", "request_body"),
    [
        ({"x-litellm-trace-id": "trace-123"}, {}),
        ({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}, {}),
        ({}, {"metadata": {"trace_id": "trace-123"}}),
        ({}, {"litellm_metadata": {"trace_id": "trace-123"}}),
        ({}, {"litellm_trace_id": "trace-123"}),
        ({}, {"litellm_trace_id": "", "metadata": {"trace_id": "trace-123"}}),
        ({}, {"litellm_trace_id": None, "metadata": {"trace_id": "trace-123"}}),
    ],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_supported_trace_carriers(
    headers: Mapping[str, str],
    request_body: dict[str, object],
) -> None:
    original_body: Final = deepcopy(request_body)
    await _run_team_trace_id_check(
        route="/chat/completions",
        headers=headers,
        request_body=request_body,
        team_metadata={"require_trace_id": True},
    )
    assert request_body == original_body


@pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"])
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None:
    assert await _run_team_trace_id_check(
        route=route,
        method="POST",
        request_body={},
        team_metadata={"require_trace_id": True},
    ) is True


@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat() -> None:
    request_body: Final[dict[str, object]] = {
        "metadata": {"trace_id": ""},
        "litellm_metadata": {"trace_id": "trace-123"},
    }
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route="/chat/completions",
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"


@pytest.mark.parametrize(
    "request_body",
    [
        {"litellm_metadata": {"trace_id": "trace-123"}},
        {"metadata": {"trace_id": "trace-123"}},
    ],
    ids=["selected-litellm-metadata", "promoted-requester-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route(
    request_body: dict[str, object],
) -> None:
    assert await _run_team_trace_id_check(
        route="/v1/responses",
        request_body=request_body,
        team_metadata={"require_trace_id": True},
    ) is True


@pytest.mark.parametrize(
    ("headers", "request_body"),
    [
        (
            {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
            {"litellm_trace_id": ""},
        ),
        (
            {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
            {"litellm_trace_id": None},
        ),
    ],
    ids=[
        "traceparent-with-empty-top-level",
        "traceparent-with-null-top-level",
    ],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trace_id_is_explicit(
    headers: Mapping[str, str],
    request_body: dict[str, object],
) -> None:
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route="/chat/completions",
            headers=headers,
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"


@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None:
    assert await _run_team_trace_id_check(
        route="/mcp/",
        request_body={},
        headers={"x-litellm-trace-id": "trace-mcp-123"},
        team_metadata={"require_trace_id": True},
    ) is True


@pytest.mark.parametrize(
    ("headers", "request_body"),
    [
        (
            {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
            {},
        ),
        ({}, {"metadata": {"trace_id": "trace-body-123"}}),
    ],
    ids=["traceparent", "body-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_routes(
    headers: Mapping[str, str],
    request_body: dict[str, object],
) -> None:
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route="/anthropic/v1/messages",
            headers=headers,
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"


@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None:
    assert await _run_team_trace_id_check(
        route="/anthropic/v1/messages",
        request_body={},
        headers={"x-litellm-trace-id": "trace-provider-123"},
        team_metadata={"require_trace_id": True},
    ) is True


@pytest.mark.parametrize(
    ("headers", "request_body"),
    [
        (
            {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
            {},
        ),
        ({}, {"metadata": {"trace_id": "trace-body-123"}}),
    ],
    ids=["traceparent", "body-metadata"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pass_through_routes(
    headers: Mapping[str, str],
    request_body: dict[str, object],
) -> None:
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route="/configured-passthrough/infer",
            headers=headers,
            request_body=request_body,
            team_metadata={"require_trace_id": True},
            registered_pass_through_endpoint=True,
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"


@pytest.mark.parametrize(
    ("route", "request_body"),
    [
        ("/mcp/", {}),
        (
            "/mcp-rest/tools/call",
            {"server_id": "math", "name": "add", "arguments": {"left": 2, "right": 3}},
        ),
        (
            "/a2a/my-agent",
            {
                "jsonrpc": "2.0",
                "id": "1",
                "method": "message/send",
                "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
            },
        ),
    ],
    ids=["streamable-mcp", "mcp-rest", "a2a"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes(
    route: str,
    request_body: dict[str, object],
) -> None:
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route=route,
            headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"},
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"


@pytest.mark.parametrize(
    ("route", "request_body"),
    [
        (
            "/mcp-rest/tools/call",
            {
                "server_id": "math",
                "name": "add",
                "arguments": {"left": 2, "right": 3},
                "metadata": {"trace_id": "trace-body-123"},
            },
        ),
        (
            "/a2a/my-agent",
            {
                "jsonrpc": "2.0",
                "id": "1",
                "method": "message/send",
                "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
                "metadata": {"trace_id": "trace-body-123"},
            },
        ),
    ],
    ids=["mcp-rest", "a2a"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_body_metadata_on_mcp_and_agent_routes(
    route: str,
    request_body: dict[str, object],
) -> None:
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route=route,
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"


@pytest.mark.parametrize(
    ("headers", "request_body"),
    [
        ({"traceparent": "invalid"}, {}),
        ({}, {"metadata": {"trace_id": ""}}),
    ],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_invalid_trace_carriers(
    headers: Mapping[str, str],
    request_body: dict[str, object],
) -> None:
    original_body: Final = deepcopy(request_body)
    exc_info: Final[pytest.ExceptionInfo[ProxyException]]
    with pytest.raises(ProxyException) as exc_info:
        await _run_team_trace_id_check(
            route="/chat/completions",
            headers=headers,
            request_body=request_body,
            team_metadata={"require_trace_id": True},
        )
    assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST)
    assert exc_info.value.param == "trace_id"
    assert request_body == original_body


@pytest.mark.parametrize(
    ("route", "method", "team_metadata"),
    [
        ("/v1/models", "GET", {"require_trace_id": True}),
        ("/chat/completions", "POST", None),
        ("/chat/completions", "POST", {"require_trace_id": False}),
        ("/key/generate", "POST", {"require_trace_id": True}),
    ],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_exempts_methods_routes_and_disabled_teams(
    route: str,
    method: str,
    team_metadata: Mapping[str, object] | None,
) -> None:
    assert await _run_team_trace_id_check(route=route, method=method, team_metadata=team_metadata) is True


@pytest.mark.asyncio
async def test_team_require_trace_id_skips_non_dict_team_metadata() -> None:
    team: Final[LiteLLM_TeamTable] = LiteLLM_TeamTable.model_construct(
        team_id="trace-id-team",
        metadata="not a mapping",
    )
    assert await _run_team_trace_id_check(route="/v1/chat/completions", team_object=team) is True


@pytest.mark.parametrize("route", ("/lens/datasets", "/lens/tracing/keys"))
def test_lens_management_does_not_require_an_inference_user_parameter(route: str) -> None:
    from starlette.requests import Request

    from litellm.proxy.auth.auth_checks import _enforce_user_param_check

    request: Final = Request({"type": "http", "method": "POST", "path": route})
    _enforce_user_param_check({"enforce_user_param": True}, request, {}, route)
    with pytest.raises(Exception, match="'user' param not passed"):
        _enforce_user_param_check({"enforce_user_param": True}, request, {}, "/v1/chat/completions")


@pytest.mark.parametrize("customer_spend, customer_budget", [(0, 10), (10, 0)])
@pytest.mark.asyncio
async def test_get_end_user_object(customer_spend, customer_budget):
    """
    Scenario 1: normal - get_end_user_object returns the cached user
    Scenario 2: user over budget - NOTE: budget enforcement now happens in 
                common_checks() via _check_end_user_budget(), not in get_end_user_object()

    This test verifies that get_end_user_object correctly retrieves the end user
    from cache. Budget enforcement is tested separately in test_check_end_user_budget().
    """
    end_user_id = "my-test-customer"
    _budget = LiteLLM_BudgetTable(max_budget=customer_budget)
    end_user_obj = LiteLLM_EndUserTable(
        user_id=end_user_id,
        spend=customer_spend,
        litellm_budget_table=_budget,
        blocked=False,
    )
    # UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts
    # and breaks get_end_user_object's typed async_get_cache path.
    _cache = UserApiKeyCache()
    _key = "end_user_id:{}".format(end_user_id)
    await _cache.async_set_cache(
        key=_key,
        value=end_user_obj,
        model_type=LiteLLM_EndUserTable,
    )
    # get_end_user_object only fetches data - it no longer enforces budget
    # Budget enforcement happens in common_checks() via _check_end_user_budget()
    result = await get_end_user_object(
        end_user_id=end_user_id,
        prisma_client="RANDOM VALUE",  # type: ignore
        user_api_key_cache=_cache,
        route="/v1/chat/completions",
    )
    assert result is not None
    assert result.user_id == end_user_id


@pytest.mark.parametrize("customer_spend, customer_budget", [(0, 10), (10, 0)])
@pytest.mark.asyncio
async def test_check_end_user_budget(customer_spend, customer_budget):
    """
    Test _check_end_user_budget enforcement:
    - Scenario 1: customer_spend=0, customer_budget=10 - should pass (under budget)
    - Scenario 2: customer_spend=10, customer_budget=0 - should fail (over budget)

    Note: Budget enforcement for end users happens in common_checks() via
    _check_end_user_budget(), not in get_end_user_object().
    """
    from litellm.proxy.auth.auth_checks import check_end_user_budget

    _budget = LiteLLM_BudgetTable(max_budget=customer_budget)
    end_user_obj = LiteLLM_EndUserTable(
        user_id="my-test-customer",
        spend=customer_spend,
        litellm_budget_table=_budget,
        blocked=False,
    )

    should_exceed = customer_spend > customer_budget

    if not should_exceed:
        await check_end_user_budget(
            end_user_obj=end_user_obj,
            route="/v1/chat/completions",
        )
        return

    with pytest.raises(litellm.BudgetExceededError) as exc_info:
        await check_end_user_budget(
            end_user_obj=end_user_obj,
            route="/v1/chat/completions",
        )
    # Verify the error has correct info
    assert exc_info.value.current_cost == customer_spend
    assert exc_info.value.max_budget == customer_budget


@pytest.mark.parametrize(
    "model, expect_to_work",
    [
        ("openai/gpt-4o-mini", True),
        ("openai/gpt-4o", False),
    ],
)
@pytest.mark.asyncio
async def test_can_key_call_model(model, expect_to_work):
    """
    If wildcard model + specific model is used, choose the specific model settings
    """
    from litellm.proxy.auth.auth_checks import can_key_call_model
    from fastapi import HTTPException

    llm_model_list = [
        {
            "model_name": "openai/*",
            "litellm_params": {
                "model": "openai/*",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
                "db_model": False,
                "access_groups": ["public-openai-models"],
            },
        },
        {
            "model_name": "openai/gpt-4o",
            "litellm_params": {
                "model": "openai/gpt-4o",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "0cfcd87f2cb12a783a466888d05c6c89df66db23e01cecd75ec0b83aed73c9ad",
                "db_model": False,
                "access_groups": ["private-openai-models"],
            },
        },
    ]
    router = litellm.Router(model_list=llm_model_list)
    args = {
        "model": model,
        "llm_model_list": llm_model_list,
        "valid_token": UserAPIKeyAuth(
            models=["public-openai-models"],
        ),
        "llm_router": router,
    }
    if expect_to_work:
        await can_key_call_model(**args)
    else:
        with pytest.raises(Exception, match='is not available for this API key') as e:
            await can_key_call_model(**args)

        print(e)


@pytest.mark.parametrize(
    "model, expect_to_work",
    [("openai/gpt-4o", False), ("openai/gpt-4o-mini", True)],
)
@pytest.mark.asyncio
async def test_can_team_call_model(model, expect_to_work):
    from litellm.proxy.auth.auth_checks import model_in_access_group
    from fastapi import HTTPException

    llm_model_list = [
        {
            "model_name": "openai/*",
            "litellm_params": {
                "model": "openai/*",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
                "db_model": False,
                "access_groups": ["public-openai-models"],
            },
        },
        {
            "model_name": "openai/gpt-4o",
            "litellm_params": {
                "model": "openai/gpt-4o",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "0cfcd87f2cb12a783a466888d05c6c89df66db23e01cecd75ec0b83aed73c9ad",
                "db_model": False,
                "access_groups": ["private-openai-models"],
            },
        },
    ]
    router = litellm.Router(model_list=llm_model_list)

    args = {
        "model": model,
        "team_models": ["public-openai-models"],
        "llm_router": router,
    }
    if expect_to_work:
        assert model_in_access_group(**args)
    else:
        assert not model_in_access_group(**args)


@pytest.mark.parametrize(
    "key_models, model, expect_to_work",
    [
        (["openai/*"], "openai/gpt-4o", True),
        (["openai/*"], "openai/gpt-4o-mini", True),
        (["openai/*"], "openaiz/gpt-4o-mini", False),
        (["bedrock/*"], "bedrock/anthropic.claude-3-5-sonnet-20240620", True),
        (["bedrock/*"], "bedrockz/anthropic.claude-3-5-sonnet-20240620", False),
        (["bedrock/us.*"], "bedrock/us.amazon.nova-micro-v1:0", True),
        (["openai/*"], "ft:gpt-4-0613", True),
        (["openai/*"], "bedrockz/ft:gpt-4-0613", False),
    ],
)
@pytest.mark.asyncio
async def test_can_key_call_model_wildcard_access(key_models, model, expect_to_work):
    from litellm.proxy._types import ProxyException
    from litellm.proxy.auth.auth_checks import can_key_call_model

    llm_model_list = [
        {
            "model_name": "openai/*",
            "litellm_params": {
                "model": "openai/*",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
                "db_model": False,
            },
        },
        {
            "model_name": "bedrock/*",
            "litellm_params": {
                "model": "bedrock/*",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "e6e7006f83029df40ebc02ddd068890253f4cd3092bcb203d3d8e6f6f606f30f",
                "db_model": False,
            },
        },
        {
            "model_name": "openai/gpt-4o",
            "litellm_params": {
                "model": "openai/gpt-4o",
                "api_key": "test-api-key",
            },
            "model_info": {
                "id": "0cfcd87f2cb12a783a466888d05c6c89df66db23e01cecd75ec0b83aed73c9ad",
                "db_model": False,
            },
        },
    ]
    router = litellm.Router(model_list=llm_model_list)

    user_api_key_object = UserAPIKeyAuth(
        models=key_models,
    )

    if expect_to_work:
        await can_key_call_model(
            model=model,
            llm_model_list=llm_model_list,
            valid_token=user_api_key_object,
            llm_router=router,
        )
    else:
        with pytest.raises(ProxyException):
            await can_key_call_model(
                model=model,
                llm_model_list=llm_model_list,
                valid_token=user_api_key_object,
                llm_router=router,
            )


@pytest.fixture
def shipped_routing_rules() -> Iterator[None]:
    from litellm.litellm_core_utils.fallback_generalizations import (
        get_fallback_generalization_rules,
        set_fallback_generalizations,
    )
    from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap

    previous_rules: Final = list(get_fallback_generalization_rules())
    set_fallback_generalizations(GetModelCostMap.load_local_model_cost_map()["fallback_generalizations"]["rules"])
    yield
    set_fallback_generalizations(previous_rules)


@pytest.mark.parametrize("object_type", ["team", "project"])
@pytest.mark.parametrize(
    "allowed_models, model, expect_to_work",
    [
        (["typesafe/*"], "jev-latest", True),
        (["typesafe/*"], "jev-1.14.0", True),
        (["jev-latest"], "jev-latest", True),
        (["openai/*"], "jev-latest", False),
        (["typesafe/*"], "foo/jev-latest", False),
        (["typesafe/*"], "jev-internal-router", False),
        (["typesafe/*"], "jev-2.0-internal", False),
        (["typesafe/*"], "jev-1.0", False),
    ],
)
def test_typesafe_provider_wildcard_covers_bare_jev_model(
    shipped_routing_rules: None,
    object_type: Literal["team", "project"],
    allowed_models: list[str],
    model: str,
    expect_to_work: bool,
) -> None:
    from litellm.proxy._types import ProxyException
    from litellm.proxy.auth.auth_checks import can_object_call_model

    if expect_to_work:
        assert can_object_call_model(model=model, llm_router=None, models=allowed_models, object_type=object_type)
        return
    with pytest.raises(ProxyException):
        can_object_call_model(model=model, llm_router=None, models=allowed_models, object_type=object_type)


@pytest.mark.parametrize(
    "key_models, model, expect_to_work",
    [
        # After a cost-map reload, add_known_models() updates anthropic_models so
        # the anthropic/* wildcard can match a newly-added Anthropic model.
        (["anthropic/*"], "claude-brand-new-model-reload-test", True),
        # Wrong provider wildcard must still be denied even after reload.
        (["openai/*"], "claude-brand-new-model-reload-test", False),
    ],
)
@pytest.mark.asyncio
async def test_wildcard_access_after_cost_map_reload(key_models, model, expect_to_work):
    """
    Regression test: after a cost-map hot-reload, calling
    add_known_models(model_cost_map=new_map) must update litellm.anthropic_models
    so that the anthropic/* wildcard correctly grants (or denies) access to
    newly-added models.

    Root cause: both reload paths in proxy_server.py only updated
    litellm.model_cost but never re-ran add_known_models(), so the provider sets
    stayed stale and wildcard matching failed for new models.

    Fix: each reload now calls litellm.add_known_models(model_cost_map=new_map)
    with the fetched map passed explicitly to avoid any reference ambiguity.
    """
    from litellm.proxy._types import ProxyException
    from litellm.proxy.auth.auth_checks import can_key_call_model

    # Build a new cost map that includes the brand-new model — exactly what
    # proxy_server.py receives from get_model_cost_map() during a reload.
    new_cost_map = dict(litellm.model_cost)
    new_cost_map[model] = {
        "litellm_provider": "anthropic",
        "max_tokens": 8192,
        "input_cost_per_token": 0.000003,
        "output_cost_per_token": 0.000015,
    }

    original_model_cost = litellm.model_cost
    litellm.model_cost = new_cost_map

    # Confirm the model is NOT yet in the provider set before reload propagation.
    assert model not in litellm.anthropic_models

    # Simulate what proxy_server.py now does after every reload.
    litellm.add_known_models(model_cost_map=new_cost_map)

    # After add_known_models(), the model must be in the set.
    assert model in litellm.anthropic_models

    llm_model_list = [
        {
            "model_name": "anthropic/*",
            "litellm_params": {"model": "anthropic/*", "api_key": "test-api-key"},
            "model_info": {"id": "test-id-anthropic-wildcard", "db_model": False},
        },
        {
            "model_name": "openai/*",
            "litellm_params": {"model": "openai/*", "api_key": "test-api-key"},
            "model_info": {"id": "test-id-openai-wildcard", "db_model": False},
        },
    ]
    router = litellm.Router(model_list=llm_model_list)
    user_api_key_object = UserAPIKeyAuth(models=key_models)

    try:
        if expect_to_work:
            await can_key_call_model(
                model=model,
                llm_model_list=llm_model_list,
                valid_token=user_api_key_object,
                llm_router=router,
            )
        else:
            with pytest.raises(ProxyException):
                await can_key_call_model(
                    model=model,
                    llm_model_list=llm_model_list,
                    valid_token=user_api_key_object,
                    llm_router=router,
                )
    finally:
        litellm.model_cost = original_model_cost
        litellm.anthropic_models.discard(model)


@pytest.mark.asyncio
async def test_add_known_models_explicit_map_updates_provider_sets():
    """
    Regression test: after a cost-map hot-reload, calling
    add_known_models(model_cost_map=new_map) with the new map passed explicitly
    must add any new provider models to the correct provider sets so that
    wildcard access checks (anthropic/*, openai/*, …) work immediately.

    This covers the proxy_server.py fix where both reload paths now call
    litellm.add_known_models(model_cost_map=new_model_cost_map) instead of
    relying on the module-level model_cost being up to date.
    """
    fake_new_model = "claude-brand-new-explicit-map-test"

    # Baseline: the model must not be in the sets before we do anything.
    assert fake_new_model not in litellm.anthropic_models

    new_cost_map = dict(litellm.model_cost)
    new_cost_map[fake_new_model] = {
        "litellm_provider": "anthropic",
        "max_tokens": 8192,
        "input_cost_per_token": 0.000003,
        "output_cost_per_token": 0.000015,
    }

    # Simulate what proxy_server.py does on reload.
    original_model_cost = litellm.model_cost
    litellm.model_cost = new_cost_map
    litellm.add_known_models(model_cost_map=new_cost_map)

    try:
        assert fake_new_model in litellm.anthropic_models, (
            "add_known_models(model_cost_map=...) did not add the new model to "
            "litellm.anthropic_models — wildcard access checks would fail."
        )
    finally:
        # Clean up: restore original state.
        litellm.model_cost = original_model_cost
        litellm.anthropic_models.discard(fake_new_model)


@pytest.mark.asyncio
async def test_is_valid_fallback_model():
    from litellm.proxy.auth.auth_checks import is_valid_fallback_model
    from litellm import Router

    router = Router(
        model_list=[
            {
                "model_name": "gpt-3.5-turbo",
                "litellm_params": {"model": "openai/gpt-3.5-turbo"},
            }
        ]
    )

    try:
        await is_valid_fallback_model(
            model="gpt-3.5-turbo", llm_router=router, user_model=None
        )
    except Exception as e:
        pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}")

    with pytest.raises(Exception, match="Invalid") as exc_info:
        await is_valid_fallback_model(
            model="gpt-4o", llm_router=router, user_model=None
        )
    e = exc_info.value
    assert "Invalid" in str(e)


@pytest.mark.parametrize(
    "token_spend, max_budget, expect_budget_error",
    [
        (5.0, 10.0, False),  # Under budget
        (10.0, 10.0, True),  # At budget limit
        (15.0, 10.0, True),  # Over budget
    ],
)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_check(
    token_spend, max_budget, expect_budget_error
):
    """
    Test if virtual key budget checks work as expected:
    1. Triggers budget alert for all cases
    2. Raises BudgetExceededError when spend >= max_budget
    """
    from litellm.proxy.auth.auth_checks import virtual_key_max_budget_check

    # Setup test data
    valid_token = UserAPIKeyAuth(
        token="test-token",
        spend=token_spend,
        max_budget=max_budget,
        user_id="test-user",
        key_alias="test-key",
    )

    user_obj = LiteLLM_UserTable(
        user_id="test-user",
        user_email="test@email.com",
        max_budget=None,
    )

    proxy_logging_obj = ProxyLogging(
        user_api_key_cache=None,
    )

    # Track if budget alert was called
    alert_called = False

    async def mock_budget_alert(*args, **kwargs):
        nonlocal alert_called
        alert_called = True

    proxy_logging_obj.budget_alerts = mock_budget_alert

    if expect_budget_error:
        with pytest.raises(litellm.BudgetExceededError) as exc_info:
            await virtual_key_max_budget_check(
                valid_token=valid_token,
                proxy_logging_obj=proxy_logging_obj,
                user_obj=user_obj,
            )
        assert exc_info.value.current_cost == token_spend
        assert exc_info.value.max_budget == max_budget
    else:
        await virtual_key_max_budget_check(
            valid_token=valid_token,
            proxy_logging_obj=proxy_logging_obj,
            user_obj=user_obj,
        )

    await asyncio.sleep(1)

    # Verify budget alert was triggered
    assert alert_called, "Budget alert should be triggered"


@pytest.mark.parametrize(
    "model, team_models, expect_to_work",
    [
        ("gpt-4", ["gpt-4"], True),  # exact match
        ("gpt-4", ["all-proxy-models"], True),  # all-proxy-models access
        ("gpt-4", ["*"], True),  # wildcard access
        ("gpt-4", ["openai/*"], True),  # openai wildcard access
        (
            "bedrock/anthropic.claude-3-5-sonnet-20240620",
            ["bedrock/*"],
            True,
        ),  # wildcard access
        (
            "bedrockz/anthropic.claude-3-5-sonnet-20240620",
            ["bedrock/*"],
            False,
        ),  # non-match wildcard access
        ("bedrock/very_new_model", ["bedrock/*"], True),  # bedrock wildcard access
        (
            "bedrock/claude-3-5-sonnet-20240620",
            ["bedrock/claude-*"],
            True,
        ),  # match on pattern
        (
            "bedrock/claude-3-6-sonnet-20240620",
            ["bedrock/claude-3-5-*"],
            False,
        ),  # don't match on pattern
        ("openai/gpt-4o", ["openai/*"], True),  # openai wildcard access
        ("gpt-4", ["gpt-3.5-turbo"], False),  # model not in allowed list
        ("claude-3", [], True),  # empty model list (allows all)
    ],
)
@pytest.mark.asyncio
async def test_can_team_access_model(model, team_models, expect_to_work):
    """
    Test cases for can_team_access_model:
    1. Exact model match
    2. all-proxy-models access
    3. Wildcard (*) access
    4. OpenAI wildcard access
    5. Model not in allowed list
    6. Empty model list
    7. None model list
    """
    try:
        team_object = LiteLLM_TeamTable(
            team_id="test-team",
            models=team_models,
        )
        result = await can_team_access_model(
            model=model,
            team_object=team_object,
            llm_router=None,
            team_model_aliases=None,
        )
        if not expect_to_work:
            pytest.fail(
                f"Expected model access check to fail for model={model}, team_models={team_models}"
            )
    except Exception as e:
        if expect_to_work:
            pytest.fail(
                f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {str(e)}"
            )


@pytest.mark.parametrize(
    "spend, soft_budget, expect_alert",
    [
        (100, 50, True),  # Over soft budget
        (50, 50, True),  # At soft budget
        (25, 50, False),  # Under soft budget
        (100, None, False),  # No soft budget set
    ],
)
@pytest.mark.asyncio
async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
    """
    Test cases for _virtual_key_soft_budget_check:
    1. Spend over soft budget
    2. Spend at soft budget
    3. Spend under soft budget
    4. No soft budget set
    """
    alert_triggered = False

    class MockProxyLogging:
        async def budget_alerts(self, type, user_info):
            nonlocal alert_triggered
            alert_triggered = True
            assert type == "soft_budget"
            assert isinstance(user_info, CallInfo)

    valid_token = UserAPIKeyAuth(
        token="test-token",
        spend=spend,
        soft_budget=soft_budget,
        user_id="test-user",
        team_id="test-team",
        key_alias="test-key",
    )

    proxy_logging_obj = MockProxyLogging()

    await virtual_key_soft_budget_check(
        valid_token=valid_token,
        proxy_logging_obj=proxy_logging_obj,
    )

    await asyncio.sleep(0.1)  # Allow time for the alert task to complete

    assert (
        alert_triggered == expect_alert
    ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"


@pytest.mark.parametrize(
    "spend, soft_budget, expect_alert, metadata, expected_alert_emails",
    [
        (
            100,
            50,
            False,
            None,
            None,
        ),  # Over soft budget, no metadata - no alert_emails configured, so no alert
        (
            50,
            50,
            False,
            None,
            None,
        ),  # At soft budget, no metadata - no alert_emails configured, so no alert
        (25, 50, False, None, None),  # Under soft budget
        (100, None, False, None, None),  # No soft budget set
        (
            100,
            50,
            True,
            {"soft_budget_alerting_emails": ["team1@example.com", "team2@example.com"]},
            ["team1@example.com", "team2@example.com"],
        ),  # Over soft budget with list of emails
        (
            100,
            50,
            True,
            {"soft_budget_alerting_emails": "team1@example.com,team2@example.com"},
            ["team1@example.com", "team2@example.com"],
        ),  # Over soft budget with comma-separated emails
        (
            100,
            50,
            True,
            {
                "soft_budget_alerting_emails": [
                    "team1@example.com",
                    "",
                    "  ",
                    "team2@example.com",
                ]
            },
            ["team1@example.com", "team2@example.com"],
        ),  # Over soft budget with empty strings filtered
    ],
)
@pytest.mark.asyncio
async def test_team_soft_budget_check(
    spend, soft_budget, expect_alert, metadata, expected_alert_emails
):
    """
    Test cases for _team_soft_budget_check:
    1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
    2. Spend at soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
    3. Spend under soft budget - should not trigger alert
    4. No soft budget set - should not trigger alert
    5. Team with alert emails in metadata (list) - should include alert_emails in CallInfo
    6. Team with alert emails in metadata (comma-separated string) - should parse and include alert_emails
    7. Team with alert emails containing empty strings - should filter them out
    """
    alert_triggered = False
    captured_call_info = None

    class MockProxyLogging:
        async def budget_alerts(self, type, user_info):
            nonlocal alert_triggered, captured_call_info
            alert_triggered = True
            captured_call_info = user_info
            assert type == "soft_budget"
            assert isinstance(user_info, CallInfo)

    valid_token = UserAPIKeyAuth(
        token="test-token",
        user_id="test-user",
        team_id="test-team",
        team_alias="test-team-alias",
        key_alias="test-key",
    )

    team_object = LiteLLM_TeamTable(
        team_id="test-team",
        spend=spend,
        soft_budget=soft_budget,
        max_budget=100.0,
        metadata=metadata,
    )

    proxy_logging_obj = MockProxyLogging()

    await _team_soft_budget_check(
        team_object=team_object,
        valid_token=valid_token,
        proxy_logging_obj=proxy_logging_obj,
    )

    await asyncio.sleep(0.1)  # Allow time for the alert task to complete

    assert (
        alert_triggered == expect_alert
    ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"

    if expect_alert:
        assert captured_call_info is not None
        assert captured_call_info.team_id == "test-team"
        assert captured_call_info.spend == spend
        assert captured_call_info.soft_budget == soft_budget
        assert captured_call_info.event_group == Litellm_EntityType.TEAM
        # Verify alert_emails if expected
        if expected_alert_emails is not None:
            assert captured_call_info.alert_emails == expected_alert_emails
        else:
            assert (
                captured_call_info.alert_emails is None
                or captured_call_info.alert_emails == []
            )


@pytest.mark.asyncio
async def test_can_user_call_model():
    from litellm.proxy.auth.auth_checks import can_user_call_model
    from litellm.proxy._types import ProxyException
    from litellm import Router

    router = Router(
        model_list=[
            {
                "model_name": "anthropic-claude",
                "litellm_params": {"model": "anthropic/anthropic-claude"},
            },
            {
                "model_name": "gpt-3.5-turbo",
                "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "test-api-key"},
            },
        ]
    )

    args = {
        "model": "anthropic-claude",
        "llm_router": router,
        "user_object": LiteLLM_UserTable(
            user_id="testuser21@mycompany.com",
            max_budget=None,
            spend=0.0042295,
            model_max_budget={},
            model_spend={},
            user_email="testuser@mycompany.com",
            models=["gpt-3.5-turbo"],
        ),
    }

    with pytest.raises(ProxyException) as e:
        await can_user_call_model(**args)

    args["model"] = "gpt-3.5-turbo"
    await can_user_call_model(**args)


@pytest.mark.asyncio
async def test_can_user_call_model_with_no_default_models():
    from litellm.proxy.auth.auth_checks import can_user_call_model
    from litellm.proxy._types import ProxyException, SpecialModelNames
    from unittest.mock import MagicMock

    args = {
        "model": "anthropic-claude",
        "llm_router": MagicMock(),
        "user_object": LiteLLM_UserTable(
            user_id="testuser21@mycompany.com",
            max_budget=None,
            spend=0.0042295,
            model_max_budget={},
            model_spend={},
            user_email="testuser@mycompany.com",
            models=[SpecialModelNames.no_default_models.value],
        ),
    }

    with pytest.raises(ProxyException) as e:
        await can_user_call_model(**args)


@pytest.mark.asyncio
async def test_get_fuzzy_user_object():
    from litellm.proxy.auth.auth_checks import _get_fuzzy_user_object
    from unittest.mock import AsyncMock, MagicMock

    # Setup mock Prisma client
    mock_prisma = MagicMock()
    mock_prisma.db = MagicMock()
    mock_prisma.db.litellm_usertable = MagicMock()

    # Mock user data
    test_user = LiteLLM_UserTable(
        user_id="test_123",
        sso_user_id="sso_123",
        user_email="test@example.com",
        organization_memberships=[],
        max_budget=None,
    )

    # Test 1: Find user by SSO ID
    mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
    result = await _get_fuzzy_user_object(
        prisma_client=mock_prisma, sso_user_id="sso_123", user_email="test@example.com"
    )
    assert result == test_user
    mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
        where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
    )

    # Test 2: SSO ID not found, find by email
    mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
    mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
    mock_prisma.db.litellm_usertable.update = AsyncMock()

    result = await _get_fuzzy_user_object(
        prisma_client=mock_prisma,
        sso_user_id="new_sso_456",
        user_email="test@example.com",
    )
    assert result == test_user
    mock_prisma.db.litellm_usertable.find_first.assert_called_with(
        where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
        include={"organization_memberships": True},
    )

    # Test 3: Verify background SSO update task when user found by email
    await asyncio.sleep(0.1)  # Allow time for background task
    mock_prisma.db.litellm_usertable.update.assert_called_with(
        where={"user_id": "test_123"}, data={"sso_user_id": "new_sso_456"}
    )

    # Test 4: User not found by either method
    mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
    mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=None)

    result = await _get_fuzzy_user_object(
        prisma_client=mock_prisma,
        sso_user_id="unknown_sso",
        user_email="unknown@example.com",
    )
    assert result is None

    # Test 5: Only email provided (no SSO ID)
    mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
    result = await _get_fuzzy_user_object(
        prisma_client=mock_prisma, user_email="test@example.com"
    )
    assert result == test_user
    mock_prisma.db.litellm_usertable.find_first.assert_called_with(
        where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
        include={"organization_memberships": True},
    )

    # Test 6: Only SSO ID provided (no email)
    mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
    result = await _get_fuzzy_user_object(
        prisma_client=mock_prisma, sso_user_id="sso_123"
    )
    assert result == test_user
    mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
        where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
    )


@pytest.mark.parametrize(
    "model, alias_map, expect_to_work",
    [
        ("gpt-4", {"gpt-4": "gpt-4-team1"}, True),  # model matches alias value
        ("gpt-5", {"gpt-4": "gpt-4-team1"}, False),
    ],
)
@pytest.mark.asyncio
async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work):
    """
    Test if can_key_call_model correctly handles model aliases in the token
    """
    from litellm.proxy.auth.auth_checks import can_key_call_model

    llm_model_list = [
        {
            "model_name": "gpt-4-team1",
            "litellm_params": {
                "model": "gpt-4",
                "api_key": "test-api-key",
            },
        }
    ]
    router = litellm.Router(model_list=llm_model_list)

    user_api_key_object = UserAPIKeyAuth(
        models=[
            "gpt-4-team1",
        ],
        team_model_aliases=alias_map,
    )

    if expect_to_work:
        await can_key_call_model(
            model=model,
            llm_model_list=llm_model_list,
            valid_token=user_api_key_object,
            llm_router=router,
        )
    else:
        with pytest.raises(Exception, match='is not available for this API key') as e:
            await can_key_call_model(
                model=model,
                llm_model_list=llm_model_list,
                valid_token=user_api_key_object,
                llm_router=router,
            )


# ---------------------------------------------------------------------------
# Access group cache helpers (_cache_access_object, _delete_cache_access_object)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_cache_access_object():
    """Test _cache_access_object stores access group in cache with correct key."""
    from litellm.proxy.auth.auth_checks import cache_access_object
    from litellm.proxy._types import LiteLLM_AccessGroupTable

    cache = DualCache()
    ag_id = "ag-test-123"
    ag_table = LiteLLM_AccessGroupTable(
        access_group_id=ag_id,
        access_group_name="test-group",
        access_model_names=["gpt-4"],
    )
    await cache_access_object(
        access_group_id=ag_id,
        access_group_table=ag_table,
        user_api_key_cache=cache,
    )
    cached = await cache.async_get_cache(key=f"access_group_id:{ag_id}")
    assert cached is not None
    if isinstance(cached, dict):
        assert cached.get("access_group_id") == ag_id
        assert cached.get("access_group_name") == "test-group"
    else:
        assert cached.access_group_id == ag_id
        assert cached.access_group_name == "test-group"


@pytest.mark.asyncio
async def test_delete_cache_access_object():
    """Test _delete_cache_access_object removes access group from in-memory cache."""
    from litellm.proxy.auth.auth_checks import delete_cache_access_object
    from litellm.proxy._types import LiteLLM_AccessGroupTable

    cache = DualCache()
    ag_id = "ag-delete-test"
    ag_table = LiteLLM_AccessGroupTable(
        access_group_id=ag_id,
        access_group_name="to-delete",
    )
    await cache.async_set_cache(key=f"access_group_id:{ag_id}", value=ag_table, ttl=60)
    await delete_cache_access_object(access_group_id=ag_id, user_api_key_cache=cache)
    cached = await cache.async_get_cache(key=f"access_group_id:{ag_id}")
    assert cached is None


# ---------------------------------------------------------------------------
# Access group resource fetchers (_get_models_from_access_groups, _get_agent_ids_from_access_groups)
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    "resource_field, access_group_data, expected",
    [
        (
            "access_model_names",
            {"access_group_id": "ag-1", "access_model_names": ["gpt-4", "claude-3"]},
            ["gpt-4", "claude-3"],
        ),
        (
            "access_agent_ids",
            {"access_group_id": "ag-2", "access_agent_ids": ["agent-a", "agent-b"]},
            ["agent-a", "agent-b"],
        ),
        (
            "access_model_names",
            {"access_group_id": "ag-3", "access_model_names": []},
            [],
        ),
    ],
)
@pytest.mark.asyncio
async def test_get_resources_from_access_groups(
    resource_field, access_group_data, expected
):
    """Test _get_resources_from_access_groups returns correct resource list from access groups."""
    from unittest.mock import AsyncMock, MagicMock, patch

    from litellm.proxy._types import LiteLLM_AccessGroupTable
    from litellm.proxy.auth.auth_checks import (
        get_agent_ids_from_access_groups,
        get_models_from_access_groups,
    )

    ag_table = LiteLLM_AccessGroupTable(
        access_group_id=access_group_data["access_group_id"],
        access_group_name="test",
        access_model_names=access_group_data.get("access_model_names", []),
        access_agent_ids=access_group_data.get("access_agent_ids", []),
    )

    with patch(
        "litellm.proxy.auth.auth_checks.get_access_object",
        new_callable=AsyncMock,
        return_value=ag_table,
    ):
        if resource_field == "access_model_names":
            result = await get_models_from_access_groups(
                access_group_ids=[access_group_data["access_group_id"]],
                prisma_client=MagicMock(),
                user_api_key_cache=DualCache(),
            )
        else:
            result = await get_agent_ids_from_access_groups(
                access_group_ids=[access_group_data["access_group_id"]],
                prisma_client=MagicMock(),
                user_api_key_cache=DualCache(),
            )
        assert sorted(result) == sorted(expected)


@pytest.mark.asyncio
async def test_get_models_from_access_groups_empty_ids():
    """Test _get_models_from_access_groups returns empty list when access_group_ids is empty."""
    from litellm.proxy.auth.auth_checks import get_models_from_access_groups

    result = await get_models_from_access_groups(access_group_ids=[])
    assert result == []


# ---------------------------------------------------------------------------
# can_team_access_model with access_group_ids fallback
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_can_team_access_model_via_access_group_ids():
    """Test can_team_access_model allows access when team has access_group_ids granting model access."""
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import can_team_access_model

    team_object = LiteLLM_TeamTable(
        team_id="test-team",
        models=[],
        access_group_ids=["ag-with-gpt4"],
    )

    with patch(
        "litellm.proxy.auth.auth_checks.get_models_from_access_groups",
        new_callable=AsyncMock,
        return_value=["gpt-4"],
    ):
        result = await can_team_access_model(
            model="gpt-4",
            team_object=team_object,
            llm_router=None,
            team_model_aliases=None,
        )
        assert result is True


@pytest.mark.asyncio
async def test_can_team_access_model_access_group_ids_denied():
    """Test can_team_access_model denies when neither team models nor access_group_ids grant access."""
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import can_team_access_model
    from litellm.proxy._types import ProxyException

    team_object = LiteLLM_TeamTable(
        team_id="test-team",
        models=["gpt-3.5-turbo"],
        access_group_ids=["ag-other"],
    )

    with patch(
        "litellm.proxy.auth.auth_checks.get_models_from_access_groups",
        new_callable=AsyncMock,
        return_value=["claude-3"],
    ):
        with pytest.raises(ProxyException):
            await can_team_access_model(
                model="gpt-4",
                team_object=team_object,
                llm_router=None,
                team_model_aliases=None,
            )


# ---------------------------------------------------------------------------
# can_key_call_model with access_group_ids fallback
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_can_key_call_model_via_access_group_ids():
    """Test can_key_call_model allows access when key has access_group_ids granting model access."""
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import can_key_call_model

    user_api_key_object = UserAPIKeyAuth(
        token="test-token",
        models=[],
        access_group_ids=["ag-with-gpt4"],
    )
    router = litellm.Router(
        model_list=[
            {
                "model_name": "gpt-4",
                "litellm_params": {"model": "openai/gpt-4", "api_key": "test"},
            }
        ]
    )

    with patch(
        "litellm.proxy.auth.auth_checks.get_models_from_access_groups",
        new_callable=AsyncMock,
        return_value=["gpt-4"],
    ):
        await can_key_call_model(
            model="gpt-4",
            llm_model_list=[],
            valid_token=user_api_key_object,
            llm_router=router,
        )


# ---------------------------------------------------------------------------
# _key_access_group_grants_model (key access group overriding team restriction)
# ---------------------------------------------------------------------------


def _patch_proxy_server_globals():
    """Patch proxy_server's prisma_client and user_api_key_cache to non-None mocks
    so the helper's None-guard doesn't short-circuit. The actual values don't
    matter because get_access_object is patched separately to return fixtures."""
    from unittest.mock import MagicMock, patch

    return [
        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
    ]


def _fake_access_group(
    access_group_id: str,
    access_model_names=None,
    assigned_team_ids=None,
    assigned_key_ids=None,
):
    from litellm.proxy._types import LiteLLM_AccessGroupTable

    return LiteLLM_AccessGroupTable(
        access_group_id=access_group_id,
        access_group_name=access_group_id,
        access_model_names=access_model_names or [],
        assigned_team_ids=assigned_team_ids or [],
        assigned_key_ids=assigned_key_ids or [],
    )


@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_team_authorized():
    """Group's assigned_team_ids includes the key's team and grants the model → True.

    This is the happy path equivalent of Andres's report: admin creates an
    access group with assigned_team_ids=[team-a], grants claude-haiku-4-5,
    attaches it to a key on team-a. Override fires.
    """
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import key_access_group_grants_model

    valid_token = UserAPIKeyAuth(
        token="test-token",
        models=[],
        access_group_ids=["premium-group"],
        team_id="team-a",
    )
    team_object = LiteLLM_TeamTable(
        team_id="team-a",
        models=["mock-success"],
        access_group_ids=[],  # deliberately not synced — the access group itself authorizes
    )

    fake_ag = _fake_access_group(
        access_group_id="premium-group",
        access_model_names=["claude-haiku-4-5"],
        assigned_team_ids=["team-a"],
    )

    patches = _patch_proxy_server_globals() + [
        patch(
            "litellm.proxy.auth.auth_checks.get_access_object",
            new_callable=AsyncMock,
            return_value=fake_ag,
        ),
    ]
    for p in patches:
        p.start()
    try:
        assert (
            await key_access_group_grants_model(
                model="claude-haiku-4-5",
                valid_token=valid_token,
                team_object=team_object,
                llm_router=None,
            )
            is True
        )
    finally:
        for p in patches:
            p.stop()


@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_key_directly_authorized():
    """Group's assigned_key_ids includes the key's token and grants the model → True.

    Per-key authorization path: an admin scopes a group directly to a key
    (assigned_key_ids) without listing the team.
    """
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import key_access_group_grants_model

    valid_token = UserAPIKeyAuth(
        token="test-token-hashed",
        models=[],
        access_group_ids=["per-key-group"],
        team_id="team-a",
    )
    team_object = LiteLLM_TeamTable(
        team_id="team-a",
        models=["mock-success"],
        access_group_ids=[],
    )

    fake_ag = _fake_access_group(
        access_group_id="per-key-group",
        access_model_names=["claude-haiku-4-5"],
        assigned_team_ids=[],
        assigned_key_ids=["test-token-hashed"],
    )

    patches = _patch_proxy_server_globals() + [
        patch(
            "litellm.proxy.auth.auth_checks.get_access_object",
            new_callable=AsyncMock,
            return_value=fake_ag,
        ),
    ]
    for p in patches:
        p.start()
    try:
        assert (
            await key_access_group_grants_model(
                model="claude-haiku-4-5",
                valid_token=valid_token,
                team_object=team_object,
                llm_router=None,
            )
            is True
        )
    finally:
        for p in patches:
            p.stop()


@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_key_has_no_groups():
    """Key with no access_group_ids → False (early return, no DB read)."""
    from litellm.proxy.auth.auth_checks import key_access_group_grants_model

    valid_token = UserAPIKeyAuth(
        token="test-token",
        models=[],
        access_group_ids=[],
        team_id="team-a",
    )
    team_object = LiteLLM_TeamTable(
        team_id="team-a",
        models=["mock-success"],
        access_group_ids=["any-group"],
    )
    assert (
        await key_access_group_grants_model(
            model="claude-haiku-4-5",
            valid_token=valid_token,
            team_object=team_object,
            llm_router=None,
        )
        is False
    )


@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_group_does_not_cover_model():
    """Group authorizes the team but does not grant the requested model → False."""
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import key_access_group_grants_model

    valid_token = UserAPIKeyAuth(
        token="test-token",
        models=[],
        access_group_ids=["basic-group"],
        team_id="team-a",
    )
    team_object = LiteLLM_TeamTable(
        team_id="team-a",
        models=["mock-success"],
        access_group_ids=[],
    )

    fake_ag = _fake_access_group(
        access_group_id="basic-group",
        access_model_names=["gpt-4o-mini"],
        assigned_team_ids=["team-a"],
    )

    patches = _patch_proxy_server_globals() + [
        patch(
            "litellm.proxy.auth.auth_checks.get_access_object",
            new_callable=AsyncMock,
            return_value=fake_ag,
        ),
    ]
    for p in patches:
        p.start()
    try:
        assert (
            await key_access_group_grants_model(
                model="claude-haiku-4-5",
                valid_token=valid_token,
                team_object=team_object,
                llm_router=None,
            )
            is False
        )
    finally:
        for p in patches:
            p.stop()


@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_group_authorizes_neither():
    """
    Bypass regression test: a team member sets a foreign access group on their
    key. The group grants the requested model but its assigned_team_ids /
    assigned_key_ids do not include this caller's team or token. Override is
    denied — the team's 401 propagates.
    """
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import key_access_group_grants_model

    valid_token = UserAPIKeyAuth(
        token="team-a-token",
        models=[],
        access_group_ids=["team-b-premium"],
        team_id="team-a",
    )
    team_object = LiteLLM_TeamTable(
        team_id="team-a",
        models=["mock-success"],
        access_group_ids=[],
    )

    fake_ag = _fake_access_group(
        access_group_id="team-b-premium",
        access_model_names=["claude-opus-4-5"],
        assigned_team_ids=["team-b"],
        assigned_key_ids=["team-b-token"],
    )

    patches = _patch_proxy_server_globals() + [
        patch(
            "litellm.proxy.auth.auth_checks.get_access_object",
            new_callable=AsyncMock,
            return_value=fake_ag,
        ),
    ]
    for p in patches:
        p.start()
    try:
        assert (
            await key_access_group_grants_model(
                model="claude-opus-4-5",
                valid_token=valid_token,
                team_object=team_object,
                llm_router=None,
            )
            is False
        )
    finally:
        for p in patches:
            p.stop()


@pytest.mark.asyncio
async def test_key_access_group_grants_model_when_get_access_object_raises():
    """Group lookup failure (404, network, etc.) is treated as no authorization."""
    from unittest.mock import AsyncMock, patch

    from litellm.proxy.auth.auth_checks import key_access_group_grants_model

    valid_token = UserAPIKeyAuth(
        token="test-token",
        models=[],
        access_group_ids=["missing-group"],
        team_id="team-a",
    )
    team_object = LiteLLM_TeamTable(
        team_id="team-a",
        models=["mock-success"],
        access_group_ids=[],
    )

    patches = _patch_proxy_server_globals() + [
        patch(
            "litellm.proxy.auth.auth_checks.get_access_object",
            new_callable=AsyncMock,
            side_effect=Exception("not found"),
        ),
    ]
    for p in patches:
        p.start()
    try:
        assert (
            await key_access_group_grants_model(
                model="claude-haiku-4-5",
                valid_token=valid_token,
                team_object=team_object,
                llm_router=None,
            )
            is False
        )
    finally:
        for p in patches:
            p.stop()


@pytest.fixture
def db_success_hook() -> Iterator[AsyncMock]:
    hook: Final = AsyncMock()
    with patch(
        "litellm.proxy.proxy_server.proxy_logging_obj",
        MagicMock(service_logging_obj=MagicMock(async_service_success_hook=hook)),
    ):
        yield hook


async def _db_service_call_types(hook: AsyncMock) -> tuple[str, ...]:
    await asyncio.sleep(0)
    return tuple(call.kwargs["call_type"] for call in hook.await_args_list if call.kwargs["service"] == ServiceTypes.DB)


def _prisma_client_serving(user_id: str) -> SimpleNamespace:
    row: Final = {
        "user_id": user_id,
        "user_role": "internal_user",
        "teams": [],
        "spend": 0.0,
        "models": [],
        "metadata": "{}",
        "allowed_cache_controls": [],
        "policies": [],
        "model_spend": "{}",
        "model_max_budget": "{}",
        "organization_memberships": [],
    }
    engine: Final = SimpleNamespace(query=AsyncMock(return_value={"data": {"result": row}}), stop=lambda: None)
    generated_client: Final = Prisma()
    generated_client._engine = engine
    return SimpleNamespace(db=PrismaWrapper(original_prisma=generated_client, iam_token_db_auth=False))


@pytest.mark.asyncio
async def test_get_user_object_cache_hit_emits_no_postgres_service_event(db_success_hook: AsyncMock) -> None:
    user_id: Final = f"cached-user-{uuid.uuid4()}"
    cache: Final = UserApiKeyCache()
    await cache.async_set_cache(key=user_id, value=LiteLLM_UserTable(user_id=user_id, user_role="internal_user"))

    result: Final = await get_user_object(
        user_id=user_id,
        prisma_client=MagicMock(),
        user_api_key_cache=cache,
        user_id_upsert=False,
        parent_otel_span="auth-span",
    )

    assert result is not None and result.user_id == user_id
    assert await _db_service_call_types(db_success_hook) == ()


@pytest.mark.asyncio
async def test_get_org_object_cache_hit_emits_no_postgres_service_event(db_success_hook: AsyncMock) -> None:
    org_id: Final = f"cached-org-{uuid.uuid4()}"
    cache: Final = UserApiKeyCache()
    await cache.async_set_cache(
        key=f"org_id:{org_id}",
        value=LiteLLM_OrganizationTable(
            organization_id=org_id, budget_id="b", models=[], created_by="t", updated_by="t"
        ),
    )

    result: Final = await get_org_object(
        org_id=org_id,
        prisma_client=MagicMock(),
        user_api_key_cache=cache,
        parent_otel_span="auth-span",
    )

    assert result is not None and result.organization_id == org_id
    assert await _db_service_call_types(db_success_hook) == ()


@pytest.mark.asyncio
async def test_get_user_object_cache_miss_emits_exactly_one_postgres_get_user_object_event(
    db_success_hook: AsyncMock,
) -> None:
    user_id: Final = f"db-user-{uuid.uuid4()}"
    prisma_client: Final = _prisma_client_serving(user_id)

    result: Final = await get_user_object(
        user_id=user_id,
        prisma_client=prisma_client,
        user_api_key_cache=UserApiKeyCache(),
        user_id_upsert=False,
        parent_otel_span="auth-span",
    )

    assert result is not None and result.user_id == user_id
    assert await _db_service_call_types(db_success_hook) == ("get_user_object",)
    assert db_success_hook.await_args_list[0].kwargs["parent_otel_span"] == "auth-span"


@pytest.mark.parametrize("entry_first", [True, False])
def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_routes_to(
    monkeypatch: pytest.MonkeyPatch, entry_first: bool
) -> None:
    """chain-entry resolves one hop to local-free and is served by local-free's own free
    deployment, so an over-budget key is waived for it; local-free by name resolves to paid-gpt
    and stays enforced. local-free's alias is a hop the router never takes for chain-entry, and
    the per-name verdict cache must not let either name's verdict leak into the other's."""
    monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
    router: Final = litellm.Router(
        model_list=[
            {
                "model_name": "local-free",
                "litellm_params": {
                    "model": "ollama/qwen3:0.6b",
                    "api_base": "http://localhost:11434",
                    "input_cost_per_token": 0,
                    "output_cost_per_token": 0,
                },
            },
            {
                "model_name": "paid-gpt",
                "litellm_params": {
                    "model": "gpt-4o",
                    "api_key": "fake",
                    "input_cost_per_token": 3e-06,
                    "output_cost_per_token": 1.5e-05,
                },
            },
        ],
        model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"},
    )
    expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False}
    order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")

    verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order}

    assert verdicts == expected
    assert {name: is_model_cost_zero(model=name, llm_router=router) for name in order} == expected


MCP_DISCOVERY_REST_ROUTES = [
    "/mcp-rest/tools/list",
    "/v1/mcp/tools",
]


@pytest.mark.parametrize("route", MCP_DISCOVERY_REST_ROUTES)
@pytest.mark.asyncio
async def test_mcp_tool_discovery_route_bypasses_team_budget(route):
    from litellm.proxy.auth.auth_checks import common_checks

    team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)

    result = await common_checks(
        request_body={},
        team_object=team_object,
        user_object=None,
        end_user_object=None,
        global_proxy_spend=None,
        general_settings={},
        route=route,
        llm_router=None,
        proxy_logging_obj=AsyncMock(),
        valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
        request=MagicMock(),
    )

    assert result is True


@pytest.mark.parametrize("method", ["initialize", "notifications/initialized", "ping", "tools/list"])
@pytest.mark.parametrize("route", ["/mcp", "/mcp/some-server"])
@pytest.mark.asyncio
async def test_mcp_jsonrpc_zero_spend_method_bypasses_team_budget(route, method):
    from litellm.proxy.auth.auth_checks import common_checks

    team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)

    result = await common_checks(
        request_body={"jsonrpc": "2.0", "id": 1, "method": method, "params": {}},
        team_object=team_object,
        user_object=None,
        end_user_object=None,
        global_proxy_spend=None,
        general_settings={},
        route=route,
        llm_router=None,
        proxy_logging_obj=AsyncMock(),
        valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
        request=MagicMock(),
    )

    assert result is True


@pytest.mark.parametrize(
    "route,request_body",
    [
        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
        ("/mcp/some-server", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
        ("/mcp-rest/tools/call", {"name": "add", "arguments": {}}),
        ("/mcp/tools/call", {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}),
        ("/mcp/tools", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
        ("/mcp/tools/list", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
    ],
)
@pytest.mark.asyncio
async def test_mcp_tool_call_route_still_enforces_team_budget(route, request_body):
    from litellm.proxy.auth.auth_checks import common_checks

    team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)

    with pytest.raises(litellm.BudgetExceededError):
        await common_checks(
            request_body=request_body,
            team_object=team_object,
            user_object=None,
            end_user_object=None,
            global_proxy_spend=None,
            general_settings={},
            route=route,
            llm_router=None,
            proxy_logging_obj=AsyncMock(),
            valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
            request=MagicMock(),
        )


@pytest.mark.parametrize(
    "route,request_body,expected",
    [
        ("/mcp-rest/tools/list", {}, True),
        ("/v1/mcp/tools", {}, True),
        ("/mcp/tools", {}, False),
        ("/mcp/tools/list", {}, False),
        ("/mcp/tools", {"method": "tools/call"}, False),
        ("/mcp/tools", {"method": "tools/list"}, True),
        ("/mcp", {"method": "initialize"}, True),
        ("/mcp", {"method": "notifications/initialized"}, True),
        ("/mcp", {"method": "ping"}, True),
        ("/mcp", {"method": "tools/list"}, True),
        ("/mcp/some-server", {"method": "tools/list"}, True),
        ("/mcp", {"method": "tools/call"}, False),
        ("/mcp/some-server", {"method": "tools/call"}, False),
        ("/mcp", {}, False),
        ("/mcp", {"method": ["tools/list"]}, False),
        ("/mcp", {"method": {"name": "ping"}}, False),
        ("/mcp-rest/tools/call", {"method": "initialize"}, False),
        ("/mcp/tools/call", {"method": "initialize"}, False),
        ("/v1/chat/completions", {"method": "initialize"}, False),
        ("/mcp-rest/other", {"method": "tools/list"}, False),
    ],
)
def test_is_mcp_discovery_request(route, request_body, expected):
    from litellm.proxy.auth.auth_checks import is_mcp_discovery_request

    assert is_mcp_discovery_request(route=route, request_body=request_body) is expected


@pytest.mark.asyncio
async def test_get_skill_names_from_access_groups_reads_cached_group_rows():
    from litellm.proxy._types import LiteLLM_AccessGroupTable
    from litellm.proxy.auth.auth_checks import cache_access_object, get_skill_names_from_access_groups
    from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache

    cache = UserApiKeyCache()
    for group_id, skills in (("group-a", ["skill-a", "shared"]), ("group-b", ["skill-b", "shared"])):
        await cache_access_object(
            access_group_id=group_id,
            access_group_table=LiteLLM_AccessGroupTable(
                access_group_id=group_id, access_group_name=group_id, access_skill_names=skills
            ),
            user_api_key_cache=cache,
        )
    untouched_prisma = MagicMock()
    untouched_prisma.db.litellm_accessgrouptable.find_unique = AsyncMock(
        side_effect=AssertionError("cache hit must not reach the DB")
    )

    assert await get_skill_names_from_access_groups(
        access_group_ids=("group-a", "group-b"),
        prisma_client=untouched_prisma,
        user_api_key_cache=cache,
        proxy_logging_obj=MagicMock(),
    ) == frozenset({"skill-a", "skill-b", "shared"})


class _Pipeline:
    def __init__(self, send: Callable[[Sequence[str]], Awaitable[None]]) -> None:
        self._send = send
        self._messages: list[str] = []

    def publish(self, channel: str, message: str) -> "_Pipeline":
        self._messages.append(message)
        return self

    async def execute(self) -> None:
        await self._send(tuple(self._messages))


class _SlowPublishRedisClient:
    def __init__(self) -> None:
        self.published: list[str] = []

    async def _send(self, messages: Sequence[str]) -> None:
        await asyncio.sleep(0.001)
        self.published.extend(messages)

    async def publish(self, channel: str, message: str) -> int:
        await self._send((message,))
        return 1

    def pipeline(self, transaction: bool = True) -> _Pipeline:
        return _Pipeline(self._send)


class _WedgedPublishRedisClient:
    def __init__(self) -> None:
        self.release = asyncio.Event()

    async def _send(self, messages: Sequence[str]) -> None:
        await self.release.wait()

    async def publish(self, channel: str, message: str) -> int:
        await self._send((message,))
        return 1

    def pipeline(self, transaction: bool = True) -> _Pipeline:
        return _Pipeline(self._send)


class _PubSubRedisCache:
    namespace = None

    def __init__(self, client: _SlowPublishRedisClient | _WedgedPublishRedisClient) -> None:
        self._client = client

    def init_pubsub_client(self) -> _SlowPublishRedisClient | _WedgedPublishRedisClient:
        return self._client


async def _drain_publishes() -> None:
    await asyncio.gather(*pubsub_module._pending_publishes)  # pyright: ignore[reportPrivateUsage]  # drain module-level tasks


@pytest.mark.asyncio
async def test_evicting_more_keys_than_the_publisher_holds_pending_broadcasts_every_one(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
    monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
    client = _SlowPublishRedisClient()
    hashed_tokens: Final = tuple(f"token-{index}" for index in range(pubsub_module._MAX_PENDING_PUBLISHES + 476))  # pyright: ignore[reportPrivateUsage]  # the cap under test

    with patch(
        "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
        return_value=_PubSubRedisCache(client),
    ):
        await delete_cache_key_objects(
            hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=None
        )
        await _drain_publishes()

    assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens)


@pytest.mark.asyncio
async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
    monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
    monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
    client = _WedgedPublishRedisClient()
    cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
    hashed_tokens: Final = tuple(
        hashlib.sha256(f"token-{index}".encode()).hexdigest()
        for index in range(600)
    )
    for hashed_token in hashed_tokens:
        cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})
    seeded: Final = tuple(
        token for token in hashed_tokens if cache.in_memory_cache_for(token).get_cache(token) is not None
    )

    with patch(
        "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
        return_value=_PubSubRedisCache(client),
    ):
        eviction: Final = asyncio.create_task(
            delete_cache_key_objects(hashed_tokens=hashed_tokens, user_api_key_cache=cache, proxy_logging_obj=None)
        )
        await asyncio.sleep(0.05)
        still_cached: Final = tuple(
            token for token in seeded if cache.in_memory_cache_for(token).get_cache(token) is not None
        )
        client.release.set()
        await eviction
        await _drain_publishes()

    assert seeded == hashed_tokens
    assert still_cached == ()


class _AuthRedisCache:
    namespace = None

    def __init__(self, entries: dict[str, object]) -> None:
        self.entries = dict(entries)

    async def async_get_cache(self, key: str, **kwargs: object) -> object | None:
        return self.entries.get(key)

    def set_cache(self, key: str, value: object, **kwargs: object) -> None:
        self.entries[key] = value

    async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
        self.entries[key] = value

    def delete_cache(self, key: str) -> None:
        self.entries.pop(key, None)

    async def async_delete_cache(self, key: str) -> None:
        self.entries.pop(key, None)

    async def delete_cache_keys(self, keys: Sequence[str]) -> None:
        for key in keys:
            self.entries.pop(key, None)


@pytest.mark.asyncio
async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_broadcast(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
    monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
    monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
    client = _WedgedPublishRedisClient()
    hashed_tokens: Final = tuple(
        hashlib.sha256(f"token-{index}".encode()).hexdigest()
        for index in range(600)
    )
    redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens})
    cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
    cache.attach_redis_cache(redis)  # pyright: ignore[reportArgumentType]  # a dict-backed fake of the Redis layer
    for hashed_token in hashed_tokens:
        cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})

    with patch(
        "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
        return_value=_PubSubRedisCache(client),
    ):
        eviction: Final = asyncio.create_task(
            delete_cache_key_objects(hashed_tokens=hashed_tokens, user_api_key_cache=cache, proxy_logging_obj=None)
        )
        await asyncio.sleep(0.05)
        refilled: Final = [token for token in hashed_tokens if await cache.async_get_cache(key=token) is not None]
        client.release.set()
        await eviction
        await _drain_publishes()

    assert refilled == []
    assert redis.entries == {}


class _SignallingPublishRedisClient(_SlowPublishRedisClient):
    def __init__(self) -> None:
        super().__init__()
        self.sent = asyncio.Event()

    async def _send(self, messages: Sequence[str]) -> None:
        await super()._send(messages)
        self.sent.set()


class _DeleteAfterBroadcastRedisCache(_AuthRedisCache):
    def __init__(self, broadcast_sent: asyncio.Event) -> None:
        super().__init__({})
        self._broadcast_sent = broadcast_sent
        self.finished = False

    async def delete_cache_keys(self, keys: Sequence[str]) -> None:
        await self._broadcast_sent.wait()
        self.finished = True


@pytest.mark.asyncio
async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
    monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
    client: Final = _SignallingPublishRedisClient()
    usage_redis: Final = _DeleteAfterBroadcastRedisCache(broadcast_sent=client.sent)
    proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=None)
    proxy_logging_obj.internal_usage_cache = InternalUsageCache(
        dual_cache=DualCache(redis_cache=usage_redis)  # pyright: ignore[reportArgumentType]  # a dict-backed fake of the Redis layer
    )
    hashed_tokens: Final = ("token-a", "token-b")

    with patch(
        "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
        return_value=_PubSubRedisCache(client),
    ):
        await delete_cache_key_objects(
            hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=proxy_logging_obj
        )
        await _drain_publishes()

    assert usage_redis.finished
    assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens)