import os
import sys
import traceback
from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping
from contextlib import asynccontextmanager
from typing import Final, Literal, NoReturn, cast
from unittest import mock

from dotenv import load_dotenv

from tests._master_key import MASTER_KEY

import litellm.proxy
import litellm.proxy.proxy_server

load_dotenv()
import io
import json

# this file is to test litellm/proxy

import asyncio
import logging

import pytest

import litellm
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding

# Configure logging
logging.basicConfig(
    level=logging.DEBUG,  # Set the desired logging level
    format="%(asctime)s - %(levelname)s - %(message)s",
)

from unittest.mock import AsyncMock, MagicMock, patch

from fastapi import FastAPI, HTTPException, WebSocket

# test /chat/completion request to the proxy
from fastapi.testclient import TestClient

from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import (  # Replace with the actual module where your FastAPI router is defined
    app,
    initialize,
    save_worker_config,
)
from litellm.proxy.utils import ProxyLogging

@pytest.fixture
def admin_mcp_proxy(monkeypatch: pytest.MonkeyPatch) -> FastAPI:
    from litellm.proxy import proxy_server
    from litellm.proxy.auth.litellm_license import LicenseCheck

    monkeypatch.setenv("LITELLM_ENABLE_ADMIN_MCP", "true")
    monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-" + "1234567890abcdef" * 4)
    for name in ("WORKER_CONFIG", "CONFIG_FILE_PATH", "DATABASE_URL", "LITELLM_LICENSE"):
        monkeypatch.delenv(name, raising=False)
    monkeypatch.setattr(proxy_server, "_license_check", LicenseCheck())
    monkeypatch.setattr(proxy_server, "premium_user", True)
    monkeypatch.setattr(proxy_server, "prisma_client", None)
    monkeypatch.setattr(proxy_server, "general_settings", {"disable_model_info_refresh": True})
    monkeypatch.setattr(proxy_server, "scheduler", None)
    return FastAPI(lifespan=proxy_server.proxy_startup_event)


@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
@pytest.mark.parametrize("failure_phase", ["startup", "serving", "shutdown", "cancelled", "license"])
async def test_admin_mcp_failure_still_closes_proxy_resources(
    admin_mcp_proxy: FastAPI,
    monkeypatch: pytest.MonkeyPatch,
    failure_phase: Literal["startup", "serving", "shutdown", "cancelled", "license"],
) -> None:
    from apscheduler.schedulers.asyncio import AsyncIOScheduler
    from apscheduler.schedulers.base import STATE_PAUSED
    from litellm_admin_mcp import server as connector_server
    from litellm_admin_mcp.gateway import Gateway

    from litellm.proxy import proxy_server
    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
    from litellm.proxy.shutdown.scheduled_jobs import AwaitableAsyncIOExecutor

    monkeypatch.setattr(proxy_server, "premium_user", failure_phase != "license")
    executor: Final = AwaitableAsyncIOExecutor()
    scheduler: Final = AsyncIOScheduler(executors={"default": executor})
    monkeypatch.setattr(proxy_server, "scheduler", scheduler)
    monkeypatch.setattr(proxy_server, "scheduler_executor", executor)
    scheduler.start()

    @asynccontextmanager
    async def failing_connector(_app: FastAPI) -> AsyncGenerator[None, None]:
        if failure_phase == "startup":
            raise RuntimeError("startup failed")
        yield
        assert scheduler.state == STATE_PAUSED
        assert GracefulShutdownManager.is_shutting_down()
        assert proxy_server.shared_aiohttp_session is not None
        assert not proxy_server.shared_aiohttp_session.closed
        if failure_phase == "shutdown":
            raise RuntimeError("shutdown failed")

    def connector_app(_gateway: Gateway) -> FastAPI:
        return FastAPI(lifespan=failing_connector)

    monkeypatch.setattr(connector_server, "create_http_app", connector_app)
    expected_error: Final = (
        HTTPException if failure_phase == "license"
        else asyncio.CancelledError if failure_phase == "cancelled"
        else RuntimeError
    )
    message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed"
    async def run_lifespan() -> None:
        async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
            assert state == {"tracing_receiver": None}
            if failure_phase == "serving":
                raise RuntimeError("serving failed")
            if failure_phase == "cancelled":
                raise asyncio.CancelledError("cancelled failed")

    with pytest.raises(expected_error, match=message):
        await run_lifespan()

    assert proxy_server.shared_aiohttp_session is not None
    assert proxy_server.shared_aiohttp_session.closed
    assert proxy_server.master_key is None
    assert all(route.name != "admin_mcp" for route in admin_mcp_proxy.routes)


@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector(
    admin_mcp_proxy: FastAPI, monkeypatch: pytest.MonkeyPatch,
) -> None:
    import httpx2

    from litellm.proxy.middleware.in_flight_requests_middleware import InFlightRequestsMiddleware
    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager

    monkeypatch.setenv("LITELLM_ADMIN_TOOLS", "list_teams")
    admin_mcp_proxy.add_middleware(InFlightRequestsMiddleware)
    started: Final = asyncio.Event()
    release: Final = asyncio.Event()

    @admin_mcp_proxy.get("/user/info")
    async def user_info() -> dict[str, object]:
        return {"user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"}}

    @admin_mcp_proxy.get("/team/list", operation_id="list_team_team_list_get")
    async def list_teams() -> dict[str, object]:
        started.set()
        await release.wait()
        return {"teams": ["completed-before-shutdown"]}

    async def complete_during_drain() -> None:
        async with asyncio.timeout(5):
            while not GracefulShutdownManager.is_shutting_down():
                await asyncio.sleep(0)
        release.set()

    async with httpx2.AsyncClient(
        transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000"
    ) as client:
        async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state:
            assert state == {"tracing_receiver": None}
            request: Final = asyncio.create_task(client.post(
                "/admin/mcp",
                headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"},
                json={"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "list_teams"}},
            ))
            await asyncio.wait_for(started.wait(), timeout=5)
            completion: Final = asyncio.create_task(complete_during_drain())
        await asyncio.wait_for(completion, timeout=5)
        response: Final = await asyncio.wait_for(request, timeout=5)

    assert response.status_code == 200, response.text
    assert response.json()["result"]["isError"] is False
    assert json.loads(response.json()["result"]["content"][0]["text"]) == {
        "teams": ["completed-before-shutdown"]
    }
    assert InFlightRequestsMiddleware.get_count() == 0


@pytest.mark.skipif(sys.version_info < (3, 12), reason="Admin MCP requires Python 3.12+")
async def test_proxy_shutdown_closes_admin_connector_when_drain_is_cancelled(
    admin_mcp_proxy: FastAPI,
) -> None:
    import httpx2

    from litellm.proxy import proxy_server
    from litellm.proxy.middleware.in_flight_requests_middleware import InFlightRequestsMiddleware
    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager

    admin_mcp_proxy.add_middleware(InFlightRequestsMiddleware)
    ready: Final = asyncio.Event()
    shutdown: Final = asyncio.Event()
    started: Final = asyncio.Event()
    release: Final = asyncio.Event()

    @admin_mcp_proxy.get("/hold")
    async def hold_request() -> dict[str, bool]:
        started.set()
        await release.wait()
        return {"complete": True}

    async def serve() -> None:
        async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy):
            ready.set()
            await shutdown.wait()

    async with httpx2.AsyncClient(
        transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000"
    ) as client:
        serving: Final = asyncio.create_task(serve())
        await asyncio.wait_for(ready.wait(), timeout=5)
        assert any(route.name == "admin_mcp" for route in admin_mcp_proxy.routes)
        request: Final = asyncio.create_task(client.get("/hold"))
        try:
            await asyncio.wait_for(started.wait(), timeout=5)
            shutdown.set()
            async with asyncio.timeout(5):
                while not GracefulShutdownManager.is_shutting_down():
                    await asyncio.sleep(0)
            serving.cancel()
            with pytest.raises(asyncio.CancelledError):
                await serving
            assert all(route.name != "admin_mcp" for route in admin_mcp_proxy.routes)
            assert proxy_server.shared_aiohttp_session is not None
            assert proxy_server.shared_aiohttp_session.closed
        finally:
            release.set()
            await asyncio.wait_for(request, timeout=5)

    assert InFlightRequestsMiddleware.get_count() == 0


# Your bearer token
token = MASTER_KEY

headers = {"Authorization": f"Bearer {token}"}

example_completion_result = {
    "choices": [
        {
            "message": {
                "content": "Whispers of the wind carry dreams to me.",
                "role": "assistant",
            }
        }
    ],
}
example_embedding_result = {
    "object": "list",
    "data": [
        {
            "object": "embedding",
            "index": 0,
            "embedding": [
                -0.006929283495992422,
                -0.005336422007530928,
                -4.547132266452536e-05,
                -0.024047505110502243,
                -0.006929283495992422,
                -0.005336422007530928,
                -4.547132266452536e-05,
                -0.024047505110502243,
                -0.006929283495992422,
                -0.005336422007530928,
                -4.547132266452536e-05,
                -0.024047505110502243,
            ],
        }
    ],
    "model": "text-embedding-3-small",
    "usage": {"prompt_tokens": 5, "total_tokens": 5},
}
example_image_generation_result = {
    "created": 1589478378,
    "data": [{"url": "https://..."}, {"url": "https://..."}],
}


def mock_patch_acompletion():
    return mock.patch(
        "litellm.proxy.proxy_server.llm_router.acompletion",
        return_value=example_completion_result,
    )


def mock_patch_aembedding():
    return mock.patch(
        "litellm.proxy.proxy_server.llm_router.aembedding",
        return_value=example_embedding_result,
    )


def mock_patch_aimage_generation():
    return mock.patch(
        "litellm.proxy.proxy_server.llm_router.aimage_generation",
        return_value=example_image_generation_result,
    )


@pytest.fixture(scope="function")
def fake_env_vars(monkeypatch):
    # Set some fake environment variables
    monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key")
    monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base")
    monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base")
    monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key")
    monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base")
    monkeypatch.setenv("REDIS_HOST", "localhost")


@pytest.fixture(scope="function")
def client_no_auth(fake_env_vars):
    # Assuming litellm.proxy.proxy_server is an object
    from litellm.proxy.proxy_server import cleanup_router_config_variables

    cleanup_router_config_variables()
    filepath = os.path.dirname(os.path.abspath(__file__))
    config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
    # initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables
    asyncio.run(initialize(config=config_fp, debug=True))
    return TestClient(app)


@mock_patch_acompletion()
def test_chat_completion(mock_acompletion, client_no_auth):
    global headers
    try:
        # Your test data
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": "hi"},
            ],
            "max_tokens": 10,
        }

        print("testing proxy server with chat completions")
        response = client_no_auth.post("/v1/chat/completions", json=test_data)
        mock_acompletion.assert_called_once_with(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "user", "content": "hi"},
            ],
            max_tokens=10,
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            specific_deployment=True,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        print(f"response - {response.text}")
        assert response.status_code == 200
        result = response.json()
        print(f"Received response: {result}")
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


def test_chat_completion_malformed_messages_returns_400(client_no_auth):
    """
    Test that malformed messages (strings instead of dicts) return 400 instead of 500.

    This test verifies that when a client sends messages as raw strings instead of
    {role, content} objects, LiteLLM returns a 400 invalid_request_error instead
    of a 500 Internal Server Error.
    """
    global headers
    try:
        # Test data with malformed messages (string instead of dict)
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                "hi how are you"
            ],  # Invalid: should be [{"role": "user", "content": "hi how are you"}]
        }

        print("testing proxy server with malformed messages")
        response = client_no_auth.post(
            "/v1/chat/completions", json=test_data, headers=headers
        )

        print(f"response status: {response.status_code}")
        print(f"response text: {response.text}")

        # Should return 400, not 500
        assert (
            response.status_code == 400
        ), f"Expected 400, got {response.status_code}. Response: {response.text}"

        # Verify error format
        result = response.json()
        assert "error" in result, "Response should contain 'error' key"
        error = result["error"]

        # Verify error type and message
        assert (
            error.get("type") == "invalid_request_error" or error.get("type") is None
        ), f"Expected invalid_request_error or None, got {error.get('type')}"
        assert (
            error.get("code") == "400" or error.get("code") == 400
        ), f"Expected code 400, got {error.get('code')}"

        # Error message should indicate invalid request format
        error_message = error.get("message", "")
        assert len(error_message) > 0, "Error message should not be empty"

    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


def test_get_settings_request_timeout(client_no_auth):
    """
    When no timeout is set, it should use the litellm.request_timeout value
    """
    # Set a known value for litellm.request_timeout
    import litellm

    # Make a GET request to /settings
    response = client_no_auth.get("/settings")

    # Check if the request was successful
    assert response.status_code == 200

    # Parse the JSON response
    settings = response.json()
    print("settings", settings)

    assert settings["litellm.request_timeout"] == litellm.request_timeout


@pytest.mark.parametrize(
    "litellm_key_header_name",
    ["x-litellm-key", None],
)
def test_add_headers_to_request(litellm_key_header_name):
    from fastapi import Request
    from starlette.datastructures import URL
    import json
    from litellm.proxy.litellm_pre_call_utils import (
        clean_headers,
        LiteLLMProxyRequestSetup,
    )

    headers = {
        "Authorization": "Bearer 1234",
        "X-Custom-Header": "Custom-Value",
        "X-Stainless-Header": "Stainless-Value",
        "anthropic-beta": "beta-value",
    }
    request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []})
    request._url = URL(url="/chat/completions")
    request._body = json.dumps({"model": "gpt-3.5-turbo"}).encode("utf-8")
    request_headers = clean_headers(headers, litellm_key_header_name)
    forwarded_headers = LiteLLMProxyRequestSetup._get_forwardable_headers(
        request_headers
    )
    assert forwarded_headers == {
        "X-Custom-Header": "Custom-Value",
        "anthropic-beta": "beta-value",
    }


@pytest.mark.parametrize(
    "litellm_key_header_name",
    ["x-litellm-key", None],
)
@pytest.mark.parametrize(
    "forward_headers",
    [True, False],
)
@mock_patch_acompletion()
def test_chat_completion_forward_headers(
    mock_acompletion, client_no_auth, litellm_key_header_name, forward_headers
):
    global headers
    try:
        if forward_headers:
            gs = getattr(litellm.proxy.proxy_server, "general_settings")
            gs["forward_client_headers_to_llm_api"] = True
            setattr(litellm.proxy.proxy_server, "general_settings", gs)
        if litellm_key_header_name is not None:
            gs = getattr(litellm.proxy.proxy_server, "general_settings")
            gs["litellm_key_header_name"] = litellm_key_header_name
            setattr(litellm.proxy.proxy_server, "general_settings", gs)
        # Your test data
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": "hi"},
            ],
            "max_tokens": 10,
        }

        headers_to_forward = {
            "X-Custom-Header": "Custom-Value",
            "X-Another-Header": "Another-Value",
        }

        if litellm_key_header_name is not None:
            headers_to_not_forward = {litellm_key_header_name: "Bearer 1234"}
        else:
            headers_to_not_forward = {"Authorization": "Bearer 1234"}

        received_headers = {**headers_to_forward, **headers_to_not_forward}

        print("testing proxy server with chat completions")
        response = client_no_auth.post(
            "/v1/chat/completions", json=test_data, headers=received_headers
        )
        if not forward_headers:
            assert "headers" not in mock_acompletion.call_args.kwargs
        else:
            assert mock_acompletion.call_args.kwargs["headers"] == {
                "x-custom-header": "Custom-Value",
                "x-another-header": "Another-Value",
            }

        print(f"response - {response.text}")
        assert response.status_code == 200
        result = response.json()
        print(f"Received response: {result}")
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


@pytest.mark.parametrize("forward_llm_auth_headers", [True, False])
@mock_patch_acompletion()
def test_chat_completion_forward_llm_provider_auth_headers(
    mock_acompletion, client_no_auth, forward_llm_auth_headers
):
    """
    Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded
    when forward_llm_provider_auth_headers=True.

    This allows clients to send their own LLM provider API keys through the proxy.
    """
    try:
        # Configure general settings
        gs = getattr(litellm.proxy.proxy_server, "general_settings")
        gs["forward_client_headers_to_llm_api"] = True
        gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers
        setattr(litellm.proxy.proxy_server, "general_settings", gs)

        # Test data
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": "hello"},
            ],
            "max_tokens": 10,
        }

        # Headers including LLM provider auth
        request_headers = {
            "Authorization": "Bearer sk-proxy-auth-123",  # Proxy auth (should be stripped)
            "x-api-key": "sk-ant-api03-test-anthropic-key",  # Anthropic API key
            "x-goog-api-key": "google-api-key-123",  # Google API key
            "X-Custom-Header": "custom-value",  # Custom header (should be forwarded)
        }

        # Make request
        response = client_no_auth.post(
            "/v1/chat/completions", json=test_data, headers=request_headers
        )

        assert response.status_code == 200

        # Check forwarded headers
        forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {})

        if forward_llm_auth_headers:
            # LLM provider auth headers should be forwarded
            assert "x-api-key" in forwarded_headers
            assert forwarded_headers["x-api-key"] == "sk-ant-api03-test-anthropic-key"
            assert "x-goog-api-key" in forwarded_headers
            assert forwarded_headers["x-goog-api-key"] == "google-api-key-123"
        else:
            # LLM provider auth headers should be stripped
            assert "x-api-key" not in forwarded_headers
            assert "x-goog-api-key" not in forwarded_headers

        # Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True)
        assert "x-custom-header" in forwarded_headers
        assert forwarded_headers["x-custom-header"] == "custom-value"

        # Proxy Authorization should never be forwarded
        assert "authorization" not in forwarded_headers

        print(
            f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}"
        )
        print(f"  Forwarded headers: {list(forwarded_headers.keys())}")

    except Exception as e:
        pytest.fail(
            f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}"
        )
    finally:
        # Clean up
        gs = getattr(litellm.proxy.proxy_server, "general_settings")
        gs.pop("forward_llm_provider_auth_headers", None)
        setattr(litellm.proxy.proxy_server, "general_settings", gs)


@mock_patch_acompletion()
@pytest.mark.asyncio
async def test_team_disable_guardrails(mock_acompletion, client_no_auth, monkeypatch):
    """
    If team not allowed to turn on/off guardrails

    Raise 403 forbidden error, if request is made by team on `/key/generate` or `/chat/completions`.
    """
    import asyncio
    import json
    import time

    from fastapi import HTTPException, Request
    from starlette.datastructures import URL

    from litellm.proxy._types import (
        LiteLLM_TeamTable,
        LiteLLM_TeamTableCachedObj,
        ProxyException,
        UserAPIKeyAuth,
    )
    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
    from litellm.proxy.proxy_server import hash_token

    user_api_key_cache: Final = UserApiKeyCache()
    _team_id = "1234"
    user_key = "sk-98765678"

    valid_token = UserAPIKeyAuth(
        team_id=_team_id,
        team_blocked=True,
        token=hash_token(user_key),
        last_refreshed_at=time.time(),
    )
    await asyncio.sleep(1)
    team_obj = LiteLLM_TeamTableCachedObj(
        team_id=_team_id,
        blocked=False,
        last_refreshed_at=time.time(),
        metadata={"guardrails": {"modify_guardrails": False}},
    )
    user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
    user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)

    monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world")

    request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions", "headers": []})
    request._url = URL(url="/chat/completions")

    body = {"metadata": {"guardrails": {"hide_secrets": False}}}
    json_bytes = json.dumps(body).encode("utf-8")

    request._body = json_bytes

    with pytest.raises(ProxyException) as exc_info:
        await user_api_key_auth(request=request, api_key="Bearer " + user_key)
    e = exc_info.value
    assert e.code == str(403)


from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler


@mock_patch_acompletion()
def test_custom_logger_failure_handler(mock_acompletion, client_no_auth, monkeypatch):
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.proxy.proxy_server import hash_token

    user_api_key_cache: Final = UserApiKeyCache()
    rpm_limit = 0

    mock_api_key = "sk-my-test-key"
    cache_value = UserAPIKeyAuth(token=hash_token(mock_api_key), rpm_limit=rpm_limit)

    user_api_key_cache.set_cache(key=hash_token(mock_api_key), value=cache_value)

    mock_logger = CustomLogger()
    mock_logger_unit_tests = CompletionCustomHandler()
    proxy_logging_obj: ProxyLogging = getattr(
        litellm.proxy.proxy_server, "proxy_logging_obj"
    )

    litellm.callbacks = [mock_logger, mock_logger_unit_tests]
    proxy_logging_obj._init_litellm_callbacks(llm_router=None)

    monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    setattr(litellm.proxy.proxy_server, "prisma_client", "FAKE-VAR")
    setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)

    with patch.object(
        mock_logger, "async_log_failure_event", new=AsyncMock()
    ) as mock_failed_alert:
        # Your test data
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": "hi"},
            ],
            "max_tokens": 10,
        }

        print("testing proxy server with chat completions")
        response = client_no_auth.post(
            "/v1/chat/completions",
            json=test_data,
            headers={"Authorization": "Bearer {}".format(mock_api_key)},
        )
        assert response.status_code == 429

        # confirm async_log_failure_event is called
        mock_failed_alert.assert_called()

        assert len(mock_logger_unit_tests.errors) == 0


@mock_patch_acompletion()
def test_engines_model_chat_completions(mock_acompletion, client_no_auth):
    global headers
    try:
        # Your test data
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": "hi"},
            ],
            "max_tokens": 10,
        }

        print("testing proxy server with chat completions")
        response = client_no_auth.post(
            "/engines/gpt-3.5-turbo/chat/completions", json=test_data
        )
        mock_acompletion.assert_called_once_with(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "user", "content": "hi"},
            ],
            max_tokens=10,
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            specific_deployment=True,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        print(f"response - {response.text}")
        assert response.status_code == 200
        result = response.json()
        print(f"Received response: {result}")
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


@mock_patch_acompletion()
def test_chat_completion_azure(mock_acompletion, client_no_auth):
    global headers
    try:
        # Your test data
        test_data = {
            "model": "azure/gpt-4.1-mini",
            "messages": [
                {"role": "user", "content": "write 1 sentence poem"},
            ],
            "max_tokens": 10,
        }

        print("testing proxy server with Azure Request /chat/completions")
        response = client_no_auth.post("/v1/chat/completions", json=test_data)

        mock_acompletion.assert_called_once_with(
            model="azure/gpt-4.1-mini",
            messages=[
                {"role": "user", "content": "write 1 sentence poem"},
            ],
            max_tokens=10,
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            specific_deployment=True,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        assert response.status_code == 200
        result = response.json()
        print(f"Received response: {result}")
        assert len(result["choices"][0]["message"]["content"]) > 0
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


# Run the test
# test_chat_completion_azure()


@mock_patch_acompletion()
def test_openai_deployments_model_chat_completions_azure(
    mock_acompletion, client_no_auth
):
    global headers
    try:
        # Your test data
        test_data = {
            "model": "azure/gpt-4.1-mini",
            "messages": [
                {"role": "user", "content": "write 1 sentence poem"},
            ],
            "max_tokens": 10,
        }

        url = "/openai/deployments/azure/gpt-4.1-mini/chat/completions"
        print(f"testing proxy server with Azure Request {url}")
        response = client_no_auth.post(url, json=test_data)

        mock_acompletion.assert_called_once_with(
            model="azure/gpt-4.1-mini",
            messages=[
                {"role": "user", "content": "write 1 sentence poem"},
            ],
            max_tokens=10,
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            specific_deployment=True,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        assert response.status_code == 200
        result = response.json()
        print(f"Received response: {result}")
        assert len(result["choices"][0]["message"]["content"]) > 0
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


# Run the test
# test_openai_deployments_model_chat_completions_azure()


### EMBEDDING
@mock_patch_aembedding()
def test_embedding(mock_aembedding, client_no_auth):
    global headers
    from litellm.proxy.proxy_server import user_custom_auth

    try:
        test_data = {
            "model": "azure/text-embedding-ada-002",
            "input": ["good morning from litellm"],
        }

        async def _pre_call_hook_side_effect(**kwargs):
            data = kwargs["data"]
            metadata = {**(data.get("metadata") or {}), "source": "unit-test"}
            data["metadata"] = metadata
            proxy_request = {**(data.get("proxy_server_request") or {})}
            proxy_request["path"] = "/v1/embeddings"
            data["proxy_server_request"] = proxy_request
            return data

        async def _post_call_success_side_effect(**kwargs):
            return kwargs["response"]

        with (
            patch.object(
                litellm.proxy.proxy_server.proxy_logging_obj,
                "pre_call_hook",
                new=AsyncMock(side_effect=_pre_call_hook_side_effect),
            ) as mock_pre_call_hook,
            patch.object(
                litellm.proxy.proxy_server.proxy_logging_obj,
                "during_call_hook",
                new=AsyncMock(return_value=None),
            ) as mock_during_hook,
            patch.object(
                litellm.proxy.proxy_server.proxy_logging_obj,
                "post_call_success_hook",
                new=AsyncMock(side_effect=_post_call_success_side_effect),
            ),
        ):
            response = client_no_auth.post("/v1/embeddings", json=test_data)

        mock_aembedding.assert_called_once_with(
            model="azure/text-embedding-ada-002",
            input=["good morning from litellm"],
            specific_deployment=True,
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        assert response.status_code == 200
        result = response.json()
        print(len(result["data"][0]["embedding"]))
        assert len(result["data"][0]["embedding"]) > 10  # this usually has len==1536 so

        call_metadata = mock_aembedding.call_args.kwargs["metadata"]
        assert call_metadata.get("source") == "unit-test"

        pre_call_kwargs = mock_pre_call_hook.await_args_list[0].kwargs
        assert (
            pre_call_kwargs.get("call_type") == "aembedding"
        ), f"expected pre_call_hook to receive call_type='aembedding', got {pre_call_kwargs.get('call_type')}"

    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


@mock_patch_aembedding()
def test_bedrock_embedding(mock_aembedding, client_no_auth):
    global headers
    from litellm.proxy.proxy_server import user_custom_auth

    try:
        test_data = {
            "model": "amazon-embeddings",
            "input": ["good morning from litellm"],
        }

        response = client_no_auth.post("/v1/embeddings", json=test_data)

        mock_aembedding.assert_called_once_with(
            model="amazon-embeddings",
            input=["good morning from litellm"],
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        assert response.status_code == 200
        print(response.status_code, response.text)
        result = response.json()
        print(len(result["data"][0]["embedding"]))
        assert len(result["data"][0]["embedding"]) > 10  # this usually has len==1536 so
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


@pytest.mark.skip(reason="AWS Suspended Account")
def test_sagemaker_embedding(client_no_auth):
    global headers
    from litellm.proxy.proxy_server import user_custom_auth

    try:
        test_data = {
            "model": "GPT-J 6B - Sagemaker Text Embedding (Internal)",
            "input": ["good morning from litellm"],
        }

        response = client_no_auth.post("/v1/embeddings", json=test_data)

        assert response.status_code == 200
        result = response.json()
        print(len(result["data"][0]["embedding"]))
        assert len(result["data"][0]["embedding"]) > 10  # this usually has len==1536 so
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


# Run the test
# test_embedding()
#### IMAGE GENERATION


@mock_patch_aimage_generation()
def test_img_gen(mock_aimage_generation, client_no_auth):
    global headers
    from litellm.proxy.proxy_server import user_custom_auth

    try:
        test_data = {
            "model": "dall-e-3",
            "prompt": "A cute baby sea otter",
            "n": 1,
            "size": "1024x1024",
            "imageConfig": {"aspectRatio": "9:16", "imageSize": "1K"},
        }

        response = client_no_auth.post("/v1/images/generations", json=test_data)

        mock_aimage_generation.assert_called_once_with(
            model="dall-e-3",
            prompt="A cute baby sea otter",
            n=1,
            size="1024x1024",
            imageConfig={"aspectRatio": "9:16", "imageSize": "1K"},
            litellm_call_id=mock.ANY,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        assert response.status_code == 200
        result = response.json()
        print(len(result["data"][0]["url"]))
        assert len(result["data"][0]["url"]) > 10
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


#### ADDITIONAL
@pytest.mark.skip(reason="test via docker tests. Requires prisma client.")
def test_add_new_model(client_no_auth):
    global headers
    try:
        test_data = {
            "model_name": "test_openai_models",
            "litellm_params": {
                "model": "gpt-3.5-turbo",
            },
            "model_info": {"description": "this is a test openai model"},
        }
        client_no_auth.post("/model/new", json=test_data, headers=headers)
        response = client_no_auth.get("/model/info", headers=headers)
        assert response.status_code == 200
        result = response.json()
        print(f"response: {result}")
        model_info = None
        for m in result["data"]:
            if m["model_name"] == "test_openai_models":
                model_info = m["model_info"]
        assert model_info["description"] == "this is a test openai model"
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")


@pytest.mark.xdist_group("proxy_heavy")
def test_health(client_no_auth):
    global headers
    import logging
    import time

    from litellm._logging import verbose_logger, verbose_proxy_logger

    verbose_proxy_logger.setLevel(logging.DEBUG)

    try:
        response = client_no_auth.get("/health")
        assert response.status_code == 200
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


# test_add_new_model()



class MyCustomHandler(CustomLogger):
    def log_pre_api_call(self, model, messages, kwargs):
        print(f"Pre-API Call")

    def log_success_event(self, kwargs, response_obj, start_time, end_time):
        print(f"On Success")
        assert kwargs["user"] == "proxy-user"
        assert kwargs["model"] == "gpt-3.5-turbo"
        assert kwargs["max_tokens"] == 10


customHandler = MyCustomHandler()


@mock_patch_acompletion()
def test_chat_completion_optional_params(mock_acompletion, client_no_auth):
    # [PROXY: PROD TEST] - DO NOT DELETE
    # This tests if all the /chat/completion params are passed to litellm
    try:
        # Your test data
        litellm.set_verbose = True
        test_data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {"role": "user", "content": "hi"},
            ],
            "max_tokens": 10,
            "user": "proxy-user",
        }

        litellm.callbacks = [customHandler]
        print("testing proxy server: optional params")
        response = client_no_auth.post("/v1/chat/completions", json=test_data)
        mock_acompletion.assert_called_once_with(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "user", "content": "hi"},
            ],
            max_tokens=10,
            user="proxy-user",
            litellm_call_id=mock.ANY,
            litellm_logging_obj=mock.ANY,
            request_timeout=mock.ANY,
            specific_deployment=True,
            metadata=mock.ANY,
            proxy_server_request=mock.ANY,
            secret_fields=mock.ANY,
        )
        assert response.status_code == 200
        result = response.json()
        print(f"Received response: {result}")
    except Exception as e:
        pytest.fail("LiteLLM Proxy test failed. Exception", e)


# Run the test
# test_chat_completion_optional_params()


# Test Reading config.yaml file
from litellm.proxy.proxy_server import ProxyConfig


@pytest.mark.skip(reason="local variable conflicts. needs to be refactored.")
@mock.patch("litellm.proxy.proxy_server.litellm.Cache")
def test_load_router_config(mock_cache, fake_env_vars):
    mock_cache.return_value.cache.__dict__ = {"redis_client": None}
    mock_cache.return_value.supported_call_types = [
        "completion",
        "acompletion",
        "embedding",
        "aembedding",
        "atranscription",
        "transcription",
    ]

    try:
        import asyncio

        print("testing reading config")
        # this is a basic config.yaml with only a model
        filepath = os.path.dirname(os.path.abspath(__file__))
        proxy_config = ProxyConfig()
        result = asyncio.run(
            proxy_config.load_config(
                router=None,
                config_file_path=f"{filepath}/example_config_yaml/simple_config.yaml",
            )
        )
        print(result)
        assert len(result[1]) == 1

        # this is a load balancing config yaml
        result = asyncio.run(
            proxy_config.load_config(
                router=None,
                config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml",
            )
        )
        print(result)
        assert len(result[1]) == 2

        # config with general settings - custom callbacks
        result = asyncio.run(
            proxy_config.load_config(
                router=None,
                config_file_path=f"{filepath}/example_config_yaml/azure_config.yaml",
            )
        )
        print(result)
        assert len(result[1]) == 2

        # tests for litellm.cache set from config
        print("testing reading proxy config for cache")
        litellm.cache = None
        asyncio.run(
            proxy_config.load_config(
                router=None,
                config_file_path=f"{filepath}/example_config_yaml/cache_no_params.yaml",
            )
        )
        assert litellm.cache is not None
        assert "redis_client" in vars(
            litellm.cache.cache
        )  # it should default to redis on proxy
        assert litellm.cache.supported_call_types == [
            "completion",
            "acompletion",
            "embedding",
            "aembedding",
            "atranscription",
            "transcription",
        ]  # init with all call types

        litellm.disable_cache()

        print("testing reading proxy config for cache with params")
        mock_cache.return_value.supported_call_types = [
            "embedding",
            "aembedding",
        ]
        asyncio.run(
            proxy_config.load_config(
                router=None,
                config_file_path=f"{filepath}/example_config_yaml/cache_with_params.yaml",
            )
        )
        assert litellm.cache is not None
        print(litellm.cache)
        print(litellm.cache.supported_call_types)
        print(vars(litellm.cache.cache))
        assert "redis_client" in vars(
            litellm.cache.cache
        )  # it should default to redis on proxy
        assert litellm.cache.supported_call_types == [
            "embedding",
            "aembedding",
        ]  # init with all call types

    except Exception as e:
        pytest.fail(
            f"Proxy: Got exception reading config: {str(e)}\n{traceback.format_exc()}"
        )


# test_load_router_config()


@pytest.mark.asyncio
async def test_team_update_redis():
    """
    Tests if team update, updates the redis cache if set
    """
    from litellm.caching.caching import DualCache, RedisCache
    from litellm.proxy._types import LiteLLM_TeamTableCachedObj
    from litellm.proxy.auth.auth_checks import cache_team_object

    proxy_logging_obj: ProxyLogging = getattr(
        litellm.proxy.proxy_server, "proxy_logging_obj"
    )

    redis_cache = RedisCache(host="localhost")

    with patch.object(
        redis_cache,
        "async_set_cache",
        new=AsyncMock(),
    ) as mock_client:
        await cache_team_object(
            team_id="1234",
            team_table=LiteLLM_TeamTableCachedObj(team_id="1234"),
            user_api_key_cache=DualCache(redis_cache=redis_cache),
            proxy_logging_obj=proxy_logging_obj,
        )

        mock_client.assert_called()


@pytest.mark.asyncio
async def test_get_team_redis(client_no_auth):
    """
    Tests if get_team_object gets value from redis cache, if set
    """
    from litellm.caching.caching import DualCache, RedisCache
    from litellm.proxy.auth.auth_checks import get_team_object

    proxy_logging_obj: ProxyLogging = getattr(
        litellm.proxy.proxy_server, "proxy_logging_obj"
    )

    redis_cache = RedisCache()

    from fastapi import HTTPException

    with patch.object(
        redis_cache,
        "async_get_cache",
        new=AsyncMock(),
    ) as mock_client:
        try:
            await get_team_object(
                team_id="1234",
                user_api_key_cache=DualCache(redis_cache=redis_cache),
                parent_otel_span=None,
                proxy_logging_obj=proxy_logging_obj,
                prisma_client=AsyncMock(),
            )
        except HTTPException:
            pass

        mock_client.assert_called_once()


import random
from litellm._uuid import uuid
from unittest.mock import PropertyMock

from litellm.proxy._types import (
    LitellmUserRoles,
    NewUserRequest,
    TeamMemberAddRequest,
    UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.team_endpoints import team_member_add
from tests.unit.proxy.management_endpoints.test_key_generate_prisma import prisma_client
@pytest.fixture
def mock_prisma_client():
    client = MagicMock()
    client.connect = AsyncMock()
    client.disconnect = AsyncMock()
    return client


@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
@pytest.mark.parametrize(
    "user_role",
    [LitellmUserRoles.INTERNAL_USER.value, LitellmUserRoles.PROXY_ADMIN.value],
)
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_create_user_default_budget(prisma_client, user_role):  # noqa: F811  # pytest fixture, not a redefinition

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    setattr(litellm, "max_internal_user_budget", 10)
    setattr(litellm, "internal_user_budget_duration", "5m")
    await litellm.proxy.proxy_server.prisma_client.connect()
    user = f"ishaan {uuid.uuid4().hex}"
    request = NewUserRequest(
        user_id=user, user_role=user_role
    )  # create a key with no budget
    with patch.object(
        litellm.proxy.proxy_server.prisma_client, "insert_data", new=AsyncMock()
    ) as mock_client:
        await new_user(
            request,
        )

        mock_client.assert_called()

        print(f"mock_client.call_args: {mock_client.call_args}")
        print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs))

        if user_role == LitellmUserRoles.INTERNAL_USER.value:
            assert (
                mock_client.call_args.kwargs["data"]["max_budget"]
                == litellm.max_internal_user_budget
            )
            assert (
                mock_client.call_args.kwargs["data"]["budget_duration"]
                == litellm.internal_user_budget_duration
            )

        else:
            assert mock_client.call_args.kwargs["data"]["max_budget"] is None
            assert mock_client.call_args.kwargs["data"]["budget_duration"] is None


def _member_add_tx_cm(team_table):
    """Transaction whose member writes land on whatever tables are mocked on `prisma_client.db`"""

    class _Tx:
        query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
        litellm_teamtable = team_table

        def __getattr__(self, table_name):
            return getattr(litellm.proxy.proxy_server.prisma_client.db, table_name)

    tx_cm = MagicMock()
    tx_cm.__aenter__ = AsyncMock(return_value=_Tx())
    tx_cm.__aexit__ = AsyncMock(return_value=None)
    return tx_cm


@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_create_team_member_add(prisma_client, new_member_method):  # noqa: F811  # pytest fixture, not a redefinition
    import time

    from fastapi import Request

    from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable
    from litellm.proxy.proxy_server import hash_token, user_api_key_cache

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    setattr(litellm, "max_internal_user_budget", 10)
    setattr(litellm, "internal_user_budget_duration", "5m")
    await litellm.proxy.proxy_server.prisma_client.connect()
    user = f"ishaan {uuid.uuid4().hex}"
    _team_id = "litellm-test-client-id-new"
    team_obj = LiteLLM_TeamTableCachedObj(
        team_id=_team_id,
        blocked=False,
        last_refreshed_at=time.time(),
        metadata={"guardrails": {"modify_guardrails": False}},
    )
    # user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)
    user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)

    setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
    if new_member_method == "user_id":
        data = {
            "team_id": _team_id,
            "member": [{"role": "user", "user_id": user}],
        }
    elif new_member_method == "user_email":
        data = {
            "team_id": _team_id,
            "member": [{"role": "user", "user_email": user}],
        }
    team_member_add_request = TeamMemberAddRequest(**data)

    with (
        patch(
            "litellm.proxy.proxy_server.prisma_client.db.litellm_usertable",
            new_callable=AsyncMock,
        ) as mock_litellm_usertable,
        patch(
            "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache",
            new=AsyncMock(return_value=team_obj),
        ) as mock_team_obj,
        patch(
            "litellm.proxy.proxy_server.prisma_client.get_data",
            new=AsyncMock(return_value=[]),
        ) as mock_get_data,
    ):

        mock_client = AsyncMock(
            return_value=LiteLLM_UserTable(
                user_id="1234", max_budget=100, user_email="1234"
            )
        )
        mock_litellm_usertable.upsert = mock_client
        mock_litellm_usertable.find_many = AsyncMock(return_value=[])
        # Mock find_first for user_email validation (returns None for new users)
        mock_litellm_usertable.find_first = AsyncMock(return_value=None)
        # Mock find_unique for user_id validation (returns None for new users)
        mock_litellm_usertable.find_unique = AsyncMock(return_value=None)
        team_mock_client = AsyncMock()
        original_val = getattr(
            litellm.proxy.proxy_server.prisma_client.db, "litellm_teamtable"
        )
        litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = team_mock_client

        team_mock_client.update = AsyncMock(
            return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
        )

        tx_cm = _member_add_tx_cm(team_mock_client)
        original_tx = litellm.proxy.proxy_server.prisma_client.tx
        litellm.proxy.proxy_server.prisma_client.tx = MagicMock(
            return_value=tx_cm
        )

        print(f"team_member_add_request={team_member_add_request}")
        await team_member_add(
            data=team_member_add_request,
            user_api_key_dict=UserAPIKeyAuth(user_role="proxy_admin"),
        )

        mock_client.assert_called()

        print(f"mock_client.call_args: {mock_client.call_args}")
        print("mock_client.call_args.kwargs: {}".format(mock_client.call_args.kwargs))

        assert (
            mock_client.call_args.kwargs["data"]["create"]["max_budget"]
            == litellm.max_internal_user_budget
        )
        assert (
            mock_client.call_args.kwargs["data"]["create"]["budget_duration"]
            == litellm.internal_user_budget_duration
        )

        litellm.proxy.proxy_server.prisma_client.db.litellm_teamtable = original_val
        litellm.proxy.proxy_server.prisma_client.tx = original_tx


@pytest.mark.parametrize("team_member_role", ["admin", "user"])
@pytest.mark.parametrize("team_route", ["/team/member_add", "/team/member_delete"])
@pytest.mark.asyncio
async def test_create_team_member_add_team_admin_user_api_key_auth(
    prisma_client, team_member_role, team_route, monkeypatch  # noqa: F811  # pytest fixture, not a redefinition
):
    import time

    from fastapi import Request

    from litellm.proxy._types import LiteLLM_TeamTableCachedObj, Member
    from litellm.proxy.proxy_server import (
        ProxyException,
        hash_token,
        user_api_key_auth,
    )

    user_api_key_cache: Final = UserApiKeyCache()

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    setattr(litellm, "max_internal_user_budget", 10)
    setattr(litellm, "internal_user_budget_duration", "5m")
    user = f"ishaan {uuid.uuid4().hex}"
    _team_id = "litellm-test-client-id-new"
    user_key = "sk-98765678"

    valid_token = UserAPIKeyAuth(
        team_id=_team_id,
        token=hash_token(user_key),
        team_member=Member(role=team_member_role, user_id=user),
        last_refreshed_at=time.time(),
    )
    user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)

    team_obj = LiteLLM_TeamTableCachedObj(
        team_id=_team_id,
        blocked=False,
        last_refreshed_at=time.time(),
        metadata={"guardrails": {"modify_guardrails": False}},
    )

    user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)

    monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)

    ## TEST IF TEAM ADMIN ALLOWED TO CALL /MEMBER_ADD ENDPOINT
    import json

    from starlette.datastructures import URL

    request = Request(scope={"type": "http", "method": "POST", "path": team_route, "headers": []})
    request._url = URL(url=team_route)

    body = {}
    json_bytes = json.dumps(body).encode("utf-8")

    request._body = json_bytes

    ## ALLOWED BY USER_API_KEY_AUTH
    await user_api_key_auth(request=request, api_key="Bearer " + user_key)


@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
@pytest.mark.parametrize("user_role", ["admin", "user"])
@pytest.mark.asyncio
async def test_create_team_member_add_team_admin(
    prisma_client, new_member_method, user_role  # noqa: F811  # pytest fixture, not a redefinition
):
    """
    Relevant issue - https://github.com/BerriAI/litellm/issues/5300

    Allow team admins to:
        - Add and remove team members
        - raise error if team member not an existing 'internal_user'
    """
    import time

    from fastapi import Request

    from litellm.proxy._types import (
        LiteLLM_TeamMembership,
        LiteLLM_TeamTableCachedObj,
        LiteLLM_UserTable,
        Member,
    )
    from litellm.proxy.proxy_server import (
        HTTPException,
        ProxyException,
        hash_token,
        user_api_key_auth,
        user_api_key_cache,
    )

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    setattr(litellm, "max_internal_user_budget", 10)
    setattr(litellm, "internal_user_budget_duration", "5m")
    user = f"ishaan {uuid.uuid4().hex}"
    _team_id = "litellm-test-client-id-new"
    user_key = "sk-98765678"
    team_admin = f"krrish {uuid.uuid4().hex}"

    valid_token = UserAPIKeyAuth(
        team_id=_team_id,
        user_id=team_admin,
        token=hash_token(user_key),
        last_refreshed_at=time.time(),
    )
    user_api_key_cache.set_cache(key=hash_token(user_key), value=valid_token)

    team_obj = LiteLLM_TeamTableCachedObj(
        team_id=_team_id,
        blocked=False,
        last_refreshed_at=time.time(),
        members_with_roles=[Member(role=user_role, user_id=team_admin)],
        metadata={"guardrails": {"modify_guardrails": False}},
    )

    user_api_key_cache.set_cache(key="team_id:{}".format(_team_id), value=team_obj)

    setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
    if new_member_method == "user_id":
        data = {
            "team_id": _team_id,
            "member": [{"role": "user", "user_id": user}],
        }
    elif new_member_method == "user_email":
        data = {
            "team_id": _team_id,
            "member": [{"role": "user", "user_email": user}],
        }
    team_member_add_request = TeamMemberAddRequest(**data)

    with (
        patch(
            "litellm.proxy.proxy_server.prisma_client.db.litellm_usertable",
            new_callable=AsyncMock,
        ) as mock_litellm_usertable,
        patch(
            "litellm.proxy.auth.auth_checks._get_team_object_from_user_api_key_cache",
            new=AsyncMock(return_value=team_obj),
        ) as mock_team_obj,
        patch(
            "litellm.proxy.proxy_server.prisma_client.get_data",
            new=AsyncMock(return_value=[]),
        ) as mock_get_data,
    ):
        mock_client = AsyncMock(
            return_value=LiteLLM_UserTable(
                user_id="1234", max_budget=100, user_email="1234"
            )
        )
        mock_litellm_usertable.upsert = mock_client
        mock_litellm_usertable.find_many = AsyncMock(return_value=[])
        # Mock find_first for user_email validation (returns None for new users)
        mock_litellm_usertable.find_first = AsyncMock(return_value=None)
        # Mock find_unique for user_id validation (returns None for new users)
        mock_litellm_usertable.find_unique = AsyncMock(return_value=None)

        team_mock_client = AsyncMock()
        team_mock_client.update = AsyncMock(
            return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
        )
        membership_mock_client = AsyncMock()
        membership_mock_client.upsert = AsyncMock(
            return_value=LiteLLM_TeamMembership(user_id="1234", team_id=_team_id)
        )

        tx_cm = _member_add_tx_cm(team_mock_client)

        with (
            patch.object(
                litellm.proxy.proxy_server.prisma_client.db,
                "litellm_teamtable",
                team_mock_client,
            ),
            patch.object(  # test-quality-ok: legacy test swaps the prisma table on the module-level client
                litellm.proxy.proxy_server.prisma_client.db,
                "litellm_teammembership",
                membership_mock_client,
            ),
            patch.object(
                litellm.proxy.proxy_server.prisma_client,
                "tx",
                MagicMock(return_value=tx_cm),
            ),
        ):
            error = None
            try:
                await team_member_add(
                    data=team_member_add_request,
                    user_api_key_dict=valid_token,
                )
            except HTTPException as e:
                error = e

            if error is not None:
                assert user_role == "user" or new_member_method == "user_id"
                assert error.status_code == 403
                return

            mock_client.assert_called()

            assert (
                mock_client.call_args.kwargs["data"]["create"]["max_budget"]
                == litellm.max_internal_user_budget
            )
            assert (
                mock_client.call_args.kwargs["data"]["create"]["budget_duration"]
                == litellm.internal_user_budget_duration
            )


@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_user_info_team_list(prisma_client):  # noqa: F811  # pytest fixture, not a redefinition
    """Assert user_info for admin calls team_list function"""
    from litellm.proxy._types import LiteLLM_UserTable

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    await litellm.proxy.proxy_server.prisma_client.connect()

    from litellm.proxy.management_endpoints.internal_user_endpoints import user_info

    with patch(
        "litellm.proxy.management_endpoints.team_endpoints.list_team",
        new_callable=AsyncMock,
    ) as mock_client:

        prisma_client.get_data = AsyncMock(
            return_value=LiteLLM_UserTable(
                user_role="proxy_admin",
                user_id="default_user_id",
                max_budget=None,
                user_email="",
            )
        )

        try:
            await user_info(
                request=MagicMock(),
                user_id=None,
                user_api_key_dict=UserAPIKeyAuth(
                    api_key=MASTER_KEY, user_id="default_user_id"
                ),
            )
        except Exception:
            pass

        mock_client.assert_called()


@pytest.mark.skip(reason="Local test")
@pytest.mark.asyncio
async def test_add_callback_via_key(prisma_client):  # noqa: F811  # pytest fixture, not a redefinition
    """
    Test if callback specified in key, is used.
    """
    global headers
    import json

    from fastapi import HTTPException, Request, Response
    from starlette.datastructures import URL

    from litellm.proxy.proxy_server import chat_completion

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    await litellm.proxy.proxy_server.prisma_client.connect()

    litellm.set_verbose = True

    try:
        # Your test data
        test_data = {
            "model": "azure/gpt-4.1-mini",
            "messages": [
                {"role": "user", "content": "write 1 sentence poem"},
            ],
            "max_tokens": 10,
            "mock_response": "Hello world",
            "api_key": "my-fake-key",
        }

        request = Request(scope={"type": "http", "method": "POST", "headers": {}})
        request._url = URL(url="/chat/completions")

        json_bytes = json.dumps(test_data).encode("utf-8")

        request._body = json_bytes

        with patch.object(
            litellm.litellm_core_utils.litellm_logging,
            "LangFuseLogger",
            new=MagicMock(),
        ) as mock_client:
            resp = await chat_completion(
                request=request,
                fastapi_response=Response(),
                user_api_key_dict=UserAPIKeyAuth(
                    metadata={
                        "allow_client_mock_response": True,
                        "logging": [
                            {
                                "callback_name": "langfuse",  # 'otel', 'langfuse', 'lunary'
                                "callback_type": "success",  # set, if required by integration - future improvement, have logging tools work for success + failure by default
                                "callback_vars": {
                                    "langfuse_public_key": "os.environ/LANGFUSE_PUBLIC_KEY",
                                    "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY",
                                    "langfuse_host": "https://us.cloud.langfuse.com",
                                },
                            }
                        ],
                    }
                ),
            )
            print(resp)
            mock_client.assert_called()
            mock_client.return_value.log_event.assert_called()
            args, kwargs = mock_client.return_value.log_event.call_args
            kwargs = kwargs["kwargs"]
            assert "user_api_key_metadata" in kwargs["litellm_params"]["metadata"]
            assert (
                "logging"
                in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"]
            )
            checked_keys = False
            for item in kwargs["litellm_params"]["metadata"]["user_api_key_metadata"][
                "logging"
            ]:
                for k, v in item["callback_vars"].items():
                    print("k={}, v={}".format(k, v))
                    if "key" in k:
                        assert "os.environ" in v
                        checked_keys = True

            assert checked_keys
    except Exception as e:
        pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "callback_type, expected_success_callbacks, expected_failure_callbacks",
    [
        ("success", ["langfuse"], []),
        ("failure", [], ["langfuse"]),
        ("success_and_failure", ["langfuse"], ["langfuse"]),
    ],
)
async def test_add_callback_via_key_litellm_pre_call_utils(
    mock_prisma_client,
    callback_type,
    expected_success_callbacks,
    expected_failure_callbacks,
):
    import json

    from fastapi import HTTPException, Request, Response
    from starlette.datastructures import URL

    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request

    setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    request = Request(scope={"type": "http", "method": "POST", "headers": {}})
    request._url = URL(url="/chat/completions")

    test_data = {
        "model": "azure/gpt-4.1-mini",
        "messages": [
            {"role": "user", "content": "write 1 sentence poem"},
        ],
        "max_tokens": 10,
        "mock_response": "Hello world",
        "api_key": "my-fake-key",
    }

    json_bytes = json.dumps(test_data).encode("utf-8")

    request._body = json_bytes

    data = {
        "data": {
            "model": "azure/gpt-4.1-mini",
            "messages": [{"role": "user", "content": "write 1 sentence poem"}],
            "max_tokens": 10,
            "mock_response": "Hello world",
            "api_key": "my-fake-key",
        },
        "request": request,
        "user_api_key_dict": UserAPIKeyAuth(
            token=None,
            key_name=None,
            key_alias=None,
            spend=0.0,
            max_budget=None,
            expires=None,
            models=[],
            aliases={},
            config={},
            user_id=None,
            team_id=None,
            max_parallel_requests=None,
            metadata={
                "allow_client_mock_response": True,
                "logging": [
                    {
                        "callback_name": "langfuse",
                        "callback_type": callback_type,
                        "callback_vars": {
                            "langfuse_public_key": "my-mock-public-key",
                            "langfuse_secret_key": "my-mock-secret-key",
                            "langfuse_host": "https://us.cloud.langfuse.com",
                        },
                    }
                ],
            },
            tpm_limit=None,
            rpm_limit=None,
            budget_duration=None,
            budget_reset_at=None,
            allowed_cache_controls=[],
            permissions={},
            model_spend={},
            model_max_budget={},
            soft_budget_cooldown=False,
            litellm_budget_table=None,
            org_id=None,
            team_spend=None,
            team_alias=None,
            team_tpm_limit=None,
            team_rpm_limit=None,
            team_max_budget=None,
            team_models=[],
            team_blocked=False,
            soft_budget=None,
            team_model_aliases=None,
            team_member_spend=None,
            team_metadata=None,
            end_user_id=None,
            end_user_tpm_limit=None,
            end_user_rpm_limit=None,
            end_user_max_budget=None,
            last_refreshed_at=None,
            api_key=None,
            user_role=None,
            allowed_model_region=None,
            parent_otel_span=None,
        ),
        "proxy_config": proxy_config,
        "general_settings": {},
        "version": "0.0.0",
    }

    new_data = await add_litellm_data_to_request(**data)
    print("NEW DATA: {}".format(new_data))

    assert "langfuse_public_key" in new_data
    assert new_data["langfuse_public_key"] == "my-mock-public-key"
    assert "langfuse_secret_key" in new_data
    assert new_data["langfuse_secret_key"] == "my-mock-secret-key"

    if expected_success_callbacks:
        assert "success_callback" in new_data
        assert new_data["success_callback"] == expected_success_callbacks

    if expected_failure_callbacks:
        assert "failure_callback" in new_data
        assert new_data["failure_callback"] == expected_failure_callbacks


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "disable_fallbacks_set",
    [
        True,
        False,
    ],
)
async def test_disable_fallbacks_by_key(disable_fallbacks_set):
    from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup

    key_metadata = {"disable_fallbacks": disable_fallbacks_set}
    existing_data = {
        "model": "azure/gpt-4.1-mini",
        "messages": [{"role": "user", "content": "write 1 sentence poem"}],
    }
    data = LiteLLMProxyRequestSetup.add_key_level_controls(
        key_metadata=key_metadata,
        data=existing_data,
        _metadata_variable_name="metadata",
    )

    assert data["disable_fallbacks"] == disable_fallbacks_set


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "callback_type, expected_success_callbacks, expected_failure_callbacks",
    [
        ("success", ["gcs_bucket"], []),
        ("failure", [], ["gcs_bucket"]),
        ("success_and_failure", ["gcs_bucket"], ["gcs_bucket"]),
    ],
)
async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket(
    mock_prisma_client,
    callback_type,
    expected_success_callbacks,
    expected_failure_callbacks,
):
    import json

    from fastapi import HTTPException, Request, Response
    from starlette.datastructures import URL

    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request

    setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    request = Request(scope={"type": "http", "method": "POST", "headers": {}})
    request._url = URL(url="/chat/completions")

    test_data = {
        "model": "azure/gpt-4.1-mini",
        "messages": [
            {"role": "user", "content": "write 1 sentence poem"},
        ],
        "max_tokens": 10,
        "mock_response": "Hello world",
        "api_key": "my-fake-key",
    }

    json_bytes = json.dumps(test_data).encode("utf-8")

    request._body = json_bytes

    data = {
        "data": {
            "model": "azure/gpt-4.1-mini",
            "messages": [{"role": "user", "content": "write 1 sentence poem"}],
            "max_tokens": 10,
            "mock_response": "Hello world",
            "api_key": "my-fake-key",
        },
        "request": request,
        "user_api_key_dict": UserAPIKeyAuth(
            token=None,
            key_name=None,
            key_alias=None,
            spend=0.0,
            max_budget=None,
            expires=None,
            models=[],
            aliases={},
            config={},
            user_id=None,
            team_id=None,
            max_parallel_requests=None,
            metadata={
                "allow_client_mock_response": True,
                "logging": [
                    {
                        "callback_name": "gcs_bucket",
                        "callback_type": callback_type,
                        "callback_vars": {
                            "gcs_bucket_name": "key-logging-project1",
                            "gcs_path_service_account": "pathrise-convert-1606954137718-a956eef1a2a8.json",
                        },
                    }
                ],
            },
            tpm_limit=None,
            rpm_limit=None,
            budget_duration=None,
            budget_reset_at=None,
            allowed_cache_controls=[],
            permissions={},
            model_spend={},
            model_max_budget={},
            soft_budget_cooldown=False,
            litellm_budget_table=None,
            org_id=None,
            team_spend=None,
            team_alias=None,
            team_tpm_limit=None,
            team_rpm_limit=None,
            team_max_budget=None,
            team_models=[],
            team_blocked=False,
            soft_budget=None,
            team_model_aliases=None,
            team_member_spend=None,
            team_metadata=None,
            end_user_id=None,
            end_user_tpm_limit=None,
            end_user_rpm_limit=None,
            end_user_max_budget=None,
            last_refreshed_at=None,
            api_key=None,
            user_role=None,
            allowed_model_region=None,
            parent_otel_span=None,
        ),
        "proxy_config": proxy_config,
        "general_settings": {},
        "version": "0.0.0",
    }

    new_data = await add_litellm_data_to_request(**data)
    print("NEW DATA: {}".format(new_data))

    assert "gcs_bucket_name" in new_data
    assert new_data["gcs_bucket_name"] == "key-logging-project1"
    assert "gcs_path_service_account" in new_data
    assert (
        new_data["gcs_path_service_account"]
        == "pathrise-convert-1606954137718-a956eef1a2a8.json"
    )

    if expected_success_callbacks:
        assert "success_callback" in new_data
        assert new_data["success_callback"] == expected_success_callbacks

    if expected_failure_callbacks:
        assert "failure_callback" in new_data
        assert new_data["failure_callback"] == expected_failure_callbacks


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "callback_type, expected_success_callbacks, expected_failure_callbacks",
    [
        ("success", ["langsmith"], []),
        ("failure", [], ["langsmith"]),
        ("success_and_failure", ["langsmith"], ["langsmith"]),
    ],
)
async def test_add_callback_via_key_litellm_pre_call_utils_langsmith(
    mock_prisma_client,
    callback_type,
    expected_success_callbacks,
    expected_failure_callbacks,
):
    import json

    from fastapi import HTTPException, Request, Response
    from starlette.datastructures import URL

    from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request

    setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    request = Request(scope={"type": "http", "method": "POST", "headers": {}})
    request._url = URL(url="/chat/completions")

    test_data = {
        "model": "azure/gpt-4.1-mini",
        "messages": [
            {"role": "user", "content": "write 1 sentence poem"},
        ],
        "max_tokens": 10,
        "mock_response": "Hello world",
        "api_key": "my-fake-key",
    }

    json_bytes = json.dumps(test_data).encode("utf-8")

    request._body = json_bytes

    data = {
        "data": {
            "model": "azure/gpt-4.1-mini",
            "messages": [{"role": "user", "content": "write 1 sentence poem"}],
            "max_tokens": 10,
            "mock_response": "Hello world",
            "api_key": "my-fake-key",
        },
        "request": request,
        "user_api_key_dict": UserAPIKeyAuth(
            token=None,
            key_name=None,
            key_alias=None,
            spend=0.0,
            max_budget=None,
            expires=None,
            models=[],
            aliases={},
            config={},
            user_id=None,
            team_id=None,
            max_parallel_requests=None,
            metadata={
                "allow_client_mock_response": True,
                "logging": [
                    {
                        "callback_name": "langsmith",
                        "callback_type": callback_type,
                        "callback_vars": {
                            "langsmith_api_key": "ls-1234",
                            "langsmith_project": "pr-brief-resemblance-72",
                            "langsmith_base_url": "https://api.smith.langchain.com",
                        },
                    }
                ],
            },
            tpm_limit=None,
            rpm_limit=None,
            budget_duration=None,
            budget_reset_at=None,
            allowed_cache_controls=[],
            permissions={},
            model_spend={},
            model_max_budget={},
            soft_budget_cooldown=False,
            litellm_budget_table=None,
            org_id=None,
            team_spend=None,
            team_alias=None,
            team_tpm_limit=None,
            team_rpm_limit=None,
            team_max_budget=None,
            team_models=[],
            team_blocked=False,
            soft_budget=None,
            team_model_aliases=None,
            team_member_spend=None,
            team_metadata=None,
            end_user_id=None,
            end_user_tpm_limit=None,
            end_user_rpm_limit=None,
            end_user_max_budget=None,
            last_refreshed_at=None,
            api_key=None,
            user_role=None,
            allowed_model_region=None,
            parent_otel_span=None,
        ),
        "proxy_config": proxy_config,
        "general_settings": {},
        "version": "0.0.0",
    }

    new_data = await add_litellm_data_to_request(**data)
    print("NEW DATA: {}".format(new_data))

    assert "langsmith_api_key" in new_data
    assert new_data["langsmith_api_key"] == "ls-1234"
    assert "langsmith_project" in new_data
    assert new_data["langsmith_project"] == "pr-brief-resemblance-72"
    assert "langsmith_base_url" in new_data
    assert new_data["langsmith_base_url"] == "https://api.smith.langchain.com"

    if expected_success_callbacks:
        assert "success_callback" in new_data
        assert new_data["success_callback"] == expected_success_callbacks

    if expected_failure_callbacks:
        assert "failure_callback" in new_data
        assert new_data["failure_callback"] == expected_failure_callbacks


@pytest.mark.parametrize("hidden", [True, False])
@pytest.mark.asyncio
async def test_model_info_alias_without_prisma(hidden):
    from litellm.proxy.proxy_server import model_info_v1

    _model_list = [
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "gpt-3.5-turbo"},
        }
    ]

    model_alias = "gpt-4"

    router = litellm.Router(
        model_list=_model_list,
        model_group_alias={
            model_alias: {
                "model": "gpt-3.5-turbo",
                "hidden": hidden,
            }
        },
    )

    setattr(litellm.proxy.proxy_server, "llm_router", router)
    setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)
    setattr(litellm.proxy.proxy_server, "prisma_client", None)

    resp = await model_info_v1(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )

    models = json.loads(resp.body)["data"]

    alias_found = any(
        m["model_name"] == model_alias
        for m in models
    )

    assert alias_found is (not hidden)

    
@pytest.mark.parametrize("hidden", [True, False])
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_proxy_model_group_alias_checks(prisma_client, hidden):  # noqa: F811  # pytest fixture, not a redefinition
    """
    Check if model group alias is returned on

    `/v1/models`
    `/v1/model/info`
    `/v1/model_group/info`
    """
    import json

    from fastapi import HTTPException, Request, Response
    from starlette.datastructures import URL

    from litellm.proxy.proxy_server import model_group_info, model_info_v1, model_list

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    await litellm.proxy.proxy_server.prisma_client.connect()

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    _model_list = [
        {
            "model_name": "gpt-3.5-turbo",
            "litellm_params": {"model": "gpt-3.5-turbo"},
        }
    ]
    model_alias = "gpt-4"
    router = litellm.Router(
        model_list=_model_list,
        model_group_alias={model_alias: {"model": "gpt-3.5-turbo", "hidden": hidden}},
    )
    setattr(litellm.proxy.proxy_server, "llm_router", router)
    setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)

    request = Request(scope={"type": "http", "method": "POST", "headers": {}})
    request._url = URL(url="/v1/models")

    resp = await model_list(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )

    if hidden:
        assert len(resp["data"]) == 1
    else:
        assert len(resp["data"]) == 2
    print(resp)

    resp = await model_info_v1(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )
    models = json.loads(resp.body)["data"]
    is_model_alias_in_list = False
    for item in models:
        if model_alias == item["model_name"]:
            is_model_alias_in_list = True

    if hidden:
        assert is_model_alias_in_list is False
    else:
        assert is_model_alias_in_list

    resp = await model_group_info(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )
    print(f"resp: {resp}")
    models = resp["data"]
    is_model_alias_in_list = False
    print(f"model_alias: {model_alias}, models: {models}")
    for item in models:
        if model_alias == item.model_group:
            is_model_alias_in_list = True

    if hidden:
        assert is_model_alias_in_list is False
    else:
        assert is_model_alias_in_list, f"models: {models}"


@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
async def test_proxy_model_group_info_rerank(prisma_client):  # noqa: F811  # pytest fixture, not a redefinition
    """
    Check if rerank model is returned on the following endpoints

    `/v1/models`
    `/v1/model/info`
    `/v1/model_group/info`
    """
    import json

    from fastapi import HTTPException, Request, Response
    from starlette.datastructures import URL

    from litellm.proxy.proxy_server import model_group_info, model_info_v1, model_list

    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
    setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
    await litellm.proxy.proxy_server.prisma_client.connect()

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    _model_list = [
        {
            "model_name": "rerank-english-v3.0",
            "litellm_params": {"model": "cohere/rerank-english-v3.0"},
            "model_info": {
                "mode": "rerank",
            },
        }
    ]
    router = litellm.Router(model_list=_model_list)
    setattr(litellm.proxy.proxy_server, "llm_router", router)
    setattr(litellm.proxy.proxy_server, "llm_model_list", _model_list)

    request = Request(scope={"type": "http", "method": "POST", "headers": {}})
    request._url = URL(url="/v1/models")

    resp = await model_list(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )

    assert len(resp["data"]) == 1
    print(resp)

    resp = await model_info_v1(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )
    models = json.loads(resp.body)["data"]
    assert models[0]["model_info"]["mode"] == "rerank"
    resp = await model_group_info(
        user_api_key_dict=UserAPIKeyAuth(models=[]),
    )

    print(resp)
    models = resp["data"]
    assert models[0].mode == "rerank"


# @pytest.mark.asyncio
# async def test_proxy_team_member_add(prisma_client):
#     """
#     Add 10 people to a team. Confirm all 10 are added.
#     """
#     from litellm.proxy.management_endpoints.team_endpoints import (
#         team_member_add,
#         new_team,
#     )
#     from litellm.proxy._types import TeamMemberAddRequest, Member, NewTeamRequest

#     setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
#     setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
#     try:

#         async def test():
#             await litellm.proxy.proxy_server.prisma_client.connect()
#             from litellm.proxy.proxy_server import user_api_key_cache

#             user_api_key_dict = UserAPIKeyAuth(
#                 user_role=LitellmUserRoles.PROXY_ADMIN,
#                 api_key=MASTER_KEY,
#                 user_id="1234",
#             )

#             new_team()
#             for _ in range(10):
#                 request = TeamMemberAddRequest(
#                     team_id="1234",
#                     member=Member(
#                         user_id="1234",
#                         user_role=LitellmUserRoles.INTERNAL_USER,
#                     ),
#                 )
#                 key = await team_member_add(
#                     request, user_api_key_dict=user_api_key_dict
#                 )

#             print(key)
#             user_id = key.user_id

#             # check /user/info to verify user_role was set correctly
#             new_user_info = await user_info(
#                 user_id=user_id, user_api_key_dict=user_api_key_dict
#             )
#             new_user_info = new_user_info.user_info
#             print("new_user_info=", new_user_info)
#             assert new_user_info["user_role"] == LitellmUserRoles.INTERNAL_USER
#             assert new_user_info["user_id"] == user_id

#             generated_key = key.key
#             bearer_token = "Bearer " + generated_key

#             assert generated_key not in user_api_key_cache.in_memory_cache.cache_dict

#             value_from_prisma = await prisma_client.get_data(
#                 token=generated_key,
#             )
#             print("token from prisma", value_from_prisma)

#             request = Request(
#                 {
#                     "type": "http",
#                     "route": api_route,
#                     "path": api_route.path,
#                     "headers": [("Authorization", bearer_token)],
#                 }
#             )

#             # use generated key to auth in
#             result = await user_api_key_auth(request=request, api_key=bearer_token)
#             print("result from user auth with new key", result)

#         asyncio.run(test())
#     except Exception as e:
#         pytest.fail(f"An exception occurred - {str(e)}")


@pytest.mark.asyncio
async def test_proxy_server_prisma_setup():
    from litellm.proxy.proxy_server import ProxyStartupEvent
    from litellm.proxy.utils import ProxyLogging
    from litellm.caching import DualCache

    user_api_key_cache = DualCache()

    with patch.object(
        litellm.proxy.proxy_server, "PrismaClient", new=MagicMock()
    ) as mock_prisma_client:
        mock_client = mock_prisma_client.return_value  # This is the mocked instance
        mock_client.connect = AsyncMock()  # Mock the connect method
        mock_client.start_view_setup_task = MagicMock()
        mock_client.health_check = AsyncMock()  # Mock the health_check method
        mock_client.start_db_health_watchdog_task = AsyncMock()
        # Mock the db attribute with start_token_refresh_task for RDS IAM token refresh
        mock_db = MagicMock()
        mock_db.start_token_refresh_task = AsyncMock()
        mock_client.db = mock_db

        prisma_client = await ProxyStartupEvent._setup_prisma_client(
            database_url="postgresql://user:pass@localhost:5432/litellm",
            proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
            user_api_key_cache=user_api_key_cache,
        )

        assert prisma_client is mock_client
        mock_client.connect.assert_called_once()
        mock_client.start_view_setup_task.assert_called_once()

        # Note: This is REALLY IMPORTANT to check that the health check is called
        # This is how we ensure the DB is ready before proceeding
        mock_client.health_check.assert_called_once()


@pytest.mark.asyncio
async def test_proxy_server_prisma_setup_invalid_db(monkeypatch):
    """
    PROD TEST: Test that proxy server startup fails when it's unable to connect to the database

    Think 2-3 times before editing / deleting this test, it's important for PROD
    """
    import httpx

    from litellm.proxy.proxy_server import ProxyStartupEvent
    from litellm.proxy.utils import ProxyLogging
    from litellm.caching import DualCache

    user_api_key_cache = DualCache()
    invalid_db_url = "postgresql://invalid:invalid@localhost:5432/nonexistent"

    monkeypatch.setenv("DATABASE_URL", invalid_db_url)

    with pytest.raises(httpx.ConnectError):
        await ProxyStartupEvent._setup_prisma_client(
            database_url=invalid_db_url,
            proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
            user_api_key_cache=user_api_key_cache,
        )


@pytest.mark.asyncio
async def test_get_ui_settings_spend_logs_threshold():
    """
    Test that get_ui_settings correctly sets DISABLE_EXPENSIVE_DB_QUERIES based on spend_logs_row_count threshold
    """
    from litellm.proxy.management_endpoints.ui_sso import get_ui_settings
    from litellm.proxy.proxy_server import proxy_state
    from fastapi import Request
    from litellm.constants import MAX_SPENDLOG_ROWS_TO_QUERY

    # Create a mock request
    mock_request = Request(
        scope={
            "type": "http",
            "headers": [],
            "method": "GET",
            "scheme": "http",
            "server": ("testserver", 80),
            "path": "/sso/get/ui_settings",
            "query_string": b"",
        }
    )

    # Test case 1: When spend_logs_row_count > MAX_SPENDLOG_ROWS_TO_QUERY
    proxy_state.set_proxy_state_variable(
        "spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY + 1
    )
    response = await get_ui_settings(mock_request)
    print("response from get_ui_settings", json.dumps(response, indent=4))
    assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is True
    assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY + 1

    # Test case 2: When spend_logs_row_count < MAX_SPENDLOG_ROWS_TO_QUERY
    proxy_state.set_proxy_state_variable(
        "spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY - 1
    )
    response = await get_ui_settings(mock_request)
    print("response from get_ui_settings", json.dumps(response, indent=4))
    assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is False
    assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY - 1

    # Test case 3: Edge case - exactly MAX_SPENDLOG_ROWS_TO_QUERY
    proxy_state.set_proxy_state_variable(
        "spend_logs_row_count", MAX_SPENDLOG_ROWS_TO_QUERY
    )
    response = await get_ui_settings(mock_request)
    print("response from get_ui_settings", json.dumps(response, indent=4))
    assert response["DISABLE_EXPENSIVE_DB_QUERIES"] is False
    assert response["NUM_SPEND_LOGS_ROWS"] == MAX_SPENDLOG_ROWS_TO_QUERY

    # Clean up
    proxy_state.set_proxy_state_variable("spend_logs_row_count", 0)


@pytest.mark.asyncio
async def test_run_background_health_check_reflects_llm_model_list(monkeypatch):
    """
    Test that _run_background_health_check reflects changes to llm_model_list in each health check iteration.
    """
    import litellm.proxy.proxy_server as proxy_server
    import copy

    test_model_list_1 = [{"model_name": "model-a"}]
    test_model_list_2 = [{"model_name": "model-b"}]
    called_model_lists = []

    async def fake_perform_health_check(model_list, details, max_concurrency=None):
        called_model_lists.append(copy.deepcopy(model_list))
        return (["healthy"], ["unhealthy"], {})

    monkeypatch.setattr(proxy_server, "health_check_interval", 1)
    monkeypatch.setattr(proxy_server, "health_check_details", None)
    monkeypatch.setattr(
        proxy_server, "llm_model_list", copy.deepcopy(test_model_list_1)
    )
    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    monkeypatch.setattr(proxy_server, "health_check_results", {})

    async def fake_sleep(interval):
        raise asyncio.CancelledError()

    monkeypatch.setattr(asyncio, "sleep", fake_sleep)

    try:
        await proxy_server._run_background_health_check()
    except asyncio.CancelledError:
        pass

    monkeypatch.setattr(
        proxy_server, "llm_model_list", copy.deepcopy(test_model_list_2)
    )

    try:
        await proxy_server._run_background_health_check()
    except asyncio.CancelledError:
        pass

    assert len(called_model_lists) >= 2
    assert called_model_lists[0] == test_model_list_1
    assert called_model_lists[1] == test_model_list_2


@pytest.mark.asyncio
async def test_background_health_check_skip_disabled_models(monkeypatch):
    """Ensure models with disable_background_health_check are skipped."""
    import litellm.proxy.proxy_server as proxy_server
    import copy

    test_model_list = [
        {"model_name": "model-a"},
        {
            "model_name": "model-b",
            "model_info": {"disable_background_health_check": True},
        },
    ]
    called_model_lists = []

    async def fake_perform_health_check(
        model_list, details, max_concurrency=None, **kwargs
    ):
        called_model_lists.append(copy.deepcopy(model_list))
        return (["healthy"], [], {})

    monkeypatch.setattr(proxy_server, "health_check_interval", 1)
    monkeypatch.setattr(proxy_server, "health_check_details", None)
    monkeypatch.setattr(proxy_server, "llm_model_list", copy.deepcopy(test_model_list))
    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    monkeypatch.setattr(proxy_server, "health_check_results", {})

    async def fake_sleep(interval):
        raise asyncio.CancelledError()

    monkeypatch.setattr(asyncio, "sleep", fake_sleep)

    try:
        await proxy_server._run_background_health_check()
    except asyncio.CancelledError:
        pass

    assert called_model_lists == [[{"model_name": "model-a"}]]


@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_legacy_three_arg_stub(
    monkeypatch,
):
    """Monkeypatched perform_health_check with only base kwargs should still run."""
    import litellm.proxy.proxy_server as proxy_server

    async def fake_perform_health_check(model_list, details, max_concurrency=None):
        return ([], [], {})

    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    result = await proxy_server._run_direct_health_check_with_instrumentation(
        [{"model_name": "m"}],
        True,
        1,
        {"enabled": True, "source": "test", "cycle_id": "c1"},
    )
    assert result == ([], [], {})


@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_accepts_instrumentation_only(
    monkeypatch,
):
    """Stub that accepts instrumentation_context but not health_check filter kwargs."""
    import litellm.proxy.proxy_server as proxy_server

    seen: list = []

    async def fake_perform_health_check(
        model_list, details, max_concurrency=None, instrumentation_context=None
    ):
        seen.append(instrumentation_context)
        return ([], [], {})

    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    await proxy_server._run_direct_health_check_with_instrumentation(
        [],
        False,
        2,
        {"enabled": True, "source": "test", "cycle_id": "c2"},
    )
    assert len(seen) == 1
    assert seen[0]["cycle_id"] == "c2"


@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_accepts_filter_only(
    monkeypatch,
):
    """Stub that accepts health_check_skip_disabled_background_models but not instrumentation."""
    import litellm.proxy.proxy_server as proxy_server

    seen: list = []

    async def fake_perform_health_check(
        model_list,
        details,
        max_concurrency=None,
        health_check_skip_disabled_background_models=False,
    ):
        seen.append(health_check_skip_disabled_background_models)
        return ([], [], {})

    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    await proxy_server._run_direct_health_check_with_instrumentation(
        [],
        True,
        None,
        {"enabled": False},
    )
    assert len(seen) == 1
    assert seen[0] is False


@pytest.mark.asyncio
async def test_run_direct_health_check_drops_only_the_rejected_kwarg(monkeypatch):
    """A callee that predates `router` must still get the skip-disabled filter: dropping the
    rejected argument alongside working ones would probe deployments the operator opted out."""
    import litellm.proxy.proxy_server as proxy_server

    seen: list[tuple[dict[str, str] | None, bool]] = []

    async def fake_perform_health_check(
        model_list,
        details,
        max_concurrency=None,
        instrumentation_context=None,
        health_check_skip_disabled_background_models=False,
    ):
        seen.append((instrumentation_context, health_check_skip_disabled_background_models))
        return ([], [], {})

    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    monkeypatch.setattr(
        proxy_server,
        "general_settings",
        {"health_check_skip_disabled_background_models": True},
    )
    await proxy_server._run_direct_health_check_with_instrumentation([], True, 1, {"cycle_id": "c3"})

    assert seen == [({"cycle_id": "c3"}, True)]


@pytest.mark.asyncio
async def test_run_direct_health_check_with_instrumentation_non_kw_typeerror_reraises(
    monkeypatch,
):
    import litellm.proxy.proxy_server as proxy_server

    async def fake_perform_health_check(**kwargs):
        raise TypeError("unsupported operand type(s)")

    monkeypatch.setattr(proxy_server, "perform_health_check", fake_perform_health_check)
    with pytest.raises(TypeError, match="unsupported operand"):
        await proxy_server._run_direct_health_check_with_instrumentation(
            [],
            True,
            1,
            {},
        )


def test_get_timeout_from_request():
    from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup

    headers = {
        "x-litellm-timeout": "90",
    }
    timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
    assert timeout == 90

    headers = {
        "x-litellm-timeout": "90.5",
    }
    timeout = LiteLLMProxyRequestSetup._get_timeout_from_request(headers)
    assert timeout == 90.5


def test_add_litellm_data_for_backend_llm_call_marks_client_side_timeout():
    """A caller-supplied x-litellm-timeout must be marked with client_side_timeout=True,
    so the router's fallback-cooldown trigger can tell it apart from a deployment
    actually timing out (a caller could otherwise force every deployment in a fallback
    chain to look unhealthy with a single near-zero timeout request)."""
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup

    user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key")

    data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
        headers={"x-litellm-timeout": "0.001"},
        request_data={},
        user_api_key_dict=user_api_key_dict,
    )
    assert data["timeout"] == 0.001
    assert data["client_side_timeout"] is True

    data_without_header = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
        headers={},
        request_data={},
        user_api_key_dict=user_api_key_dict,
    )
    assert "client_side_timeout" not in data_without_header


@pytest.mark.parametrize(
    "request_data",
    [
        {"timeout": 0.001},
        {"request_timeout": 0.001},
        {"stream_timeout": 0.001},
    ],
)
def test_add_litellm_data_for_backend_llm_call_marks_client_side_timeout_from_body(
    request_data,
):
    """Router._get_timeout resolves the effective timeout from kwargs["timeout"],
    kwargs["request_timeout"], or kwargs["stream_timeout"], and a caller can supply any
    of those directly in the request body, not just via the x-litellm-timeout header.
    Missing this would let a caller force a 408 on every deployment in a fallback chain
    without it being recognized as caller-controlled, cooling down deployments other
    tenants rely on."""
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup

    user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key")

    data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
        headers={},
        request_data=request_data,
        user_api_key_dict=user_api_key_dict,
    )
    assert data["client_side_timeout"] is True


def test_add_litellm_data_for_backend_llm_call_ignores_forged_client_side_timeout():
    """The caller-supplied client_side_timeout key itself must never be trusted verbatim:
    the marker is always recomputed from the actual timeout sources, so a caller can't
    forge client_side_timeout=True to dodge cooldown on a real deployment failure."""
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup

    user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key")

    data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call(
        headers={},
        request_data={"client_side_timeout": True},
        user_api_key_dict=user_api_key_dict,
    )
    assert "client_side_timeout" not in data


@pytest.mark.parametrize(
    "ui_exists, ui_has_content",
    [
        (True, True),  # UI path exists and has content
        (True, False),  # UI path exists but is empty
        (False, False),  # UI path doesn't exist
    ],
)
def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content):
    """
    Test the non-root Docker UI path detection logic.

    Tests that when LITELLM_NON_ROOT is set to "true":
    - If UI path exists and has content, it should be used
    - If UI path doesn't exist or is empty, proper error logging occurs
    """
    import tempfile
    import shutil
    from unittest.mock import MagicMock

    # Create a temporary directory to act as /tmp/litellm_ui
    test_ui_path = tmp_path / "litellm_ui"

    if ui_exists:
        test_ui_path.mkdir(parents=True, exist_ok=True)
        if ui_has_content:
            # Create some dummy files to simulate built UI
            (test_ui_path / "index.html").write_text("<html></html>")
            (test_ui_path / "app.js").write_text("console.log('test');")

    # Mock the environment variable and os.path operations
    monkeypatch.setenv("LITELLM_NON_ROOT", "true")

    # Create a mock logger to capture log messages
    mock_logger = MagicMock()

    # We need to reimport or reload the relevant code section
    # Since this is module-level code, we'll test the logic directly
    ui_path = None
    non_root_ui_path = str(test_ui_path)

    # Simulate the logic from proxy_server.py lines 909-920
    if os.getenv("LITELLM_NON_ROOT", "").lower() == "true":
        if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path):
            mock_logger.info(
                f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
            )
            mock_logger.info(
                f"UI files found: {len(os.listdir(non_root_ui_path))} items"
            )
            ui_path = non_root_ui_path
        else:
            mock_logger.error(
                f"UI not found at {non_root_ui_path}. UI will not be available."
            )
            mock_logger.error(
                f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}"
            )

    # Verify behavior based on test parameters
    if ui_exists and ui_has_content:
        # UI should be found and used
        assert ui_path == non_root_ui_path
        assert mock_logger.info.call_count == 2
        mock_logger.info.assert_any_call(
            f"Using pre-built UI for non-root Docker: {non_root_ui_path}"
        )
        # Verify the second info call mentions the number of items
        info_calls = [call[0][0] for call in mock_logger.info.call_args_list]
        assert any("UI files found:" in call and "items" in call for call in info_calls)
        assert mock_logger.error.call_count == 0
    else:
        # UI should not be found, error should be logged
        assert ui_path is None
        assert mock_logger.error.call_count == 2
        mock_logger.error.assert_any_call(
            f"UI not found at {non_root_ui_path}. UI will not be available."
        )
        # Verify the second error call has path existence info
        error_calls = [call[0][0] for call in mock_logger.error.call_args_list]
        assert any("Path exists:" in call for call in error_calls)
        assert mock_logger.info.call_count == 0


@pytest.mark.asyncio
async def test_get_config_callbacks_with_all_types(client_no_auth):
    """
    Test that /get/config/callbacks returns all three callback types:
    - success_callback with type="success"
    - failure_callback with type="failure"
    - callbacks (success_and_failure) with type="success_and_failure"
    """
    from litellm.proxy.proxy_server import ProxyConfig

    # Create a mock config with all three callback types
    mock_config_data = {
        "litellm_settings": {
            "success_callback": ["langfuse", "braintrust"],
            "failure_callback": ["sentry"],
            "callbacks": ["otel", "langsmith"],
        },
        "environment_variables": {
            "LANGFUSE_PUBLIC_KEY": "test-public-key",
            "LANGFUSE_SECRET_KEY": "test-secret-key",
            "LANGFUSE_HOST": "https://test.langfuse.com",
            "BRAINTRUST_API_KEY": "test-braintrust-key",
            "OTEL_EXPORTER": "otlp",
            "OTEL_ENDPOINT": "http://localhost:4317",
            "LANGSMITH_API_KEY": "test-langsmith-key",
        },
        "general_settings": {},
    }

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    with patch.object(
        proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
    ):
        response = client_no_auth.get("/get/config/callbacks")

        assert response.status_code == 200
        result = response.json()

        # Verify response structure
        assert "status" in result
        assert result["status"] == "success"
        assert "callbacks" in result

        callbacks = [cb for cb in result["callbacks"] if not cb.get("read_only", False)]

        # Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure)
        assert len(callbacks) == 5

        # Group callbacks by type
        success_callbacks = [cb for cb in callbacks if cb.get("type") == "success"]
        failure_callbacks = [cb for cb in callbacks if cb.get("type") == "failure"]
        success_and_failure_callbacks = [
            cb for cb in callbacks if cb.get("type") == "success_and_failure"
        ]

        # Verify all callbacks have required fields
        for callback in callbacks:
            assert "name" in callback
            assert "variables" in callback
            assert "type" in callback
            assert callback["type"] in ["success", "failure", "success_and_failure"]

        # Verify success callbacks
        assert len(success_callbacks) == 2
        success_names = [cb["name"] for cb in success_callbacks]
        assert "langfuse" in success_names
        assert "braintrust" in success_names

        # Verify failure callbacks
        assert len(failure_callbacks) == 1
        assert failure_callbacks[0]["name"] == "sentry"

        # Verify success_and_failure callbacks
        assert len(success_and_failure_callbacks) == 2
        success_and_failure_names = [cb["name"] for cb in success_and_failure_callbacks]
        assert "otel" in success_and_failure_names
        assert "langsmith" in success_and_failure_names


@pytest.mark.asyncio
async def test_get_config_callbacks_environment_variables(client_no_auth):
    """
    Test that /get/config/callbacks correctly includes environment variables
    for each callback type. Under ``client_no_auth`` the resolved role is
    not ``PROXY_ADMIN``, so values matched by the redaction helper come back
    as ``"REDACTED"`` and other values pass through verbatim.
    """
    from litellm.proxy.proxy_server import ProxyConfig

    # Create a mock config with callbacks and their env vars
    mock_config_data = {
        "litellm_settings": {
            "success_callback": ["langfuse"],
            "failure_callback": [],
            "callbacks": ["otel"],
        },
        "environment_variables": {
            "LANGFUSE_PUBLIC_KEY": "test-public-key",
            "LANGFUSE_SECRET_KEY": "test-secret-key",
            "LANGFUSE_HOST": "https://cloud.langfuse.com",
            "OTEL_EXPORTER": "otlp",
            "OTEL_ENDPOINT": "http://localhost:4317",
            "OTEL_HEADERS": "key=value",
        },
        "general_settings": {},
    }

    proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")

    with patch.object(
        proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data)
    ):
        response = client_no_auth.get("/get/config/callbacks")

        assert response.status_code == 200
        result = response.json()

        callbacks = result["callbacks"]

        # Find langfuse callback (success type)
        langfuse_callback = next(
            (cb for cb in callbacks if cb["name"] == "langfuse"), None
        )
        assert langfuse_callback is not None
        assert langfuse_callback["type"] == "success"
        assert "variables" in langfuse_callback

        langfuse_vars = langfuse_callback["variables"]
        assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars
        assert langfuse_vars["LANGFUSE_PUBLIC_KEY"] == "REDACTED"
        assert "LANGFUSE_SECRET_KEY" in langfuse_vars
        assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "REDACTED"
        assert "LANGFUSE_HOST" in langfuse_vars
        assert langfuse_vars["LANGFUSE_HOST"] == "https://cloud.langfuse.com"

        # Find otel callback (success_and_failure type)
        otel_callback = next((cb for cb in callbacks if cb["name"] == "otel"), None)
        assert otel_callback is not None
        assert otel_callback["type"] == "success_and_failure"
        assert "variables" in otel_callback

        otel_vars = otel_callback["variables"]
        assert "OTEL_EXPORTER" in otel_vars
        assert otel_vars["OTEL_EXPORTER"] == "otlp"
        assert "OTEL_ENDPOINT" in otel_vars
        assert otel_vars["OTEL_ENDPOINT"] == "http://localhost:4317"
        assert "OTEL_HEADERS" in otel_vars
        assert otel_vars["OTEL_HEADERS"] == "REDACTED"


@pytest.mark.asyncio
async def test_update_config_success_callback_normalization(monkeypatch):
    """
    Ensure success_callback values are normalized to lowercase when updating config.
    This prevents delete_callback (which searches lowercase) from failing on mixed case inputs like 'SQS'.
    """
    import litellm.proxy.proxy_server as proxy_server
    from litellm.proxy._types import ConfigYAML

    monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock())

    existing_litellm_settings = {"success_callback": ["langfuse"]}

    class FakeRow:
        def __init__(self, name, value):
            self.param_name = name
            self.param_value = value

    upserted = {}

    async def fake_find_first(where=None):
        if where and where.get("param_name") == "litellm_settings":
            return FakeRow("litellm_settings", existing_litellm_settings)
        return None

    async def fake_upsert(where=None, data=None):
        upserted[where["param_name"]] = json.loads(data["update"]["param_value"])

    class MockPrisma:
        def __init__(self):
            self.db = MagicMock()
            self.db.litellm_config = MagicMock()
            self.db.litellm_config.find_first = AsyncMock(side_effect=fake_find_first)
            self.db.litellm_config.upsert = AsyncMock(side_effect=fake_upsert)

    monkeypatch.setattr(proxy_server, "prisma_client", MockPrisma())

    class MockProxyConfig:
        async def add_deployment(self, prisma_client=None, proxy_logging_obj=None):  # noqa: F811  # pytest fixture, not a redefinition
            return None

        def reject_config_owned_writes(self, *, section_name, changed_keys):
            return None

    monkeypatch.setattr(proxy_server, "proxy_config", MockProxyConfig())

    config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]})
    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth

    admin_user = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test"
    )
    request = MagicMock()
    request.json = AsyncMock(return_value={"litellm_settings": {"success_callback": ["SQS", "sQs"]}})
    await proxy_server.update_config(config_update, request=request, user_api_key_dict=admin_user)

    assert (
        "litellm_settings" in upserted
    ), "litellm_config.upsert was not called for litellm_settings"
    callbacks = upserted["litellm_settings"]["success_callback"]

    # Deduped and normalized
    assert "sqs" in callbacks
    assert "SQS" not in callbacks
    assert "sQs" not in callbacks
    # Existing callback should still be present
    assert "langfuse" in callbacks


@pytest.mark.parametrize(
    "data",
    [
        {
            "model": {
                "model_name": "azure/gpt-4.1-mini",
                "litellm_params": {"model": "azure/gpt-4.1-mini"},
                "model_info": {"base_model": "gpt-4.1-mini"},
            },
            "expected": "gpt-4.1-mini",
        },
        {
            "model": {
                "model_name": "openai/gpt-4.1-mini",
                "litellm_params": {"model": "openai/gpt-4.1-mini"},
            },
            "expected": "openai/gpt-4.1-mini",
        },
        {
            "model": {
                "model_name": "openai/gpt-4.1-mini",
                "litellm_params": {"model": "openai/gpt-4.1-mini"},
                "model_info": {"base_model": "gpt-4.1-mini"},
            },
            "expected": "gpt-4.1-mini",
        },
        {
            "model": {
                "model_name": "claude-sonnet-4-5-20250929",
                "litellm_params": {"model": "anthropic/claude-sonnet-4-5@20250929"},
                "model_info": {"base_model": "anthropic/claude-sonnet-4-5-20250929"},
            },
            "expected": "anthropic/claude-sonnet-4-5-20250929",
        },
        {
            "model": {
                "model_name": "gemini-2.5-flash-001",
                "litellm_params": {"model": "gemini/gemini-2.5-flash@001"},
                "model_info": {"base_model": "gemini-2.5-flash-001"},
            },
            "expected": "gemini-2.5-flash-001",
        },
    ],
)
def test_get_litellm_model_info(data):
    from litellm.proxy.proxy_server import get_litellm_model_info

    model = data["model"]
    get_info_mock = MagicMock()

    with mock.patch(
        "litellm.get_model_info",
        new=get_info_mock,
    ):
        get_litellm_model_info(model=model)
        get_info_mock.assert_called_once_with(data["expected"])


def test_live_session_accept_failure_releases_reservations(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from litellm.proxy import proxy_server

    released: Final[list[str]] = []

    async def release_budget(user_api_key_dict: object) -> None:
        released.append("budget")

    async def release_slot(user_api_key_dict: object) -> None:
        released.append("slot")

    async def receive() -> dict[str, object]:
        return {"type": "websocket.connect"}

    async def send_failing_accept(message: Mapping[str, object]) -> None:
        if message["type"] == "websocket.accept":
            raise RuntimeError("accept failed")

    scope: Final = {
        "type": "websocket",
        "path": "/v1/live/sessions",
        "raw_path": b"/v1/live/sessions",
        "root_path": "",
        "scheme": "ws",
        "query_string": b"",
        "headers": [(b"host", b"testserver"), (b"authorization", b"Bearer test-api-key")],
        "client": ("testclient", 50000),
        "server": ("testserver", 80),
        "subprotocols": [],
    }
    route_call: Final = AsyncMock()
    monkeypatch.setattr(proxy_server, "route_request", route_call)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot)

    with pytest.raises(RuntimeError, match="accept failed"):
        asyncio.run(client_no_auth.app(scope, receive, send_failing_accept))

    assert released == ["budget", "slot"]
    route_call.assert_not_awaited()


def _override_live_session_authorizer(
    monkeypatch: pytest.MonkeyPatch,
    authorizer: Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]],
) -> None:
    from litellm.proxy import proxy_server

    monkeypatch.setitem(
        proxy_server.app.dependency_overrides,
        proxy_server.live_session_model_authorizer,
        lambda: authorizer,
    )


def _start_live_session_and_wait_for_reply(client: TestClient) -> None:
    with client.websocket_connect("/v1/live/sessions", headers={"Authorization": "Bearer test-api-key"}) as websocket:
        websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}})
        websocket.receive_text()


def test_live_session_model_check_error_releases_reservations(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from litellm.proxy import proxy_server

    released: Final[list[str]] = []

    async def release_budget(user_api_key_dict: object) -> None:
        released.append("budget")

    async def release_slot(user_api_key_dict: object) -> None:
        released.append("slot")

    async def failing_authorizer(websocket: object, model: str) -> UserAPIKeyAuth:
        raise RuntimeError("agent lookup failed")

    route_call: Final = AsyncMock()
    monkeypatch.setattr(proxy_server, "route_request", route_call)
    _override_live_session_authorizer(monkeypatch, failing_authorizer)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot)

    with pytest.raises(RuntimeError, match="agent lookup failed"):
        _start_live_session_and_wait_for_reply(client_no_auth)

    assert released == ["budget", "slot"]
    route_call.assert_not_awaited()


class _PreCallAborted(BaseException):
    pass


def test_live_session_pre_call_abort_releases_reservations(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from litellm.proxy import proxy_server
    from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing

    frame_auth: Final = UserAPIKeyAuth(token="frame-authorized")
    released_for_frame: Final[list[str]] = []

    async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None:
        if user_api_key_dict is frame_auth:
            released_for_frame.append("budget")

    async def release_slot(user_api_key_dict: UserAPIKeyAuth) -> None:
        if user_api_key_dict is frame_auth:
            released_for_frame.append("slot")

    async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
        return frame_auth

    class _AbortingPreCallProcessing(ProxyBaseLLMRequestProcessing):
        async def common_processing_pre_call_logic(self, *args: object, **kwargs: object) -> NoReturn:
            raise _PreCallAborted

    route_call: Final = AsyncMock()
    monkeypatch.setattr(proxy_server, "route_request", route_call)
    _override_live_session_authorizer(monkeypatch, authorize)
    monkeypatch.setattr(proxy_server, "ProxyBaseLLMRequestProcessing", _AbortingPreCallProcessing)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", release_slot)

    with pytest.raises(_PreCallAborted):
        _start_live_session_and_wait_for_reply(client_no_auth)

    assert released_for_frame == ["budget", "slot"]
    route_call.assert_not_awaited()


@pytest.mark.parametrize("path", ["/openai/v1/live/sessions", "/v1/live/sessions", "/live/sessions"])
def test_live_session_invalid_start_is_rejected_without_routing(
    path: str,
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from litellm.proxy import proxy_server
    from starlette.websockets import WebSocketDisconnect

    async def no_op(user_api_key_dict: object) -> None:
        return None

    route_call: Final = AsyncMock()
    monkeypatch.setattr(proxy_server, "route_request", route_call)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)

    with client_no_auth.websocket_connect(path, headers={"Authorization": "Bearer test-api-key"}) as websocket:
        websocket.send_text('{"type":"session.start","session":{"model":""}}')
        assert websocket.receive_json() == {
            "type": "error",
            "error": {
                "type": "invalid_request_error",
                "code": "invalid_session_start",
                "message": "First message must be a session.start JSON object with a non-empty session.model.",
            },
        }
        with pytest.raises(WebSocketDisconnect) as error:
            websocket.receive_json()
        assert error.value.code == 1008

    route_call.assert_not_awaited()


@pytest.mark.parametrize(
    ("denial_type", "expected_code"),
    (("user_model_access_denied", "model_not_allowed"), ("budget_exceeded", "session_not_allowed")),
)
def test_live_session_frame_authorization_denial_closes_with_policy_code(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
    denial_type: str,
    expected_code: str,
) -> None:
    from litellm.proxy import proxy_server
    from litellm.proxy._types import ProxyException
    from starlette.websockets import WebSocketDisconnect

    async def no_op(user_api_key_dict: object) -> None:
        return None

    authorized_models: Final[list[str]] = []

    async def deny(websocket: object, model: str) -> UserAPIKeyAuth:
        authorized_models.append(model)
        raise ProxyException("model access denied for this user", denial_type, "model", 403)

    route_call: Final = AsyncMock()
    monkeypatch.setattr(proxy_server, "route_request", route_call)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
    _override_live_session_authorizer(monkeypatch, deny)

    with client_no_auth.websocket_connect(
        "/v1/live/sessions",
        headers={"Authorization": "Bearer test-api-key"},
    ) as websocket:
        websocket.send_json({"type": "session.start", "session": {"model": "gpt-4o"}})
        assert websocket.receive_json() == {
            "type": "error",
            "error": {
                "type": "invalid_request_error",
                "code": expected_code,
                "message": "model access denied for this user",
            },
        }
        with pytest.raises(WebSocketDisconnect) as error:
            websocket.receive_json()
        assert error.value.code == 1008

    assert authorized_models == ["gpt-4o"]
    route_call.assert_not_awaited()


def test_live_session_valid_start_routes_with_model_and_frame(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from types import SimpleNamespace

    from litellm.proxy import proxy_server
    from starlette.websockets import WebSocket

    async def no_op(user_api_key_dict: object) -> None:
        return None

    async def complete_call() -> None:
        return None

    captured: Final = SimpleNamespace(data=None, route_type=None)

    async def route(
        *,
        data: dict[str, object],
        route_type: str,
        llm_router: object,
        user_model: str | None,
    ) -> Coroutine[object, object, None]:
        captured.data = data
        captured.route_type = route_type
        await cast(WebSocket, data["websocket"]).send_text("routed")
        return complete_call()

    frame_auth: Final = UserAPIKeyAuth(token="frame-authorized")

    async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
        return frame_auth

    monkeypatch.setattr(proxy_server, "route_request", route)
    _override_live_session_authorizer(monkeypatch, authorize)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
    start_frame: Final[dict[str, object]] = {
        "type": "session.start",
        "session": {"model": "gpt-3.5-turbo", "instructions": "Stay brief"},
    }

    with client_no_auth.websocket_connect(
        "/v1/live/sessions",
        headers={"Authorization": "Bearer test-api-key"},
    ) as websocket:
        websocket.send_json(start_frame)
        assert websocket.receive_text() == "routed"

    routed_data: Final = cast(dict[str, object], captured.data)
    assert captured.route_type == "_arealtime"
    assert routed_data["model"] == "gpt-3.5-turbo"
    assert routed_data["live_session_start"] == start_frame
    assert routed_data["user_api_key_dict"] is frame_auth


def test_live_session_ignores_query_model_and_authorizes_frame_model(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from fastapi import Request

    from litellm.proxy import proxy_server
    from starlette.websockets import WebSocket

    async def no_op(user_api_key_dict: object) -> None:
        return None

    async def complete_call() -> None:
        return None

    async def route(
        *,
        data: dict[str, object],
        route_type: str,
        llm_router: object,
        user_model: str | None,
    ) -> Coroutine[object, object, None]:
        await cast(WebSocket, data["websocket"]).send_text("routed")
        return complete_call()

    upgrade_auth: Final = UserAPIKeyAuth(token="upgrade", models=["gpt-live-1"])

    async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth:
        if await request.body() != b'{"model": ""}':
            raise HTTPException(status_code=403, detail="Model is not allowed")
        return upgrade_auth

    auth_call: Final = AsyncMock(side_effect=authenticate)
    monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", auth_call)
    frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", models=["gpt-live-1"])
    authorize: Final = AsyncMock(return_value=frame_auth)
    monkeypatch.setattr(proxy_server, "route_request", route)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
    _override_live_session_authorizer(monkeypatch, authorize)

    with client_no_auth.websocket_connect(
        "/v1/live/sessions?model=gpt-realtime",
        headers={"Authorization": "Bearer test-api-key"},
    ) as websocket:
        websocket.send_json({"type": "session.start", "session": {"model": "gpt-live-1"}})
        assert websocket.receive_text() == "routed"

    request_arg: Final = auth_call.await_args.kwargs["request"]
    assert json.loads(asyncio.run(request_arg.body())) == {"model": ""}
    assert authorize.await_args.args[1] == "gpt-live-1"


def test_live_session_reauthorized_token_still_needs_the_model_on_the_key(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from starlette.websockets import WebSocketDisconnect

    from litellm.proxy import proxy_server

    released_tokens: Final[list[str | None]] = []

    async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None:
        released_tokens.append(user_api_key_dict.token)

    async def no_op(user_api_key_dict: object) -> None:
        return None

    async def custom_auth_without_common_checks(websocket: object, model: str) -> UserAPIKeyAuth:
        return UserAPIKeyAuth(token="custom-auth", models=["gpt-4o"])

    route_call: Final = AsyncMock()
    monkeypatch.setattr(proxy_server, "route_request", route_call)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
    _override_live_session_authorizer(monkeypatch, custom_auth_without_common_checks)

    with client_no_auth.websocket_connect(
        "/v1/live/sessions",
        headers={"Authorization": "Bearer test-api-key"},
    ) as websocket:
        websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}})
        error_event: Final = websocket.receive_json()
        with pytest.raises(WebSocketDisconnect) as error:
            websocket.receive_json()

    assert error_event["error"]["code"] == "model_not_allowed", error_event
    assert error.value.code == 1008
    assert "custom-auth" in released_tokens
    route_call.assert_not_awaited()


def test_live_session_releases_the_upgrade_reservation_before_the_frame_reserves(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from starlette.websockets import WebSocket

    from litellm.proxy import proxy_server

    released_reservations: Final[list[object]] = []

    async def release_budget(user_api_key_dict: UserAPIKeyAuth) -> None:
        released_reservations.append(user_api_key_dict.budget_reservation)

    async def no_op(user_api_key_dict: object) -> None:
        return None

    async def complete_call() -> None:
        return None

    async def route(
        *,
        data: dict[str, object],
        route_type: str,
        llm_router: object,
        user_model: str | None,
    ) -> Coroutine[object, object, None]:
        await cast(WebSocket, data["websocket"]).send_text("routed")
        return complete_call()

    frame_auth: Final = UserAPIKeyAuth(token="frame-authorized", budget_reservation={"reserved_for": "frame model"})

    async def upgrade_auth() -> UserAPIKeyAuth:
        return UserAPIKeyAuth(token="upgrade", budget_reservation={"reserved_for": "query model"})

    released_before_frame_auth: Final[list[tuple[object, ...]]] = []

    async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
        released_before_frame_auth.append(tuple(released_reservations))
        return frame_auth

    monkeypatch.setattr(proxy_server, "route_request", route)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", release_budget)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)
    monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth_live_websocket, upgrade_auth)
    _override_live_session_authorizer(monkeypatch, authorize)

    with client_no_auth.websocket_connect(
        "/v1/live/sessions?model=gpt-4o",
        headers={"Authorization": "Bearer test-api-key"},
    ) as websocket:
        websocket.send_json({"type": "session.start", "session": {"model": "gpt-3.5-turbo"}})
        assert websocket.receive_text() == "routed"

    assert released_before_frame_auth == [({"reserved_for": "query model"},)], released_reservations


def test_live_session_preserves_pretty_printed_start_frame(
    client_no_auth: TestClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from types import SimpleNamespace

    from starlette.websockets import WebSocket

    from litellm.proxy import proxy_server

    async def no_op(user_api_key_dict: object) -> None:
        return None

    async def complete_call() -> None:
        return None

    captured: Final = SimpleNamespace(data=None, route_type=None)

    async def route(
        *,
        data: dict[str, object],
        route_type: str,
        llm_router: object,
        user_model: str | None,
    ) -> Coroutine[object, object, None]:
        captured.data = data
        captured.route_type = route_type
        await cast(WebSocket, data["websocket"]).send_text("routed")
        return complete_call()

    frame_auth: Final = UserAPIKeyAuth(token="frame-authorized")

    async def authorize(websocket: object, model: str) -> UserAPIKeyAuth:
        return frame_auth

    monkeypatch.setattr(proxy_server, "route_request", route)
    _override_live_session_authorizer(monkeypatch, authorize)
    monkeypatch.setattr(proxy_server, "_release_realtime_budget_reservation", no_op)
    monkeypatch.setattr(proxy_server, "_release_realtime_max_parallel_slot", no_op)

    start_frame: Final = "\r\n".join(
        [
            "{",
            '  "type": "session.start",',
            '  "session": {"model": "gpt-live-1", "instructions": "a\\nb"}',
            "}",
        ]
    )
    expected_start: Final[dict[str, object]] = {
        "type": "session.start",
        "session": {"model": "gpt-live-1", "instructions": "a\nb"},
    }

    with client_no_auth.websocket_connect(
        "/v1/live/sessions",
        headers={"Authorization": "Bearer test-api-key"},
    ) as websocket:
        websocket.send_text(start_frame)
        assert websocket.receive_text() == "routed"

    routed_data: Final = cast(dict[str, object], captured.data)
    assert captured.route_type == "_arealtime"
    assert routed_data["model"] == "gpt-live-1"
    assert routed_data["live_session_start"] == expected_start