import json
import os
import sys

import litellm
import pytest
import yaml
from fastapi.testclient import TestClient

sys.path.insert(
    0, os.path.abspath("../../../..")
)  # Adds the parent directory to the system path

from unittest.mock import AsyncMock, MagicMock, patch

from fastapi import HTTPException

from litellm.proxy._types import (
    GenerateKeyRequest,
    LiteLLM_BudgetTable,
    LiteLLM_OrganizationTable,
    LiteLLM_TeamTableCachedObj,
    LiteLLM_UserTable,
    LiteLLM_VerificationToken,
    LitellmUserRoles,
    Member,
    ProxyException,
    ResetSpendRequest,
    UpdateKeyRequest,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.management_endpoints.key_management_endpoints import (
    _check_org_key_limits,
    _check_team_key_limits,
    _common_key_generation_helper,
    _enforce_upperbound_key_params,
    _get_and_validate_existing_key,
    _list_key_helper,
    _persist_deleted_verification_tokens,
    _process_single_key_update,
    _save_deleted_verification_token_records,
    _transform_verification_tokens_to_deleted_records,
    _validate_max_budget,
    _validate_reset_spend_value,
    _validate_update_key_data,
    can_modify_verification_token,
    check_org_key_model_specific_limits,
    check_team_key_model_specific_limits,
    delete_verification_tokens,
    generate_key_fn,
    generate_key_helper_fn,
    key_aliases,
    key_generation_check,
    list_keys,
    prepare_key_update_data,
    reset_key_spend_fn,
    validate_key_list_check,
    validate_key_team_change,
)
from litellm.proxy.proxy_server import app

client = TestClient(app)


@pytest.mark.asyncio
async def test_list_keys():
    mock_prisma_client = AsyncMock()
    mock_find_many = AsyncMock(return_value=[])
    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
    args = {
        "prisma_client": mock_prisma_client,
        "page": 1,
        "size": 50,
        "user_id": "cda88cb4-cc2c-4e8c-b871-dc71ca111b00",
        "team_id": None,
        "organization_id": None,
        "key_alias": None,
        "key_hash": None,
        "exclude_team_id": None,
        "return_full_object": True,
        "admin_team_ids": ["28bd3181-02c5-48f2-b408-ce790fb3d5ba"],
    }
    try:
        result = await _list_key_helper(**args)
    except Exception as e:
        print(f"error: {e}")

    mock_find_many.assert_called_once()

    where_condition = mock_find_many.call_args.kwargs["where"]
    print(f"where_condition: {where_condition}")
    assert json.dumps({"team_id": {"not": "litellm-dashboard"}}) in json.dumps(
        where_condition
    )


@pytest.mark.asyncio
async def test_list_keys_include_created_by_keys():
    """
    Test that include_created_by_keys parameter correctly includes keys created by the user
    and applies specific filtering to both user's own keys and created_by keys.
    """
    mock_prisma_client = AsyncMock()
    mock_find_many = AsyncMock(return_value=[])
    mock_count = AsyncMock(return_value=0)
    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
    mock_prisma_client.db.litellm_verificationtoken.count = mock_count

    test_user_id = "user-123"
    test_org_id = "org-456"
    test_key_alias = "test-alias"
    test_key_hash = "hashed-token-789"

    # Test Case 1: include_created_by_keys=True with specific filters
    args = {
        "prisma_client": mock_prisma_client,
        "page": 1,
        "size": 50,
        "user_id": test_user_id,
        "team_id": None,
        "organization_id": test_org_id,
        "key_alias": test_key_alias,
        "key_hash": test_key_hash,
        "exclude_team_id": None,
        "return_full_object": True,
        "admin_team_ids": None,
        "include_created_by_keys": True,
    }

    try:
        result = await _list_key_helper(**args)
    except Exception as e:
        print(f"error: {e}")

    mock_find_many.assert_called_once()
    mock_count.assert_called_once()

    where_condition = mock_find_many.call_args.kwargs["where"]
    print(f"where_condition with include_created_by_keys=True: {where_condition}")

    # Verify the structure contains AND with OR conditions
    assert "AND" in where_condition
    assert "OR" in where_condition["AND"][1]

    or_conditions = where_condition["AND"][1]["OR"]

    # Should have 2 OR conditions: user's own keys and created_by keys
    assert len(or_conditions) == 2

    # First condition should be user's own keys with all filters applied
    user_condition = None
    created_by_condition = None

    for condition in or_conditions:
        if "user_id" in condition:
            user_condition = condition
        elif "created_by" in condition:
            created_by_condition = condition

    assert user_condition is not None, "User condition should be present"
    assert created_by_condition is not None, "Created by condition should be present"

    # Verify user condition has all the filters
    assert user_condition["user_id"] == test_user_id
    assert user_condition["organization_id"] == test_org_id
    assert user_condition["key_alias"] == test_key_alias
    assert user_condition["token"] == test_key_hash

    # Verify created_by condition only has the created_by filter (no other filters applied)
    # This is the current behavior - created_by keys don't inherit other filters
    assert created_by_condition["created_by"] == test_user_id
    assert (
        len(created_by_condition) == 1
    ), "Created by condition should only have created_by field"

    # Reset mocks for Test Case 2
    mock_find_many.reset_mock()
    mock_count.reset_mock()

    # Test Case 2: include_created_by_keys=False should not include created_by condition
    args["include_created_by_keys"] = False

    try:
        result = await _list_key_helper(**args)
    except Exception as e:
        print(f"error: {e}")

    where_condition_no_created_by = mock_find_many.call_args.kwargs["where"]
    print(
        f"where_condition with include_created_by_keys=False: {where_condition_no_created_by}"
    )

    # Should not have OR conditions when include_created_by_keys=False and no admin_team_ids
    # The user condition should be merged directly into the where clause
    assert "created_by" not in json.dumps(where_condition_no_created_by)

    # Reset mocks for Test Case 3
    mock_find_many.reset_mock()
    mock_count.reset_mock()

    # Test Case 3: include_created_by_keys=True with exclude_team_id
    args.update(
        {
            "include_created_by_keys": True,
            "exclude_team_id": "excluded-team-123",
            "team_id": None,  # Make sure no specific team is set
        }
    )

    try:
        result = await _list_key_helper(**args)
    except Exception as e:
        print(f"error: {e}")

    where_condition_with_exclude = mock_find_many.call_args.kwargs["where"]
    print(f"where_condition with exclude_team_id: {where_condition_with_exclude}")

    or_conditions_with_exclude = where_condition_with_exclude["AND"][1]["OR"]

    # Find the user condition and created_by condition
    user_condition_with_exclude = None
    created_by_condition_with_exclude = None

    for condition in or_conditions_with_exclude:
        if "user_id" in condition:
            user_condition_with_exclude = condition
        elif "created_by" in condition:
            created_by_condition_with_exclude = condition

    # Verify exclude_team_id is applied to user condition
    assert (
        user_condition_with_exclude is not None
    ), "User condition with exclude should be present"
    assert user_condition_with_exclude["team_id"] == {"not": "excluded-team-123"}

    # Verify created_by condition still only has created_by filter
    assert (
        created_by_condition_with_exclude is not None
    ), "Created by condition with exclude should be present"
    assert created_by_condition_with_exclude["created_by"] == test_user_id
    assert len(created_by_condition_with_exclude) == 1


@pytest.mark.asyncio
async def test_key_token_handling(monkeypatch):
    """
    Test that token handling in key generation follows the expected behavior:
    1. token field should not equal key field
    2. if token_id exists, it should equal token field
    """
    mock_prisma_client = AsyncMock()
    mock_insert_data = AsyncMock(
        return_value=MagicMock(
            token="hashed_token_123", litellm_budget_table=None, object_permission=None
        )
    )
    mock_prisma_client.insert_data = mock_insert_data
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )
    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=MagicMock(
            token="hashed_token_123", litellm_budget_table=None, object_permission=None
        )
    )

    from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_fn,
    )
    from litellm.proxy.proxy_server import prisma_client

    # Use monkeypatch to set the prisma_client
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Test key generation
    response = await generate_key_fn(
        data=GenerateKeyRequest(),
        user_api_key_dict=UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234"
        ),
    )

    # Verify token handling
    assert response.key != response.token, "Token should not equal key"
    if hasattr(response, "token_id"):
        assert (
            response.token == response.token_id
        ), "Token should equal token_id if token_id exists"


@pytest.mark.asyncio
async def test_budget_reset_and_expires_at_first_of_month(monkeypatch):
    """
    Test that when budget_duration, duration, and key_budget_duration are "1mo":
    - budget_reset_at is set to first of next month (standardized reset time)
    - expires is set to approximately 1 month from creation time (exact duration)
    """
    mock_prisma_client = AsyncMock()
    mock_insert_data = AsyncMock(
        return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None)
    )
    mock_prisma_client.insert_data = mock_insert_data
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )
    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None)
    )

    from datetime import datetime, timedelta, timezone

    import pytest

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_helper_fn,
    )
    from litellm.proxy.proxy_server import prisma_client

    # Use monkeypatch to set the prisma_client
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Test key generation with budget_duration="1mo", duration="1mo", key_budget_duration="1mo"
    response = await generate_key_helper_fn(
        request_type="user",
        budget_duration="1mo",
        duration="1mo",
        key_budget_duration="1mo",
        user_id="test_user",
    )

    print(f"response: {response}\n")
    # Get the current date
    now = datetime.now(timezone.utc)

    # Calculate expected reset date (first of next month) for budget_reset_at
    if now.month == 12:
        expected_month = 1
        expected_year = now.year + 1
    else:
        expected_month = now.month + 1
        expected_year = now.year

    # Verify budget_reset_at is set to first of next month (standardized reset time)
    budget_reset_at = response.get("budget_reset_at")
    assert budget_reset_at is not None, "budget_reset_at not found in response"
    assert (
        budget_reset_at.year == expected_year
    ), f"Expected year {expected_year}, got {budget_reset_at.year} for budget_reset_at"
    assert (
        budget_reset_at.month == expected_month
    ), f"Expected month {expected_month}, got {budget_reset_at.month} for budget_reset_at"
    assert (
        budget_reset_at.day == 1
    ), f"Expected day 1, got {budget_reset_at.day} for budget_reset_at"

    # Verify expires is set to approximately 1 month from creation time (exact duration, not standardized)
    expires = response.get("expires")
    assert expires is not None, "expires not found in response"
    # expires should be approximately 1 month from now (same day next month, same time)
    # Allow for some variance due to test execution time (subtract 1 second buffer for timing)
    expected_expires_min = now + timedelta(days=28, seconds=-1)
    expected_expires_max = now + timedelta(days=32)
    assert (
        expected_expires_min <= expires <= expected_expires_max
    ), f"Expected expires to be approximately 1 month from now, got {expires}"


@pytest.mark.asyncio
async def test_key_expiration_exact_duration_hours(monkeypatch):
    """
    Test that key expiration uses exact duration addition, not standardized reset times.
    Specifically tests the bug where "12h" duration would expire at midnight instead of 12 hours from creation.
    """
    mock_prisma_client = AsyncMock()
    mock_insert_data = AsyncMock(
        return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None)
    )
    mock_prisma_client.insert_data = mock_insert_data
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )
    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None)
    )

    from datetime import datetime, timedelta, timezone

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_helper_fn,
    )

    # Use monkeypatch to set the prisma_client
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Test key generation with duration="12h"
    # This should expire exactly 12 hours from creation, not at the next midnight/noon boundary
    response = await generate_key_helper_fn(
        request_type="user",
        duration="12h",
        user_id="test_user",
    )

    expires = response.get("expires")
    assert expires is not None, "expires not found in response"

    # Calculate expected expiration (approximately 12 hours from now)
    # Allow for small variance due to test execution time
    now = datetime.now(timezone.utc)
    expected_expires_min = now + timedelta(hours=11, minutes=59)
    expected_expires_max = now + timedelta(hours=12, minutes=1)

    assert (
        expected_expires_min <= expires <= expected_expires_max
    ), f"Expected expires to be approximately 12 hours from now ({now}), got {expires}. Duration should be exact, not aligned to time boundaries."

    # Verify it's NOT aligned to hour boundaries (e.g., not exactly at :00 minutes)
    # If created at 2:30 PM, it should expire at 2:30 AM, not midnight
    expires_minute = expires.minute
    expires_second = expires.second
    # If the expiration is exactly at :00:00, it might be aligned (though could be coincidence)
    # More importantly, verify the duration is correct
    time_diff = expires - now
    hours_diff = time_diff.total_seconds() / 3600
    assert (
        11.9 <= hours_diff <= 12.1
    ), f"Expected expiration to be approximately 12 hours from creation, got {hours_diff} hours"


@pytest.mark.asyncio
async def test_key_generation_with_object_permission(monkeypatch):
    """Ensure /key/generate correctly handles `object_permission` input by
    1. Creating a record in litellm_objectpermissiontable
    2. Passing the returned `object_permission_id` into the key insert payload
    """
    # --- Setup mocked prisma client ---
    mock_prisma_client = AsyncMock()

    # identity helper for jsonify_object (used inside generate_key_helper_fn)
    mock_prisma_client.jsonify_object = lambda data: data  # type: ignore

    # Mock the prisma_client.db.litellm_objectpermissiontable.create call
    mock_object_permission_create = AsyncMock(
        return_value=MagicMock(object_permission_id="objperm123")
    )
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable.create = (
        mock_object_permission_create
    )

    # Mock prisma_client.insert_data for both user and key tables
    async def _insert_data_side_effect(*args, **kwargs):  # type: ignore
        table_name = kwargs.get("table_name")
        if table_name == "user":
            # minimal attributes accessed later in generate_key_helper_fn
            return MagicMock(models=[], spend=0)
        elif table_name == "key":
            return MagicMock(
                token="hashed_token_456",
                litellm_budget_table=None,
                object_permission=None,
            )
        return MagicMock()

    mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect)

    # Attach the mocked prisma client to the proxy_server module
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # --- Import objects after monkeypatching ---
    from litellm.proxy._types import (
        GenerateKeyRequest,
        LiteLLM_ObjectPermissionBase,
        LitellmUserRoles,
    )
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_fn,
    )

    # --- Call generate_key_fn with object_permission ---
    request_data = GenerateKeyRequest(
        object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["my-vector"])
    )

    await generate_key_fn(
        data=request_data,
        user_api_key_dict=UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-1234",
            user_id="user-1",
        ),
    )

    # --- Assertions ---
    # 1. Object permission creation was triggered
    mock_object_permission_create.assert_called_once()

    # 2. Key insert received the generated object_permission_id
    key_insert_calls = [
        call.kwargs
        for call in mock_prisma_client.insert_data.call_args_list
        if call.kwargs.get("table_name") == "key"
    ]
    assert len(key_insert_calls) == 1
    assert key_insert_calls[0]["data"].get("object_permission_id") == "objperm123"


@pytest.mark.asyncio
async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch):
    """Ensure generate_key_helper_fn passes access_group_ids into the key insert payload."""
    mock_prisma_client = AsyncMock()
    mock_prisma_client.jsonify_object = lambda data: data  # type: ignore
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock(
        return_value=MagicMock(object_permission_id=None)
    )

    captured_key_data = {}

    async def _insert_data_side_effect(*args, **kwargs):
        table_name = kwargs.get("table_name")
        if table_name == "user":
            return MagicMock(models=[], spend=0)
        elif table_name == "key":
            captured_key_data.update(kwargs.get("data", {}))
            return MagicMock(
                token="hashed_token_789",
                litellm_budget_table=None,
                object_permission=None,
                created_at=None,
                updated_at=None,
            )
        return MagicMock()

    mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_helper_fn,
    )

    await generate_key_helper_fn(
        request_type="key",
        table_name="key",
        user_id="test-user",
        access_group_ids=["ag-1", "ag-2"],
    )

    assert captured_key_data.get("access_group_ids") == ["ag-1", "ag-2"]


@pytest.mark.asyncio
async def test_key_generation_with_mcp_tool_permissions(monkeypatch):
    """
    Test that /key/generate correctly handles mcp_tool_permissions in object_permission.

    This test verifies that:
    1. mcp_tool_permissions is accepted in the object_permission field
    2. The field is properly stored in the LiteLLM_ObjectPermissionTable
    3. The key is correctly linked to the object_permission record
    """
    mock_prisma_client = AsyncMock()
    mock_prisma_client.jsonify_object = lambda data: data

    # Track what data is passed to create
    created_permission_data = {}

    async def mock_create(**kwargs):
        created_permission_data.update(kwargs.get("data", {}))
        return MagicMock(object_permission_id="objperm_mcp_123")

    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable.create = mock_create

    async def _insert_data_side_effect(*args, **kwargs):
        table_name = kwargs.get("table_name")
        if table_name == "user":
            return MagicMock(models=[], spend=0)
        elif table_name == "key":
            return MagicMock(
                token="hashed_token_789",
                litellm_budget_table=None,
                object_permission=None,
            )
        return MagicMock()

    mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
        AsyncMock(
            side_effect=lambda object_permission=None, **kwargs: object_permission
        ),
    )

    from litellm.proxy._types import (
        GenerateKeyRequest,
        LiteLLM_ObjectPermissionBase,
        LitellmUserRoles,
    )
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_fn,
    )

    # Create request with mcp_tool_permissions
    request_data = GenerateKeyRequest(
        object_permission=LiteLLM_ObjectPermissionBase(
            mcp_servers=["server_1"],
            mcp_tool_permissions={"server_1": ["tool1", "tool2", "tool3"]},
        )
    )

    await generate_key_fn(
        data=request_data,
        user_api_key_dict=UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-1234",
            user_id="user-mcp-1",
        ),
    )

    # Verify mcp_tool_permissions was stored (serialized to JSON string for GraphQL compatibility)
    assert "mcp_tool_permissions" in created_permission_data
    import json

    assert json.loads(created_permission_data["mcp_tool_permissions"]) == {
        "server_1": ["tool1", "tool2", "tool3"]
    }
    assert created_permission_data["mcp_servers"] == ["server_1"]


@pytest.mark.asyncio
async def test_key_update_object_permissions_existing_permission(monkeypatch):
    """
    Test updating object permissions when a key already has an existing object_permission_id.

    This test verifies that when updating vector stores for a key that already has an
    object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated
    with the new permissions and the object_permission_id remains the same.
    """
    from unittest.mock import AsyncMock, MagicMock

    import pytest

    from litellm.proxy._types import (
        LiteLLM_ObjectPermissionBase,
        LiteLLM_VerificationToken,
    )
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _handle_update_object_permission,
    )

    # Mock prisma client
    mock_prisma_client = AsyncMock()
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Mock existing key with object_permission_id
    existing_key_row = LiteLLM_VerificationToken(
        token="test_token_hash",
        object_permission_id="existing_perm_id_123",
        user_id="user123",
        team_id=None,
    )

    # Mock existing object permission record
    existing_object_permission = MagicMock()
    existing_object_permission.model_dump.return_value = {
        "object_permission_id": "existing_perm_id_123",
        "vector_stores": ["old_store_1", "old_store_2"],
    }

    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
        return_value=existing_object_permission
    )

    # Mock upsert operation
    updated_permission = MagicMock()
    updated_permission.object_permission_id = "existing_perm_id_123"
    mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
        return_value=updated_permission
    )

    # Test data with new object permission
    data_json = {
        "object_permission": LiteLLM_ObjectPermissionBase(
            vector_stores=["new_store_1", "new_store_2", "new_store_3"]
        ).model_dump(exclude_unset=True, exclude_none=True),
        "user_id": "user123",
    }

    # Call the function
    result = await _handle_update_object_permission(
        data_json=data_json,
        existing_key_row=existing_key_row,
    )

    # Verify the object_permission was removed from data_json and object_permission_id was set
    assert "object_permission" not in result
    assert result["object_permission_id"] == "existing_perm_id_123"

    # Verify database operations were called correctly
    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
        where={"object_permission_id": "existing_perm_id_123"}
    )
    mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()


@pytest.mark.asyncio
async def test_key_update_object_permissions_no_existing_permission(monkeypatch):
    """
    Test creating object permissions when a key has no existing object_permission_id.

    This test verifies that when updating object permissions for a key that has
    object_permission_id set to None, a new entry is created in the
    LiteLLM_ObjectPermissionTable and the key is updated with the new object_permission_id.
    """
    from unittest.mock import AsyncMock, MagicMock

    import pytest

    from litellm.proxy._types import (
        LiteLLM_ObjectPermissionBase,
        LiteLLM_VerificationToken,
    )
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _handle_update_object_permission,
    )

    # Mock prisma client
    mock_prisma_client = AsyncMock()
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    existing_key_row_no_perm = LiteLLM_VerificationToken(
        token="test_token_hash_2",
        object_permission_id=None,
        user_id="user456",
        team_id=None,
    )

    # Mock find_unique to return None (no existing permission)
    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
        return_value=None
    )

    # Mock upsert to create new record
    new_permission = MagicMock()
    new_permission.object_permission_id = "new_perm_id_456"
    mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
        return_value=new_permission
    )

    data_json = {
        "object_permission": LiteLLM_ObjectPermissionBase(
            vector_stores=["brand_new_store"]
        ).model_dump(exclude_unset=True, exclude_none=True),
        "user_id": "user456",
    }

    result = await _handle_update_object_permission(
        data_json=data_json,
        existing_key_row=existing_key_row_no_perm,
    )

    # Verify new object_permission_id was set
    assert "object_permission" not in result
    assert result["object_permission_id"] == "new_perm_id_456"
    # Verify upsert was called to create new record
    mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()


@pytest.mark.asyncio
async def test_key_update_object_permissions_missing_permission_record(monkeypatch):
    """
    Test creating object permissions when existing object_permission_id record is not found.

    This test verifies that when updating object permissions for a key that has an
    object_permission_id but the corresponding record cannot be found in the database,
    a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions.
    """
    from unittest.mock import AsyncMock, MagicMock

    import pytest

    from litellm.proxy._types import (
        LiteLLM_ObjectPermissionBase,
        LiteLLM_VerificationToken,
    )
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _handle_update_object_permission,
    )

    # Mock prisma client
    mock_prisma_client = AsyncMock()
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    existing_key_row_missing_perm = LiteLLM_VerificationToken(
        token="test_token_hash_3",
        object_permission_id="missing_perm_id_789",
        user_id="user789",
        team_id=None,
    )

    # Mock find_unique to return None (permission record not found)
    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
        return_value=None
    )

    # Mock upsert to create new record
    new_permission = MagicMock()
    new_permission.object_permission_id = "recreated_perm_id_789"
    mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
        return_value=new_permission
    )

    data_json = {
        "object_permission": LiteLLM_ObjectPermissionBase(
            vector_stores=["recreated_store"]
        ).model_dump(exclude_unset=True, exclude_none=True),
        "user_id": "user789",
    }

    result = await _handle_update_object_permission(
        data_json=data_json,
        existing_key_row=existing_key_row_missing_perm,
    )

    # Verify new object_permission_id was set
    assert "object_permission" not in result
    assert result["object_permission_id"] == "recreated_perm_id_789"

    # Verify find_unique was called with the missing permission ID
    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
        where={"object_permission_id": "missing_perm_id_789"}
    )

    # Verify upsert was called to create new record
    mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()


@pytest.mark.asyncio
async def test_key_update_object_permission_does_not_add_null_fields():
    """
    Updating a key with an object_permission that only sets a subset of fields
    must not normalize the unset list fields to ``None``.

    The UI always submits object_permission with empty MCP/vector lists, even for
    a TPM/RPM-only edit. ``models``/``blocked_tools``/``search_tools`` are
    non-nullable array columns, so emitting them as ``None`` makes the downstream
    Prisma write fail. The normalized object_permission must keep the same field
    set the caller provided.
    """
    data = UpdateKeyRequest(
        key="sk-test-key",
        tpm_limit=123,
        rpm_limit=456,
        object_permission={
            "vector_stores": [],
            "mcp_servers": [],
            "mcp_access_groups": [],
            "mcp_toolsets": [],
            "agents": [],
            "agent_access_groups": [],
        },
    )
    provided_fields = set(data.object_permission.model_fields_set)

    existing_key_row = MagicMock()
    existing_key_row.user_id = "admin_user"
    existing_key_row.token = "hashed_token"
    existing_key_row.team_id = None
    existing_key_row.organization_id = None
    existing_key_row.project_id = None

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-admin",
        user_id="admin_user",
    )

    await _validate_update_key_data(
        data=data,
        existing_key_row=existing_key_row,
        user_api_key_dict=user_api_key_dict,
        llm_router=None,
        premium_user=False,
        prisma_client=AsyncMock(),
        user_api_key_cache=MagicMock(),
    )

    normalized = data.object_permission.model_dump(exclude_unset=True)
    assert set(normalized.keys()) == provided_fields
    assert "models" not in normalized
    assert "blocked_tools" not in normalized
    assert "search_tools" not in normalized
    assert "mcp_tool_permissions" not in normalized


@pytest.mark.asyncio
async def test_key_info_returns_object_permission(monkeypatch):
    """
    Test that /key/info correctly returns the object_permission relation.

    This test verifies that when calling /key/info for a key with object_permission_id,
    the response includes the full object_permission object with fields like
    mcp_access_groups, mcp_servers, vector_stores, agents, etc.

    Regression test for bug where object_permission_id was returned but not the
    related object_permission object.
    """
    from unittest.mock import AsyncMock, MagicMock

    import pytest

    from litellm.proxy._types import LiteLLM_VerificationToken
    from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn

    # Mock prisma client
    mock_prisma_client = AsyncMock()
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Mock key with object_permission_id
    test_key_token = "hashed_test_token_123"
    test_object_permission_id = "objperm_info_test_123"

    mock_key_info = MagicMock(spec=LiteLLM_VerificationToken)
    mock_key_info.token = test_key_token
    mock_key_info.object_permission_id = test_object_permission_id
    mock_key_info.user_id = "user123"
    mock_key_info.team_id = None
    mock_key_info.litellm_budget_table = None

    # Mock the dict/model_dump methods
    mock_key_info.model_dump.return_value = {
        "token": test_key_token,
        "object_permission_id": test_object_permission_id,
        "user_id": "user123",
        "team_id": None,
        "litellm_budget_table": None,
    }
    mock_key_info.dict.return_value = mock_key_info.model_dump.return_value

    # Mock find_unique for the key lookup
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_key_info
    )

    # Mock object permission record
    mock_object_permission = MagicMock()
    mock_object_permission.model_dump.return_value = {
        "object_permission_id": test_object_permission_id,
        "mcp_access_groups": ["test_group_1", "test_group_2"],
        "mcp_servers": ["server_1"],
        "vector_stores": ["vs_1", "vs_2"],
        "agents": ["agent_1"],
    }
    mock_object_permission.dict.return_value = (
        mock_object_permission.model_dump.return_value
    )

    # Mock find_unique for object permission lookup
    mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
        return_value=mock_object_permission
    )

    # Create user API key dict
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-test-key-456",
    )

    # Call info_key_fn
    result = await info_key_fn(
        key="sk-test-key-456",
        user_api_key_dict=user_api_key_dict,
    )

    # Assertions
    assert "info" in result
    assert "object_permission_id" in result["info"]
    assert result["info"]["object_permission_id"] == test_object_permission_id

    # CRITICAL: Verify that object_permission object is included in response
    assert "object_permission" in result["info"], (
        "object_permission field missing from /key/info response. "
        "Expected full object_permission object to be attached."
    )

    # Verify object_permission contains the expected fields
    obj_perm = result["info"]["object_permission"]
    assert obj_perm["object_permission_id"] == test_object_permission_id
    assert obj_perm["mcp_access_groups"] == ["test_group_1", "test_group_2"]
    assert obj_perm["mcp_servers"] == ["server_1"]
    assert obj_perm["vector_stores"] == ["vs_1", "vs_2"]
    assert obj_perm["agents"] == ["agent_1"]

    # Verify the object permission was actually queried from database
    mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
        where={"object_permission_id": test_object_permission_id}
    )


@pytest.mark.asyncio
async def test_get_new_token_with_valid_key(monkeypatch):
    """Test get_new_token function when provided with a valid key that starts with 'sk-'"""
    from unittest.mock import AsyncMock

    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        get_new_token,
    )

    # Mock get_ui_settings_cached to return setting disabled (custom keys allowed)
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={}),
    )

    # Test with valid new_key
    data = RegenerateKeyRequest(new_key="sk-test123456789")
    result = await get_new_token(data)

    assert result == "sk-test123456789"


@pytest.mark.asyncio
async def test_get_new_token_with_invalid_key(monkeypatch):
    """Test get_new_token function when provided with an invalid key that doesn't start with 'sk-'"""
    from unittest.mock import AsyncMock

    from fastapi import HTTPException

    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        get_new_token,
    )

    # Mock get_ui_settings_cached to return setting disabled (custom keys allowed)
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={}),
    )

    # Test with invalid new_key (doesn't start with 'sk-')
    data = RegenerateKeyRequest(new_key="invalid-key-123")

    with pytest.raises(HTTPException) as exc_info:
        await get_new_token(data)

    assert exc_info.value.status_code == 400
    assert "New key must start with 'sk-'" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_check_custom_key_allowed_when_disabled(monkeypatch):
    """_check_custom_key_allowed raises 403 when disable_custom_api_keys is true."""
    from unittest.mock import AsyncMock

    from fastapi import HTTPException

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _check_custom_key_allowed,
    )

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={"disable_custom_api_keys": True}),
    )

    with pytest.raises(HTTPException) as exc_info:
        await _check_custom_key_allowed("sk-custom-key-123")

    assert exc_info.value.status_code == 403
    assert "disabled" in str(exc_info.value.detail).lower()


@pytest.mark.asyncio
async def test_check_custom_key_allowed_when_enabled(monkeypatch):
    """_check_custom_key_allowed does nothing when disable_custom_api_keys is false."""
    from unittest.mock import AsyncMock

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _check_custom_key_allowed,
    )

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={"disable_custom_api_keys": False}),
    )

    # Should not raise
    await _check_custom_key_allowed("sk-custom-key-123")


@pytest.mark.asyncio
async def test_check_custom_key_allowed_when_unset(monkeypatch):
    """_check_custom_key_allowed does nothing when setting is not present."""
    from unittest.mock import AsyncMock

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _check_custom_key_allowed,
    )

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={}),
    )

    # Should not raise
    await _check_custom_key_allowed("sk-custom-key-123")


@pytest.mark.asyncio
async def test_check_custom_key_allowed_none_key_always_passes(monkeypatch):
    """_check_custom_key_allowed does nothing when key is None, even if setting is on."""
    from unittest.mock import AsyncMock

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _check_custom_key_allowed,
    )

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={"disable_custom_api_keys": True}),
    )

    # Should not raise — None means auto-generate
    await _check_custom_key_allowed(None)


@pytest.mark.asyncio
async def test_get_new_token_rejected_when_custom_keys_disabled(monkeypatch):
    """get_new_token raises 403 when new_key is set and disable_custom_api_keys is true."""
    from unittest.mock import AsyncMock

    from fastapi import HTTPException

    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        get_new_token,
    )

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={"disable_custom_api_keys": True}),
    )

    data = RegenerateKeyRequest(new_key="sk-custom-regen-key")

    with pytest.raises(HTTPException) as exc_info:
        await get_new_token(data)

    assert exc_info.value.status_code == 403


@pytest.mark.asyncio
async def test_get_new_token_auto_generates_when_custom_keys_disabled(monkeypatch):
    """get_new_token auto-generates a key when new_key is None, even if setting is on."""
    from unittest.mock import AsyncMock

    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        get_new_token,
    )

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_ui_settings_cached",
        AsyncMock(return_value={"disable_custom_api_keys": True}),
    )

    data = RegenerateKeyRequest()  # no new_key
    result = await get_new_token(data)

    assert result.startswith("sk-")


@pytest.mark.asyncio
async def test_generate_service_account_requires_team_id():
    with pytest.raises(HTTPException):
        await _common_key_generation_helper(
            data=GenerateKeyRequest(
                metadata={"service_account_id": "sa"},
                team_id=None,
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
            ),
            litellm_changed_by=None,
            team_table=None,
        )


@pytest.mark.asyncio
async def test_generate_service_account_works_with_team_id():
    from unittest.mock import patch

    # Mock the database and router dependencies from proxy_server
    with (
        patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
        patch("litellm.proxy.proxy_server.llm_router") as mock_router,
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
        ) as mock_generate_key,
    ):

        # Configure mocks
        mock_prisma.return_value = AsyncMock()
        mock_router.return_value = None
        # Mock the response from generate_key_helper_fn
        mock_generate_key.return_value = {
            "key": "sk-test-key",
            "expires": None,
            "user_id": "test-user",
            "team_id": "IJ",
        }

        # This should not raise an exception since team_id is provided
        await _common_key_generation_helper(
            data=GenerateKeyRequest(
                metadata={"service_account_id": "sa"},
                team_id="IJ",
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
            ),
            litellm_changed_by=None,
            team_table=None,
        )


@pytest.mark.asyncio
async def test_update_service_account_requires_team_id():
    data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"})
    existing_key = LiteLLM_VerificationToken(token="hashed", team_id=None)

    with pytest.raises(HTTPException):
        await prepare_key_update_data(data=data, existing_key_row=existing_key)


@pytest.mark.asyncio
async def test_update_service_account_works_with_team_id():
    data = UpdateKeyRequest(
        key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ"
    )
    existing_key = LiteLLM_VerificationToken(token="hashed")

    await prepare_key_update_data(data=data, existing_key_row=existing_key)


@pytest.mark.asyncio
async def test_update_preserves_service_account_id_when_metadata_replaced():
    """
    Regression: /key/update wholesale-replaced metadata, silently dropping
    service_account_id. The pre-call check then treated the key as a regular
    key and bypassed service_account_settings.enforced_params.
    """
    data = UpdateKeyRequest(
        key="sk-1",
        metadata={"unrelated": "value"},
        team_id="IJ",
    )
    existing_key = LiteLLM_VerificationToken(
        token="hashed",
        team_id="IJ",
        metadata={"service_account_id": "sa-123"},
    )

    result = await prepare_key_update_data(data=data, existing_key_row=existing_key)

    assert result["metadata"]["service_account_id"] == "sa-123"
    assert result["metadata"]["unrelated"] == "value"


@pytest.mark.asyncio
async def test_update_rejects_service_account_id_overwrite():
    """
    Once assigned, a key's service_account_id is an identity marker — rebinding
    it via update would break spend attribution. Reject rather than silently
    ignore so scripted callers surface the bug.
    """
    data = UpdateKeyRequest(
        key="sk-1",
        metadata={"service_account_id": "sa-new"},
        team_id="IJ",
    )
    existing_key = LiteLLM_VerificationToken(
        token="hashed",
        team_id="IJ",
        metadata={"service_account_id": "sa-old"},
    )

    with pytest.raises(HTTPException) as exc_info:
        await prepare_key_update_data(data=data, existing_key_row=existing_key)
    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_update_allows_matching_service_account_id():
    """Resending the same value (e.g. UI round-trip) is a no-op, not a conflict."""
    data = UpdateKeyRequest(
        key="sk-1",
        metadata={"service_account_id": "sa-123", "other": "value"},
        team_id="IJ",
    )
    existing_key = LiteLLM_VerificationToken(
        token="hashed",
        team_id="IJ",
        metadata={"service_account_id": "sa-123"},
    )

    result = await prepare_key_update_data(data=data, existing_key_row=existing_key)

    assert result["metadata"]["service_account_id"] == "sa-123"
    assert result["metadata"]["other"] == "value"


@pytest.mark.asyncio
async def test_update_rejects_explicit_null_service_account_id():
    """
    Explicit null is an attempt to clear — not an omission. Silently ignoring
    it would let a caller think they cleared the field when they didn't, so
    treat it the same as any other rebind attempt and return 400.
    """
    data = UpdateKeyRequest(
        key="sk-1",
        metadata={"service_account_id": None},
        team_id="IJ",
    )
    existing_key = LiteLLM_VerificationToken(
        token="hashed",
        team_id="IJ",
        metadata={"service_account_id": "sa-old"},
    )

    with pytest.raises(HTTPException) as exc_info:
        await prepare_key_update_data(data=data, existing_key_row=existing_key)
    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_update_rejects_whole_metadata_null_on_service_account_key():
    """
    `metadata: null` on a key with a reserved field would have dereferenced None
    inside the reserved-field loop and returned a 500. Must surface as a 400
    since the effect would be to clear an immutable field.
    """
    data = UpdateKeyRequest(key="sk-1", metadata=None, team_id="IJ")
    existing_key = LiteLLM_VerificationToken(
        token="hashed",
        team_id="IJ",
        metadata={"service_account_id": "sa-123"},
    )

    with pytest.raises(HTTPException) as exc_info:
        await prepare_key_update_data(data=data, existing_key_row=existing_key)
    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_update_without_metadata_still_preserves_existing():
    """Omitting metadata entirely must not drop existing metadata fields."""
    data = UpdateKeyRequest(key="sk-1", max_budget=100)
    existing_key = LiteLLM_VerificationToken(
        token="hashed",
        team_id="IJ",
        metadata={"service_account_id": "sa-123", "other": "kept"},
    )

    result = await prepare_key_update_data(data=data, existing_key_row=existing_key)

    assert result["metadata"]["service_account_id"] == "sa-123"
    assert result["metadata"]["other"] == "kept"


@pytest.mark.asyncio
async def test_prepare_key_update_data_encrypts_callback_vars(monkeypatch):
    """/key/update must encrypt callback_vars values before they reach the DB."""
    from litellm.proxy.common_utils.callback_utils import decrypt_callback_vars

    monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-32-bytes-aaaaaaaaaaaaaa")
    data = UpdateKeyRequest(
        key="sk-1",
        metadata={
            "logging": [
                {
                    "callback_name": "langfuse",
                    "callback_type": "success",
                    "callback_vars": {
                        "langfuse_public_key": "pk-real",
                        "langfuse_secret_key": "sk-real",
                    },
                }
            ]
        },
    )
    existing_key = LiteLLM_VerificationToken(token="hashed")

    result = await prepare_key_update_data(data=data, existing_key_row=existing_key)

    cv = result["metadata"]["logging"][0]["callback_vars"]
    assert cv["langfuse_secret_key"] != "sk-real"
    assert cv["langfuse_public_key"] != "pk-real"
    recovered = decrypt_callback_vars(result["metadata"])["logging"][0]["callback_vars"]
    assert recovered["langfuse_secret_key"] == "sk-real"
    assert recovered["langfuse_public_key"] == "pk-real"


@pytest.mark.asyncio
async def test_prepare_key_update_data_duration_never_expires():
    """Test that duration="-1" sets expires to None (never expires)."""
    from litellm.proxy._types import UpdateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        prepare_key_update_data,
    )

    # Mock existing key
    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=["gpt-3.5-turbo"],
        user_id="test-user",
        team_id=None,
        auto_rotate=False,
        rotation_interval=None,
        metadata={},
    )

    # Test setting duration to "-1" (never expires)
    update_request = UpdateKeyRequest(key="test-token", duration="-1")

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    # Verify that expires is set to None
    assert result["expires"] is None


@pytest.mark.asyncio
async def test_prepare_key_update_data_duration_none_never_expires():
    """Test that duration=None sets expires to None (never expires)."""
    from litellm.proxy._types import UpdateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        prepare_key_update_data,
    )

    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=["gpt-3.5-turbo"],
        user_id="test-user",
        team_id=None,
        auto_rotate=False,
        rotation_interval=None,
        metadata={},
    )

    update_request = UpdateKeyRequest(key="test-token", duration=None)

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    assert result["expires"] is None


@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_requires_team_id():
    """
    Test that validate_team_id_used_in_service_account_request raises HTTPException
    when team_id is None for service account key generation.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        validate_team_id_used_in_service_account_request,
    )

    mock_prisma_client = AsyncMock()

    # Test that HTTPException is raised when team_id is None
    with pytest.raises(HTTPException) as exc_info:
        await validate_team_id_used_in_service_account_request(
            team_id=None,
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.status_code == 400
    assert "team_id is required for service account keys" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_requires_prisma_client():
    """
    Test that validate_team_id_used_in_service_account_request raises HTTPException
    when prisma_client is None for service account key generation.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        validate_team_id_used_in_service_account_request,
    )

    # Test that HTTPException is raised when prisma_client is None
    with pytest.raises(HTTPException) as exc_info:
        await validate_team_id_used_in_service_account_request(
            team_id="test-team-id",
            prisma_client=None,
        )

    assert exc_info.value.status_code == 400
    assert "prisma_client is required for service account keys" in str(
        exc_info.value.detail
    )


@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_checks_team_exists():
    """
    Test that validate_team_id_used_in_service_account_request validates that
    the team_id exists in the database for service account key generation.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        validate_team_id_used_in_service_account_request,
    )

    mock_prisma_client = AsyncMock()

    # Mock the database query to return None (team doesn't exist)
    mock_find_unique = AsyncMock(return_value=None)
    mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique

    # Test that HTTPException is raised when team doesn't exist in DB
    with pytest.raises(HTTPException) as exc_info:
        await validate_team_id_used_in_service_account_request(
            team_id="non-existent-team-id",
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.status_code == 400
    assert "team_id does not exist in the database" in str(exc_info.value.detail)

    # Verify the database was queried with the correct parameters
    mock_find_unique.assert_called_once_with(where={"team_id": "non-existent-team-id"})


@pytest.mark.asyncio
async def test_validate_team_id_used_in_service_account_request_success():
    """
    Test that validate_team_id_used_in_service_account_request returns True
    when team_id exists in the database for service account key generation.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        validate_team_id_used_in_service_account_request,
    )

    mock_prisma_client = AsyncMock()

    # Mock the database query to return a team object (team exists)
    mock_team = {"team_id": "existing-team-id", "team_name": "Test Team"}
    mock_find_unique = AsyncMock(return_value=mock_team)
    mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique

    # Test that function returns True when team exists
    result = await validate_team_id_used_in_service_account_request(
        team_id="existing-team-id",
        prisma_client=mock_prisma_client,
    )

    assert result is True

    # Verify the database was queried with the correct parameters
    mock_find_unique.assert_called_once_with(where={"team_id": "existing-team-id"})


@pytest.mark.asyncio
async def test_generate_service_account_key_endpoint_validation():
    """
    Test that the /key/service-account/generate endpoint properly validates
    team_id requirement and team existence in database.
    """
    from unittest.mock import patch

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_service_account_key_fn,
    )

    # Test case 1: Missing team_id
    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
        # Mock prisma_client to be not None so we can reach team_id validation
        mock_prisma_instance = AsyncMock()
        mock_prisma.return_value = mock_prisma_instance

        with pytest.raises(HTTPException) as exc_info:
            await generate_service_account_key_fn(
                data=GenerateKeyRequest(team_id=None),
                user_api_key_dict=UserAPIKeyAuth(
                    user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
                ),
                litellm_changed_by=None,
            )

        assert exc_info.value.status_code == 400
        assert "team_id is required for service account keys" in str(
            exc_info.value.detail
        )

    # Test case 2: Team doesn't exist in database
    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
        # Mock team not found
        mock_find_unique = AsyncMock(return_value=None)
        mock_prisma.db.litellm_teamtable.find_unique = mock_find_unique

        with pytest.raises(HTTPException) as exc_info:
            await generate_service_account_key_fn(
                data=GenerateKeyRequest(team_id="non-existent-team"),
                user_api_key_dict=UserAPIKeyAuth(
                    user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"
                ),
                litellm_changed_by=None,
            )

        assert exc_info.value.status_code == 400
        assert "team_id does not exist in the database" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
    """
    Test that the unblock_key endpoint correctly handles both sk- prefixed tokens
    and hashed tokens by properly converting sk- tokens to hashed format before
    database operations.
    """
    from unittest.mock import AsyncMock, MagicMock

    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key

    # Mock dependencies
    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    # Use a proper 64-character hex hash for testing
    test_hashed_token = (
        "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
    )

    # Mock the key record that will be returned from database
    mock_key_record = MagicMock()
    mock_key_record.token = test_hashed_token
    mock_key_record.blocked = False
    mock_key_record.model_dump_json.return_value = (
        f'{{"token": "{test_hashed_token}", "blocked": false}}'
    )

    # Mock database operations
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_key_record
    )
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=mock_key_record
    )

    # Mock hash_token function
    def mock_hash_token(token):
        if token == "sk-test123456789":
            return test_hashed_token
        return token

    # Apply monkeypatch
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
    monkeypatch.setattr(
        "litellm.store_audit_logs", False
    )  # Disable audit logs for simpler test

    # Mock get_key_object and _cache_key_object
    async def mock_delete_cache_key_object(**kwargs):
        pass

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
        mock_delete_cache_key_object,
    )

    # Create mock request and user auth
    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
    )

    # Test Case 1: Using sk- prefixed token
    sk_token_request = BlockKeyRequest(key="sk-test123456789")

    result = await unblock_key(
        data=sk_token_request,
        http_request=mock_request,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by=None,
    )

    # Verify that the database update was called with hashed token
    mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with(
        where={"token": test_hashed_token}, data={"blocked": False}
    )

    assert result == mock_key_record

    # Reset mocks for second test
    mock_prisma_client.db.litellm_verificationtoken.update.reset_mock()

    # Test Case 2: Using already hashed token
    hashed_token_request = BlockKeyRequest(key=test_hashed_token)

    result = await unblock_key(
        data=hashed_token_request,
        http_request=mock_request,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by=None,
    )

    # Verify that the database update was called with the same hashed token
    mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with(
        where={"token": test_hashed_token}, data={"blocked": False}
    )

    assert result == mock_key_record


@pytest.mark.asyncio
async def test_unblock_key_invalid_key_format(monkeypatch):
    """
    Test that unblock_key properly validates key format and raises appropriate errors
    for invalid keys.
    """
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key
    from litellm.proxy.utils import ProxyException

    # Mock prisma_client to avoid DB connection error
    mock_prisma_client = AsyncMock()
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Mock request and user auth
    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
    )

    # Test with invalid key format
    invalid_key_request = BlockKeyRequest(key="invalid-key-format")

    with pytest.raises(ProxyException) as exc_info:
        await unblock_key(
            data=invalid_key_request,
            http_request=mock_request,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc_info.value.code == "400"
    assert "Invalid key format" in str(exc_info.value.message)


@pytest.mark.asyncio
async def test_block_key_nonexistent_key_returns_404(monkeypatch):
    """
    Test that block_key returns 404 (not misleading 401) when the key
    doesn't exist in the database, even when the caller is authenticated
    as a proxy admin.

    Previously, block_key would call get_key_object() for cache refresh,
    which raised a 401 ProxyException with 'Authentication Error' — making
    it look like an auth failure when it was really a missing-key error.
    """
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import block_key

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    # find_unique returns None → key does not exist
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )

    def mock_hash_token(token):
        return "abcd1234" * 8  # 64-char hex

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
    monkeypatch.setattr("litellm.store_audit_logs", False)

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
    )

    data = BlockKeyRequest(key="sk-does-not-exist-key")

    with pytest.raises(ProxyException) as exc_info:
        await block_key(
            data=data,
            http_request=mock_request,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc_info.value.code == "404"
    assert "not found" in str(exc_info.value.message).lower()
    # Must NOT contain "Authentication Error"
    assert "Authentication Error" not in str(exc_info.value.message)
    # update should never be called since the key doesn't exist
    mock_prisma_client.db.litellm_verificationtoken.update.assert_not_called()


@pytest.mark.asyncio
async def test_unblock_key_nonexistent_key_returns_404(monkeypatch):
    """
    Test that unblock_key returns 404 (not misleading 401) when the key
    doesn't exist in the database.
    """
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        unblock_key,
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    # find_unique returns None → key does not exist
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )

    def mock_hash_token(token):
        return "abcd1234" * 8

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
    monkeypatch.setattr("litellm.store_audit_logs", False)

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
    )

    data = BlockKeyRequest(key="sk-does-not-exist-key")

    with pytest.raises(ProxyException) as exc_info:
        await unblock_key(
            data=data,
            http_request=mock_request,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc_info.value.code == "404"
    assert "not found" in str(exc_info.value.message).lower()
    assert "Authentication Error" not in str(exc_info.value.message)
    mock_prisma_client.db.litellm_verificationtoken.update.assert_not_called()


@pytest.mark.asyncio
async def test_update_key_nonexistent_key_returns_404(monkeypatch):
    """
    Test that update_key_fn returns 404 (not misleading 401) when the body
    key doesn't exist in the database, even when the caller is authenticated
    as a proxy admin via the Authorization header.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    # find_unique returns None → key does not exist
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
    )

    data = UpdateKeyRequest(key="sk-does-not-exist-key")

    with pytest.raises(ProxyException) as exc_info:
        await update_key_fn(
            request=mock_request,
            data=data,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc_info.value.code == "404"
    assert "not found" in str(exc_info.value.message).lower()
    assert "Authentication Error" not in str(exc_info.value.message)


@pytest.mark.asyncio
async def test_block_key_existing_key_succeeds(monkeypatch):
    """
    Test that block_key successfully blocks an existing key and
    invalidates the cache entry.
    """
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import block_key

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    test_hashed_token = (
        "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
    )

    mock_key_record = MagicMock()
    mock_key_record.token = test_hashed_token
    mock_key_record.blocked = False
    mock_key_record.model_dump_json.return_value = (
        f'{{"token": "{test_hashed_token}", "blocked": false}}'
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_key_record
    )
    mock_updated_record = MagicMock()
    mock_updated_record.token = test_hashed_token
    mock_updated_record.blocked = True
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=mock_updated_record
    )

    def mock_hash_token(token):
        if token.startswith("sk-"):
            return test_hashed_token
        return token

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
    monkeypatch.setattr("litellm.store_audit_logs", False)

    # Mock _delete_cache_key_object
    async def mock_delete_cache_key_object(**kwargs):
        pass

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
        mock_delete_cache_key_object,
    )

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
    )

    data = BlockKeyRequest(key="sk-test123456789")

    result = await block_key(
        data=data,
        http_request=mock_request,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by=None,
    )

    # Verify the key was found and updated
    mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with(
        where={"token": test_hashed_token}
    )
    mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once_with(
        where={"token": test_hashed_token}, data={"blocked": True}
    )
    assert result == mock_updated_record


@pytest.mark.asyncio
async def test_validate_key_team_change_with_member_permissions():
    """
    Test validate_key_team_change function with team member permissions.

    This test covers the new logic that allows team members with specific
    permissions to update keys, not just team admins.
    """
    from unittest.mock import AsyncMock, MagicMock, patch

    from litellm.proxy._types import KeyManagementRoutes

    # Create mock objects
    mock_key = MagicMock()
    mock_key.user_id = "test-user-123"
    mock_key.models = ["gpt-4"]
    mock_key.tpm_limit = None
    mock_key.rpm_limit = None

    mock_team = MagicMock()
    mock_team.team_id = "test-team-456"
    mock_team.members_with_roles = []
    mock_team.tpm_limit = None
    mock_team.rpm_limit = None

    mock_change_initiator = MagicMock()
    mock_change_initiator.user_id = "test-user-123"

    mock_router = MagicMock()

    # Mock the member object returned by _get_user_in_team
    mock_member_object = MagicMock()

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model",
        new_callable=AsyncMock,
    ):
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team"
        ) as mock_get_user:
            with patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin"
            ) as mock_is_admin:
                with patch(
                    "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint"
                ) as mock_has_perms:

                    mock_get_user.return_value = mock_member_object
                    mock_is_admin.return_value = False
                    mock_has_perms.return_value = True

                    # This should not raise an exception due to member permissions
                    await validate_key_team_change(
                        key=mock_key,
                        team=mock_team,
                        change_initiated_by=mock_change_initiator,
                        llm_router=mock_router,
                    )

                    # Verify the permission check was called with correct parameters
                    mock_has_perms.assert_called_once_with(
                        team_member_object=mock_member_object,
                        team_table=mock_team,
                        route=KeyManagementRoutes.KEY_UPDATE.value,
                    )


@pytest.mark.asyncio
async def test_validate_key_team_change_skips_all_team_models_sentinel():
    """
    Test that validate_key_team_change skips the 'all-team-models' sentinel
    value when checking if the target team can access the key's models.

    Keys with models=["all-team-models"] mean "use whatever models the team
    allows", so moving them to any team should not fail model validation.
    """
    from unittest.mock import AsyncMock, MagicMock, patch

    mock_key = MagicMock()
    mock_key.user_id = "test-user-123"
    mock_key.models = ["all-team-models"]
    mock_key.tpm_limit = None
    mock_key.rpm_limit = None

    mock_team = MagicMock()
    mock_team.team_id = "test-team-456"
    mock_team.models = ["gpt-4", "claude-3"]
    mock_team.members_with_roles = []
    mock_team.tpm_limit = None
    mock_team.rpm_limit = None

    mock_change_initiator = MagicMock()
    mock_change_initiator.user_id = "test-user-123"
    mock_change_initiator.user_role = LitellmUserRoles.PROXY_ADMIN.value

    mock_router = MagicMock()

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model",
        new_callable=AsyncMock,
    ) as mock_can_access:
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team"
        ) as mock_get_user:
            mock_get_user.return_value = MagicMock()

            await validate_key_team_change(
                key=mock_key,
                team=mock_team,
                change_initiated_by=mock_change_initiator,
                llm_router=mock_router,
            )

            # can_team_access_model should NOT have been called since
            # "all-team-models" is a sentinel that should be skipped
            mock_can_access.assert_not_called()


def test_key_rotation_fields_helper():
    """
    Test the key data update logic for rotation fields.

    This test focuses on the core logic that adds rotation fields to key_data
    when auto_rotate is enabled, without the complexity of full key generation.
    """
    # Test Case 1: With rotation enabled
    key_data = {"models": ["gpt-3.5-turbo"], "user_id": "test-user"}

    auto_rotate = True
    rotation_interval = "30d"

    # Simulate the rotation logic from generate_key_helper_fn
    if auto_rotate and rotation_interval:
        key_data.update(
            {"auto_rotate": auto_rotate, "rotation_interval": rotation_interval}
        )

    # Verify rotation fields are added
    assert key_data["auto_rotate"] == True
    assert key_data["rotation_interval"] == "30d"
    assert key_data["models"] == ["gpt-3.5-turbo"]  # Original fields preserved

    # Test Case 2: Without rotation enabled
    key_data2 = {"models": ["gpt-4"], "user_id": "test-user"}

    auto_rotate2 = False
    rotation_interval2 = None

    # Simulate the rotation logic
    if auto_rotate2 and rotation_interval2:
        key_data2.update(
            {"auto_rotate": auto_rotate2, "rotation_interval": rotation_interval2}
        )

    # Verify rotation fields are NOT added
    assert "auto_rotate" not in key_data2
    assert "rotation_interval" not in key_data2
    assert key_data2["models"] == ["gpt-4"]  # Original fields preserved

    # Test Case 3: auto_rotate=True but no interval
    key_data3 = {"models": ["claude-3"], "user_id": "test-user"}

    auto_rotate3 = True
    rotation_interval3 = None

    # Simulate the rotation logic
    if auto_rotate3 and rotation_interval3:
        key_data3.update(
            {"auto_rotate": auto_rotate3, "rotation_interval": rotation_interval3}
        )

    # Verify rotation fields are NOT added (missing interval)
    assert "auto_rotate" not in key_data3
    assert "rotation_interval" not in key_data3


@pytest.mark.asyncio
async def test_update_key_fn_auto_rotate_enable():
    """Test that update_key_fn properly handles enabling auto rotation."""
    from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        prepare_key_update_data,
    )

    # Mock existing key
    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=["gpt-3.5-turbo"],
        user_id="test-user",
        team_id=None,
        auto_rotate=False,
        rotation_interval=None,
        metadata={},
    )

    # Test enabling auto rotation
    update_request = UpdateKeyRequest(
        key="test-token", auto_rotate=True, rotation_interval="30d"
    )

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    # Verify rotation fields are included
    assert result["auto_rotate"] is True
    assert result["rotation_interval"] == "30d"


@pytest.mark.asyncio
async def test_update_key_fn_auto_rotate_disable():
    """Test that update_key_fn properly handles disabling auto rotation."""
    from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        prepare_key_update_data,
    )

    # Mock existing key with rotation enabled
    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=["gpt-3.5-turbo"],
        user_id="test-user",
        team_id=None,
        auto_rotate=True,
        rotation_interval="30d",
        metadata={},
    )

    # Test disabling auto rotation
    update_request = UpdateKeyRequest(key="test-token", auto_rotate=False)

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    # Verify auto_rotate is set to False
    assert result["auto_rotate"] is False


@pytest.mark.asyncio
async def test_check_team_key_limits_no_existing_keys():
    """
    Test _check_team_key_limits when team has no existing keys.
    Should allow any TPM/RPM limits within team bounds.
    """
    # Mock prisma client
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )

    # Create team table with limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=10000,
        rpm_limit=1000,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request with limits within team bounds
    data = GenerateKeyRequest(
        tpm_limit=5000,
        rpm_limit=500,
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
    )

    # Should not raise any exception
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )

    # Verify database was queried
    mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
        where={"team_id": "test-team-123"}
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_with_existing_keys_within_bounds():
    """
    Test _check_team_key_limits when team has existing keys but total allocation
    is still within team limits.
    """
    # Create mock existing keys
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 3000
    existing_key1.rpm_limit = 200

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = 2000
    existing_key2.rpm_limit = 300

    existing_key3 = MagicMock()
    existing_key3.tpm_limit = None  # Should be ignored in calculation
    existing_key3.rpm_limit = None  # Should be ignored in calculation

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2, existing_key3]
    )

    # Create team table with limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-456",
        team_alias="test-team",
        tpm_limit=10000,  # Total: 3000 + 2000 + 4000 (new) = 9000 < 10000 ✓
        rpm_limit=1000,  # Total: 200 + 300 + 400 (new) = 900 < 1000 ✓
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request that would still be within bounds
    data = GenerateKeyRequest(
        tpm_limit=4000,
        rpm_limit=400,
    )

    # Should not raise any exception
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_tpm_overallocation():
    """
    Test _check_team_key_limits when new key would cause TPM overallocation.
    Should raise HTTPException with appropriate error message.
    """
    # Create mock existing keys with high TPM usage
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 6000
    existing_key1.rpm_limit = 100
    existing_key1.metadata = {}

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = 3000
    existing_key2.rpm_limit = 200
    existing_key2.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2]
    )

    # Create team table with limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-789",
        team_alias="test-team",
        tpm_limit=10000,  # Allocated: 6000 + 3000 = 9000, New: 2000, Total: 11000 > 10000 ✗
        rpm_limit=1000,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request that would exceed TPM limits
    data = GenerateKeyRequest(
        tpm_limit=2000,
        rpm_limit=100,
        tpm_limit_type="guaranteed_throughput",
    )

    # Should raise HTTPException for TPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        await _check_team_key_limits(
            team_table=team_table,
            data=data,
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated TPM limit=9000 + Key TPM limit=2000 is greater than team TPM limit=10000"
        in str(exc_info.value.detail)
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_rpm_overallocation():
    """
    Test _check_team_key_limits when new key would cause RPM overallocation.
    Should raise HTTPException with appropriate error message.
    """
    # Create mock existing keys with high RPM usage
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 1000
    existing_key1.rpm_limit = 600
    existing_key1.metadata = {}

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = 2000
    existing_key2.rpm_limit = 300
    existing_key2.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2]
    )

    # Create team table with limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-101",
        team_alias="test-team",
        tpm_limit=10000,
        rpm_limit=1000,  # Allocated: 600 + 300 = 900, New: 200, Total: 1100 > 1000 ✗
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request that would exceed RPM limits
    data = GenerateKeyRequest(
        tpm_limit=1000,
        rpm_limit=200,
        rpm_limit_type="guaranteed_throughput",
    )

    # Should raise HTTPException for RPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        await _check_team_key_limits(
            team_table=team_table,
            data=data,
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated RPM limit=900 + Key RPM limit=200 is greater than team RPM limit=1000"
        in str(exc_info.value.detail)
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_on_update_excludes_self():
    """
    Test that _check_team_key_limits excludes the key being updated from the
    allocated totals. Without this, the key's current limits would be
    double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
    """
    from litellm.proxy._types import hash_token as _ht

    # The key being updated is returned by find_many with its current limits.
    # In the DB, token is stored as a SHA-256 hash of the raw key.
    self_key = MagicMock()
    self_key.token = _ht("sk-self-team-key")
    self_key.tpm_limit = 6000
    self_key.rpm_limit = 600
    self_key.metadata = {}

    # Another key in the team
    other_key = MagicMock()
    other_key.token = _ht("sk-other-team-key")
    other_key.tpm_limit = 3000
    other_key.rpm_limit = 300
    other_key.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[self_key, other_key]
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-self",
        team_alias="test-team",
        tpm_limit=10000,
        rpm_limit=1000,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Updating the key to 7000 TPM. Other key uses 3000, so total = 10000 <= 10000.
    # Without the fix, this would be 6000 (self) + 3000 (other) + 7000 = 16000 > 10000.
    data = UpdateKeyRequest(
        key="sk-self-team-key",
        tpm_limit=7000,
        rpm_limit=700,
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
    )

    # Should not raise - the key's own limits should be excluded from the sum
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_no_team_limits():
    """
    Test _check_team_key_limits when team has no TPM/RPM limits set.
    Should allow any key limits since there are no team constraints.
    """
    # Create mock existing keys
    existing_key = MagicMock()
    existing_key.tpm_limit = 5000
    existing_key.rpm_limit = 500

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key]
    )

    # Create team table with no limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-202",
        team_alias="test-team",
        tpm_limit=None,  # No team limit
        rpm_limit=None,  # No team limit
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request with any limits
    data = GenerateKeyRequest(
        tpm_limit=10000,  # High limit should be allowed
        rpm_limit=2000,  # High limit should be allowed
    )

    # Should not raise any exception
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_no_key_limits():
    """
    Test _check_team_key_limits when new key has no TPM/RPM limits.
    Should not raise any exceptions since no limits are being allocated.
    """
    # Create mock existing keys
    existing_key = MagicMock()
    existing_key.tpm_limit = 8000
    existing_key.rpm_limit = 800

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key]
    )

    # Create team table with limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-303",
        team_alias="test-team",
        tpm_limit=10000,
        rpm_limit=1000,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request with no limits
    data = GenerateKeyRequest(
        tpm_limit=None,  # No limit being set
        rpm_limit=None,  # No limit being set
    )

    # Should not raise any exception
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_mixed_scenarios():
    """
    Test _check_team_key_limits with mixed scenarios:
    - Some existing keys have limits, others don't
    - New key has only one type of limit
    - Team has only one type of limit
    """
    # Create mock existing keys with mixed limits
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 3000
    existing_key1.rpm_limit = None  # No RPM limit

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = None  # No TPM limit
    existing_key2.rpm_limit = 400

    existing_key3 = MagicMock()
    existing_key3.tpm_limit = 2000
    existing_key3.rpm_limit = 300

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2, existing_key3]
    )

    # Create team table with only TPM limit
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-404",
        team_alias="test-team",
        tpm_limit=10000,  # Allocated: 3000 + 0 + 2000 = 5000, New: 4000, Total: 9000 < 10000 ✓
        rpm_limit=None,  # No team RPM limit
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request with only TPM limit
    data = GenerateKeyRequest(
        tpm_limit=4000,
        rpm_limit=None,  # No RPM limit being set
    )

    # Should not raise any exception
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


@pytest.mark.asyncio
async def test_check_team_key_limits_exact_boundary():
    """
    Test _check_team_key_limits when allocation exactly matches team limits.
    Should allow the allocation (boundary case).
    """
    # Create mock existing keys
    existing_key = MagicMock()
    existing_key.tpm_limit = 7000
    existing_key.rpm_limit = 700

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key]
    )

    # Create team table with limits
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-505",
        team_alias="test-team",
        tpm_limit=10000,  # Allocated: 7000, New: 3000, Total: 10000 = 10000 ✓
        rpm_limit=1000,  # Allocated: 700, New: 300, Total: 1000 = 1000 ✓
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
    )

    # Create request that exactly matches remaining capacity
    data = GenerateKeyRequest(
        tpm_limit=3000,
        rpm_limit=300,
    )

    # Should not raise any exception (exact boundary should be allowed)
    await _check_team_key_limits(
        team_table=team_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


def test_check_team_key_model_specific_limits_no_limits():
    """
    Test check_team_key_model_specific_limits when no model-specific limits are set.
    Should return without raising any exceptions.
    """
    # Create existing key with no model-specific limits
    existing_key = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user",
        team_id="test-team-123",
        metadata={},
    )

    keys = [existing_key]

    # Create team table
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=10000,
        rpm_limit=1000,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
        metadata={},
    )

    # Create request with no model-specific limits
    data = GenerateKeyRequest(
        model_rpm_limit=None,
        model_tpm_limit=None,
    )

    # Should not raise any exception
    check_team_key_model_specific_limits(
        keys=keys,
        team_table=team_table,
        data=data,
    )


def test_check_team_key_model_specific_limits_rpm_overallocation():
    """
    Test check_team_key_model_specific_limits when model-specific RPM would cause overallocation.
    Should raise HTTPException with appropriate error message.
    """
    # Create existing keys with model-specific RPM limits
    existing_key1 = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user-1",
        team_id="test-team-456",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 500,
                "gpt-3.5-turbo": 300,
            }
        },
    )

    existing_key2 = LiteLLM_VerificationToken(
        token="test-token-2",
        user_id="test-user-2",
        team_id="test-team-456",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 300,
            }
        },
    )

    keys = [existing_key1, existing_key2]

    # Create team table with RPM limit
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-456",
        team_alias="test-team",
        tpm_limit=10000,
        rpm_limit=1000,  # Total team RPM limit
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
        metadata={},
    )

    # Create request that would exceed model-specific RPM limits
    # Existing gpt-4: 500 + 300 = 800, New: 300, Total: 1100 > 1000 (team limit)
    data = GenerateKeyRequest(
        model_rpm_limit={
            "gpt-4": 300,  # This would cause overallocation
        },
        model_tpm_limit=None,
    )

    # Should raise HTTPException for model-specific RPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        check_team_key_model_specific_limits(
            keys=keys,
            team_table=team_table,
            data=data,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated RPM limit=800 + Key RPM limit=300 is greater than team RPM limit=1000"
        in str(exc_info.value.detail)
    )


def test_check_team_key_model_specific_limits_team_model_rpm_overallocation():
    """
    Test check_team_key_model_specific_limits when team has model-specific RPM limits
    in metadata and key allocation would exceed those limits.

    This tests the scenario where team_table.metadata["model_rpm_limit"] is set
    with per-model limits, not just a global team RPM limit.
    """
    # Create existing keys with model-specific RPM limits
    existing_key1 = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user-1",
        team_id="test-team-789",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 300,
                "gpt-3.5-turbo": 200,
            }
        },
    )

    existing_key2 = LiteLLM_VerificationToken(
        token="test-token-2",
        user_id="test-user-2",
        team_id="test-team-789",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 250,
            }
        },
    )

    keys = [existing_key1, existing_key2]

    # Create team table with model-specific RPM limits in metadata
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-789",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
        metadata={
            "model_rpm_limit": {
                "gpt-4": 700,  # Team-level model-specific limit for gpt-4
                "gpt-3.5-turbo": 500,
            }
        },
    )

    # Create request that would exceed team's model-specific RPM limits
    # Existing gpt-4: 300 + 250 = 550, New: 200, Total: 750 > 700 (team model-specific limit)
    data = GenerateKeyRequest(
        model_rpm_limit={
            "gpt-4": 200,  # This would cause overallocation against team model-specific limit
        },
        model_tpm_limit=None,
    )

    # Should raise HTTPException for team model-specific RPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        check_team_key_model_specific_limits(
            keys=keys,
            team_table=team_table,
            data=data,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated RPM limit=550 + Key RPM limit=200 is greater than team RPM limit=700"
        in str(exc_info.value.detail)
    )


def test_check_team_key_model_specific_limits_team_model_tpm_overallocation():
    """
    Test check_team_key_model_specific_limits when team has model-specific TPM limits
    in metadata and key allocation would exceed those limits.

    This tests the scenario where team_table.metadata["model_tpm_limit"] is set
    with per-model limits, not just a global team TPM limit.
    """
    # Create existing keys with model-specific TPM limits
    existing_key1 = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user-1",
        team_id="test-team-101",
        metadata={
            "model_tpm_limit": {
                "gpt-4": 5000,
                "claude-3": 3000,
            }
        },
    )

    existing_key2 = LiteLLM_VerificationToken(
        token="test-token-2",
        user_id="test-user-2",
        team_id="test-team-101",
        metadata={
            "model_tpm_limit": {
                "gpt-4": 3500,
            }
        },
    )

    keys = [existing_key1, existing_key2]

    # Create team table with model-specific TPM limits in metadata
    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-101",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=100.0,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[],
        metadata={
            "model_tpm_limit": {
                "gpt-4": 10000,  # Team-level model-specific limit for gpt-4
                "claude-3": 8000,
            }
        },
    )

    # Create request that would exceed team's model-specific TPM limits
    # Existing gpt-4: 5000 + 3500 = 8500, New: 2000, Total: 10500 > 10000 (team model-specific limit)
    data = GenerateKeyRequest(
        model_rpm_limit=None,
        model_tpm_limit={
            "gpt-4": 2000,  # This would cause overallocation against team model-specific limit
        },
    )

    # Should raise HTTPException for team model-specific TPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        check_team_key_model_specific_limits(
            keys=keys,
            team_table=team_table,
            data=data,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated TPM limit=8500 + Key TPM limit=2000 is greater than team TPM limit=10000"
        in str(exc_info.value.detail)
    )


@pytest.mark.asyncio
async def test_generate_key_with_object_permission():
    """
    Test that /key/generate correctly handles object_permission by:
    1. Creating a record in litellm_objectpermissiontable
    2. Passing the returned object_permission_id into the key insert payload
    3. NOT passing the object_permission dict to the key table
    """
    from unittest.mock import patch

    from litellm.proxy._types import (
        GenerateKeyRequest,
        LiteLLM_ObjectPermissionBase,
        LitellmUserRoles,
    )
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _common_key_generation_helper,
    )

    # Mock prisma client
    mock_prisma_client = MagicMock()
    mock_prisma_client.jsonify_object = lambda x: x

    # Mock object permission creation
    mock_object_perm_create = AsyncMock(
        return_value=MagicMock(object_permission_id="objperm_key_456")
    )
    mock_prisma_client.db.litellm_objectpermissiontable = MagicMock()
    mock_prisma_client.db.litellm_objectpermissiontable.create = mock_object_perm_create

    # Mock key insertion
    mock_key_insert = AsyncMock(
        return_value=MagicMock(
            token="hashed_token_123",
            litellm_budget_table=None,
            created_at="2024-01-01T00:00:00Z",
            updated_at="2024-01-01T00:00:00Z",
        )
    )
    mock_prisma_client.insert_data = mock_key_insert

    # Create request with object_permission
    key_request = GenerateKeyRequest(
        models=["gpt-4"],
        object_permission=LiteLLM_ObjectPermissionBase(
            vector_stores=["vector_store_1"],
            mcp_servers=["mcp_server_1"],
        ),
    )

    mock_admin_auth = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        user_id="admin_user",
    )

    # Patch the prisma_client and other dependencies
    with (
        patch(
            "litellm.proxy.proxy_server.prisma_client",
            mock_prisma_client,
        ),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch(
            "litellm.proxy.proxy_server.premium_user",
            False,
        ),
        patch(
            "litellm.proxy.proxy_server.litellm_proxy_admin_name",
            "admin",
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
            new=AsyncMock(
                side_effect=lambda object_permission=None, **kwargs: object_permission
            ),
        ),
    ):
        # Execute
        result = await _common_key_generation_helper(
            data=key_request,
            user_api_key_dict=mock_admin_auth,
            litellm_changed_by=None,
            team_table=None,
        )

    # Verify object permission creation was called
    mock_object_perm_create.assert_awaited_once()

    # Verify key insertion was called
    assert mock_key_insert.call_count == 1
    key_insert_kwargs = mock_key_insert.call_args.kwargs
    key_data = key_insert_kwargs["data"]

    # Verify object_permission_id is in the key data
    assert key_data.get("object_permission_id") == "objperm_key_456"

    # Verify object_permission dict is NOT in the key data
    assert "object_permission" not in key_data


@pytest.mark.asyncio
async def test_generate_key_team_member_inherits_org_skips_membership_check():
    """Regression: a team member creating a key for an org-scoped team must not
    be blocked by the org-membership check.

    When ``organization_id`` is inherited from the key's team (via
    ``apply_enterprise_key_management_params`` -> ``add_team_organization_id``),
    the caller already passed team-level authorization. Requiring an explicit
    ``LiteLLM_OrganizationMembership`` row on top of that broke the normal admin
    workflow (admins only add users to teams). This asserts the org-membership
    check is skipped when the org id came from the caller's team.
    """
    from unittest.mock import AsyncMock, MagicMock, patch

    from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _common_key_generation_helper,
    )

    org_id = "org-from-team"

    # Team belongs to an org; caller is a team member but NOT an explicit member
    # of that organization (the regression scenario).
    mock_team_table = MagicMock()
    mock_team_table.organization_id = org_id
    mock_team_table.metadata = None

    mock_validate_org = AsyncMock()
    mock_generate_key = AsyncMock(
        return_value={
            "key": "sk-test-key",
            "expires": None,
            "user_id": "alice",
            "team_id": "team-1",
        }
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", True),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org",
            mock_validate_org,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object",
            new_callable=AsyncMock,
            return_value=MagicMock(litellm_budget_table=None),
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            mock_generate_key,
        ),
    ):
        result = await _common_key_generation_helper(
            data=GenerateKeyRequest(
                user_id="alice",
                team_id="team-1",
                organization_id=org_id,
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_id="alice",
                user_role=LitellmUserRoles.INTERNAL_USER.value,
            ),
            litellm_changed_by=None,
            team_table=mock_team_table,
        )

    # Key creation proceeded for the team member ...
    mock_generate_key.assert_awaited_once()
    assert result is not None
    # ... and the org-membership check was bypassed because organization_id was
    # inherited from the caller's team.
    mock_validate_org.assert_not_awaited()


@pytest.mark.asyncio
async def test_generate_key_foreign_org_without_team_still_enforces_membership():
    """VERIA-55: a caller assigning a key to an organization that was NOT
    inherited from a team must still pass the org-membership check.

    This guards the IDOR fix: ``team_table is None`` (or an org id that does not
    match the team) means the org id did not come from team context, so the
    explicit membership validation must run.
    """
    from unittest.mock import AsyncMock, MagicMock, patch

    from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _common_key_generation_helper,
    )

    foreign_org_id = "someone-elses-org"

    mock_validate_org = AsyncMock()
    mock_generate_key = AsyncMock(
        return_value={
            "key": "sk-test-key",
            "expires": None,
            "user_id": "alice",
            "team_id": None,
        }
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", True),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org",
            mock_validate_org,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object",
            new_callable=AsyncMock,
            return_value=MagicMock(litellm_budget_table=None),
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            mock_generate_key,
        ),
    ):
        await _common_key_generation_helper(
            data=GenerateKeyRequest(
                user_id="alice",
                organization_id=foreign_org_id,
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_id="alice",
                user_role=LitellmUserRoles.INTERNAL_USER.value,
            ),
            litellm_changed_by=None,
            team_table=None,
        )

    # No team context -> the org-membership check must still run.
    mock_validate_org.assert_awaited_once()
    assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id


@pytest.mark.asyncio
async def test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership():
    """VERIA-55: when a team is present but its organization_id differs from the
    organization_id on the key request, the org-membership check must still run."""
    from unittest.mock import AsyncMock, MagicMock, patch

    from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _common_key_generation_helper,
    )

    team_org_id = "other-org"
    foreign_org_id = "someone-elses-org"

    mock_team_table = MagicMock()
    mock_team_table.organization_id = team_org_id
    mock_team_table.metadata = None

    mock_validate_org = AsyncMock()
    mock_generate_key = AsyncMock(
        return_value={
            "key": "sk-test-key",
            "expires": None,
            "user_id": "alice",
            "team_id": "team-1",
        }
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", True),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm_enterprise.proxy.management_endpoints.key_management_endpoints.apply_enterprise_key_management_params",
            side_effect=lambda data, team_table: data,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org",
            mock_validate_org,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object",
            new_callable=AsyncMock,
            return_value=MagicMock(litellm_budget_table=None),
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            mock_generate_key,
        ),
    ):
        await _common_key_generation_helper(
            data=GenerateKeyRequest(
                user_id="alice",
                team_id="team-1",
                organization_id=foreign_org_id,
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_id="alice",
                user_role=LitellmUserRoles.INTERNAL_USER.value,
            ),
            litellm_changed_by=None,
            team_table=mock_team_table,
        )

    mock_validate_org.assert_awaited_once()
    assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id


# ============================================
# Organization Key Limit Tests
# ============================================


@pytest.mark.asyncio
async def test_check_org_key_limits_no_existing_keys():
    """
    Test _check_org_key_limits when organization has no existing keys.
    Should allow any TPM/RPM limits within organization bounds.
    """
    # Mock prisma client
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )

    # Create organization table with limits in budget table
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-123",
        organization_alias="test-org",
        budget_id="budget-123",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-123",
            tpm_limit=20000,
            rpm_limit=2000,
        ),
    )

    # Create request with limits within organization bounds
    data = GenerateKeyRequest(
        tpm_limit=10000,
        rpm_limit=1000,
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
    )

    # Should not raise any exception
    await _check_org_key_limits(
        org_table=org_table,
        data=data,
        prisma_client=mock_prisma_client,
    )

    # Verify database was queried
    mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
        where={"organization_id": "test-org-123"}
    )


@pytest.mark.asyncio
async def test_check_org_key_limits_with_existing_keys_within_bounds():
    """
    Test _check_org_key_limits when organization has existing keys but total allocation
    is still within organization limits.
    """
    # Create mock existing keys
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 5000
    existing_key1.rpm_limit = 400
    existing_key1.metadata = {}

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = 3000
    existing_key2.rpm_limit = 500
    existing_key2.metadata = {}

    existing_key3 = MagicMock()
    existing_key3.tpm_limit = None  # Should be ignored in calculation
    existing_key3.rpm_limit = None  # Should be ignored in calculation
    existing_key3.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2, existing_key3]
    )

    # Create organization table with limits
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-456",
        organization_alias="test-org",
        budget_id="budget-456",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-456",
            tpm_limit=20000,  # Total: 5000 + 3000 + 8000 (new) = 16000 < 20000 ✓
            rpm_limit=2000,  # Total: 400 + 500 + 800 (new) = 1700 < 2000 ✓
        ),
    )

    # Create request that would still be within bounds
    data = GenerateKeyRequest(
        tpm_limit=8000,
        rpm_limit=800,
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
    )

    # Should not raise any exception
    await _check_org_key_limits(
        org_table=org_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


@pytest.mark.asyncio
async def test_check_org_key_limits_tpm_overallocation():
    """
    Test _check_org_key_limits when new key would cause TPM overallocation.
    Should raise HTTPException with appropriate error message.
    """
    # Create mock existing keys with high TPM usage
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 12000
    existing_key1.rpm_limit = 500
    existing_key1.metadata = {}

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = 6000
    existing_key2.rpm_limit = 400
    existing_key2.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2]
    )

    # Create organization table with limits
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-789",
        organization_alias="test-org",
        budget_id="budget-789",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-789",
            tpm_limit=20000,  # Allocated: 12000 + 6000 = 18000, New: 3000, Total: 21000 > 20000 ✗
            rpm_limit=2000,
        ),
    )

    # Create request that would exceed TPM limits
    data = GenerateKeyRequest(
        tpm_limit=3000,
        rpm_limit=500,
        tpm_limit_type="guaranteed_throughput",
    )

    # Should raise HTTPException for TPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        await _check_org_key_limits(
            org_table=org_table,
            data=data,
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated TPM limit=18000 + Key TPM limit=3000 is greater than organization TPM limit=20000"
        in str(exc_info.value.detail)
    )


@pytest.mark.asyncio
async def test_check_org_key_limits_rpm_overallocation():
    """
    Test _check_org_key_limits when new key would cause RPM overallocation.
    Should raise HTTPException with appropriate error message.
    """
    # Create mock existing keys with high RPM usage
    existing_key1 = MagicMock()
    existing_key1.tpm_limit = 5000
    existing_key1.rpm_limit = 1200
    existing_key1.metadata = {}

    existing_key2 = MagicMock()
    existing_key2.tpm_limit = 3000
    existing_key2.rpm_limit = 600
    existing_key2.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key1, existing_key2]
    )

    # Create organization table with limits
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-101",
        organization_alias="test-org",
        budget_id="budget-101",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-101",
            tpm_limit=20000,
            rpm_limit=2000,  # Allocated: 1200 + 600 = 1800, New: 300, Total: 2100 > 2000 ✗
        ),
    )

    # Create request that would exceed RPM limits
    data = GenerateKeyRequest(
        tpm_limit=5000,
        rpm_limit=300,
        rpm_limit_type="guaranteed_throughput",
    )

    # Should raise HTTPException for RPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        await _check_org_key_limits(
            org_table=org_table,
            data=data,
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated RPM limit=1800 + Key RPM limit=300 is greater than organization RPM limit=2000"
        in str(exc_info.value.detail)
    )


@pytest.mark.asyncio
async def test_check_org_key_limits_no_org_limits():
    """
    Test _check_org_key_limits when organization has no TPM/RPM limits set.
    Should allow any key limits since there are no organization constraints.
    """
    # Create mock existing keys
    existing_key = MagicMock()
    existing_key.tpm_limit = 10000
    existing_key.rpm_limit = 1000
    existing_key.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key]
    )

    # Create organization table with no limits
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-202",
        organization_alias="test-org",
        budget_id="budget-202",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-202",
            tpm_limit=None,  # No organization limit
            rpm_limit=None,  # No organization limit
        ),
    )

    # Create request with any limits
    data = GenerateKeyRequest(
        tpm_limit=50000,  # High limit should be allowed
        rpm_limit=5000,  # High limit should be allowed
        tpm_limit_type="guaranteed_throughput",
    )

    # Should not raise any exception
    await _check_org_key_limits(
        org_table=org_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


def test_check_org_key_model_specific_limits_no_limits():
    """
    Test check_org_key_model_specific_limits when no model-specific limits are set.
    Should return without raising any exceptions.
    """
    # Create existing key with no model-specific limits
    existing_key = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user",
        organization_id="test-org-123",
        metadata={},
    )

    keys = [existing_key]

    # Create organization table
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-123",
        organization_alias="test-org",
        budget_id="budget-123",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        metadata={},
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-123",
            tpm_limit=20000,
            rpm_limit=2000,
        ),
    )

    # Create request with no model-specific limits
    data = GenerateKeyRequest(
        model_rpm_limit=None,
        model_tpm_limit=None,
    )

    # Should not raise any exception
    check_org_key_model_specific_limits(
        keys=keys,
        org_table=org_table,
        data=data,
    )


def test_check_org_key_model_specific_limits_rpm_overallocation():
    """
    Test check_org_key_model_specific_limits when model-specific RPM would cause overallocation.
    Should raise HTTPException with appropriate error message.
    """
    # Create existing keys with model-specific RPM limits
    existing_key1 = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user-1",
        organization_id="test-org-456",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 800,
                "gpt-3.5-turbo": 500,
            }
        },
    )

    existing_key2 = LiteLLM_VerificationToken(
        token="test-token-2",
        user_id="test-user-2",
        organization_id="test-org-456",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 600,
            }
        },
    )

    keys = [existing_key1, existing_key2]

    # Create organization table with RPM limit
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-456",
        organization_alias="test-org",
        budget_id="budget-456",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        metadata={},
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-456",
            tpm_limit=20000,
            rpm_limit=2000,  # Total organization RPM limit
        ),
    )

    # Create request that would exceed model-specific RPM limits
    # Existing gpt-4: 800 + 600 = 1400, New: 700, Total: 2100 > 2000 (org limit)
    data = GenerateKeyRequest(
        model_rpm_limit={
            "gpt-4": 700,  # This would cause overallocation
        },
        model_tpm_limit=None,
    )

    # Should raise HTTPException for model-specific RPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        check_org_key_model_specific_limits(
            keys=keys,
            org_table=org_table,
            data=data,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated RPM limit=1400 + Key RPM limit=700 is greater than organization RPM limit=2000"
        in str(exc_info.value.detail)
    )


def test_check_org_key_model_specific_limits_org_model_rpm_overallocation():
    """
    Test check_org_key_model_specific_limits when organization has model-specific RPM limits
    in metadata and key allocation would exceed those limits.

    This tests the scenario where org_table.metadata["model_rpm_limit"] is set
    with per-model limits, not just a global organization RPM limit.
    """
    # Create existing keys with model-specific RPM limits
    existing_key1 = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user-1",
        organization_id="test-org-789",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 600,
                "gpt-3.5-turbo": 400,
            }
        },
    )

    existing_key2 = LiteLLM_VerificationToken(
        token="test-token-2",
        user_id="test-user-2",
        organization_id="test-org-789",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 500,
            }
        },
    )

    keys = [existing_key1, existing_key2]

    # Create organization table with model-specific RPM limits in metadata
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-789",
        organization_alias="test-org",
        budget_id="budget-789",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        metadata={
            "model_rpm_limit": {
                "gpt-4": 1400,  # Organization-level model-specific limit for gpt-4
                "gpt-3.5-turbo": 1000,
            }
        },
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-789",
            tpm_limit=None,
            rpm_limit=None,
        ),
    )

    # Create request that would exceed organization's model-specific RPM limits
    # Existing gpt-4: 600 + 500 = 1100, New: 400, Total: 1500 > 1400 (org model-specific limit)
    data = GenerateKeyRequest(
        model_rpm_limit={
            "gpt-4": 400,  # This would cause overallocation against org model-specific limit
        },
        model_tpm_limit=None,
    )

    # Should raise HTTPException for organization model-specific RPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        check_org_key_model_specific_limits(
            keys=keys,
            org_table=org_table,
            data=data,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated RPM limit=1100 + Key RPM limit=400 is greater than organization RPM limit=1400"
        in str(exc_info.value.detail)
    )


def test_check_org_key_model_specific_limits_org_model_tpm_overallocation():
    """
    Test check_org_key_model_specific_limits when organization has model-specific TPM limits
    in metadata and key allocation would exceed those limits.

    This tests the scenario where org_table.metadata["model_tpm_limit"] is set
    with per-model limits, not just a global organization TPM limit.
    """
    # Create existing keys with model-specific TPM limits
    existing_key1 = LiteLLM_VerificationToken(
        token="test-token-1",
        user_id="test-user-1",
        organization_id="test-org-101",
        metadata={
            "model_tpm_limit": {
                "gpt-4": 10000,
                "claude-3": 6000,
            }
        },
    )

    existing_key2 = LiteLLM_VerificationToken(
        token="test-token-2",
        user_id="test-user-2",
        organization_id="test-org-101",
        metadata={
            "model_tpm_limit": {
                "gpt-4": 7000,
            }
        },
    )

    keys = [existing_key1, existing_key2]

    # Create organization table with model-specific TPM limits in metadata
    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-101",
        organization_alias="test-org",
        budget_id="budget-101",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        metadata={
            "model_tpm_limit": {
                "gpt-4": 20000,  # Organization-level model-specific limit for gpt-4
                "claude-3": 15000,
            }
        },
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-101",
            tpm_limit=None,
            rpm_limit=None,
        ),
    )

    # Create request that would exceed organization's model-specific TPM limits
    # Existing gpt-4: 10000 + 7000 = 17000, New: 4000, Total: 21000 > 20000 (org model-specific limit)
    data = GenerateKeyRequest(
        model_rpm_limit=None,
        model_tpm_limit={
            "gpt-4": 4000,  # This would cause overallocation against org model-specific limit
        },
    )

    # Should raise HTTPException for organization model-specific TPM overallocation
    with pytest.raises(HTTPException) as exc_info:
        check_org_key_model_specific_limits(
            keys=keys,
            org_table=org_table,
            data=data,
        )

    assert exc_info.value.status_code == 400
    assert (
        "Allocated TPM limit=17000 + Key TPM limit=4000 is greater than organization TPM limit=20000"
        in str(exc_info.value.detail)
    )


def test_transform_verification_tokens_to_deleted_records():
    from datetime import datetime, timezone

    user_api_key_dict = UserAPIKeyAuth(
        user_id="user-123",
        api_key="sk-test",
        user_role=LitellmUserRoles.PROXY_ADMIN.value,
    )

    key1 = LiteLLM_VerificationToken(
        token="hashed-token-1",
        user_id="user-123",
        team_id="team-456",
        key_alias="test-key-1",
        spend=100.0,
        max_budget=1000.0,
        models=["gpt-4"],
        aliases={},
        config={},
        permissions={},
        metadata={"test": "value"},
        model_max_budget={},
        model_spend={},
        soft_budget_cooldown=False,
        allowed_routes=[],
    )

    key2 = LiteLLM_VerificationToken(
        token="hashed-token-2",
        user_id="user-789",
        team_id=None,
        key_alias="test-key-2",
        spend=50.0,
        max_budget=500.0,
        models=["gpt-3.5-turbo"],
        aliases={"alias": "model"},
        config={"config": "value"},
        permissions={"permission": True},
        metadata={},
        model_max_budget={"gpt-4": {"budget_limit": 100.0}},
        model_spend={},
        soft_budget_cooldown=False,
        allowed_routes=[],
    )

    records = _transform_verification_tokens_to_deleted_records(
        keys=[key1, key2],
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by="admin-user",
    )

    assert len(records) == 2
    assert all("deleted_at" in record for record in records)
    assert all("deleted_by" in record for record in records)
    assert all("deleted_by_api_key" in record for record in records)
    assert all("litellm_changed_by" in record for record in records)
    assert all(record["deleted_by"] == "user-123" for record in records)
    assert all(
        record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records
    )
    assert all(record["litellm_changed_by"] == "admin-user" for record in records)

    record1 = records[0]
    assert record1["token"] == "hashed-token-1"
    assert record1["user_id"] == "user-123"
    assert record1["team_id"] == "team-456"
    assert isinstance(record1["aliases"], str)
    assert isinstance(record1["config"], str)
    assert isinstance(record1["permissions"], str)
    assert isinstance(record1["metadata"], str)
    assert "litellm_budget_table" not in record1
    assert "litellm_organization_table" not in record1
    assert "object_permission" not in record1
    assert "id" not in record1

    record2 = records[1]
    assert record2["token"] == "hashed-token-2"
    assert isinstance(record2["model_max_budget"], str)


def test_transform_verification_tokens_to_deleted_records_empty_list():
    user_api_key_dict = UserAPIKeyAuth(
        user_id="user-123",
        api_key="sk-test",
        user_role=LitellmUserRoles.PROXY_ADMIN.value,
    )

    records = _transform_verification_tokens_to_deleted_records(
        keys=[],
        user_api_key_dict=user_api_key_dict,
    )

    assert records == []


@pytest.mark.asyncio
async def test_save_deleted_verification_token_records():
    mock_prisma_client = AsyncMock()
    mock_create_many = AsyncMock()
    mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
        mock_create_many
    )

    records = [
        {
            "token": "hashed-token-1",
            "user_id": "user-123",
            "deleted_at": "2024-01-01T00:00:00Z",
            "deleted_by": "admin",
        },
        {
            "token": "hashed-token-2",
            "user_id": "user-456",
            "deleted_at": "2024-01-01T00:00:00Z",
            "deleted_by": "admin",
        },
    ]

    await _save_deleted_verification_token_records(
        records=records, prisma_client=mock_prisma_client
    )

    mock_create_many.assert_called_once_with(data=records)


@pytest.mark.asyncio
async def test_save_deleted_verification_token_records_empty_list():
    mock_prisma_client = AsyncMock()
    mock_create_many = AsyncMock()
    mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
        mock_create_many
    )

    await _save_deleted_verification_token_records(
        records=[], prisma_client=mock_prisma_client
    )

    mock_create_many.assert_not_called()


@pytest.mark.asyncio
async def test_persist_deleted_verification_tokens():
    mock_prisma_client = AsyncMock()
    mock_create_many = AsyncMock()
    mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
        mock_create_many
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_id="user-123",
        api_key="sk-test",
        user_role=LitellmUserRoles.PROXY_ADMIN.value,
    )

    key = LiteLLM_VerificationToken(
        token="hashed-token-1",
        user_id="user-123",
        team_id="team-456",
        key_alias="test-key",
        spend=100.0,
        max_budget=1000.0,
        models=["gpt-4"],
        aliases={},
        config={},
        permissions={},
        metadata={},
        model_max_budget={},
        model_spend={},
        soft_budget_cooldown=False,
        allowed_routes=[],
    )

    await _persist_deleted_verification_tokens(
        keys=[key],
        prisma_client=mock_prisma_client,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by="admin-user",
    )

    mock_create_many.assert_called_once()
    call_args = mock_create_many.call_args
    assert "data" in call_args.kwargs
    records = call_args.kwargs["data"]
    assert len(records) == 1
    assert records[0]["token"] == "hashed-token-1"
    assert records[0]["deleted_by"] == "user-123"
    assert records[0]["litellm_changed_by"] == "admin-user"


@pytest.mark.asyncio
async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    user_api_key_dict = UserAPIKeyAuth(
        user_id="admin-user",
        api_key="sk-admin",
        user_role=LitellmUserRoles.PROXY_ADMIN.value,
    )

    key1 = LiteLLM_VerificationToken(
        token="hashed-token-1",
        user_id="user-123",
        team_id="team-456",
        key_alias="test-key-1",
        spend=100.0,
        max_budget=1000.0,
        models=["gpt-4"],
        aliases={},
        config={},
        permissions={},
        metadata={},
        model_max_budget={},
        model_spend={},
        soft_budget_cooldown=False,
        allowed_routes=[],
    )

    key2 = LiteLLM_VerificationToken(
        token="hashed-token-2",
        user_id="user-789",
        team_id=None,
        key_alias="test-key-2",
        spend=50.0,
        max_budget=500.0,
        models=["gpt-3.5-turbo"],
        aliases={},
        config={},
        permissions={},
        metadata={},
        model_max_budget={},
        model_spend={},
        soft_budget_cooldown=False,
        allowed_routes=[],
    )

    mock_find_many = AsyncMock(return_value=[key1, key2])
    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many

    # delete_data returns {"deleted_keys": ...} from utils.py line 3049
    # The function at line 2410 assigns it to deleted_tokens
    # Then at line 2444 returns {"deleted_keys": deleted_tokens}
    # So if delete_data returns {"deleted_keys": list}, then result would be nested
    # But looking at the error, it seems like delete_data might return just the list
    # Or the code extracts it. Let's return the list directly since that's what the test expects
    mock_delete_data = AsyncMock(return_value=["hashed-token-1", "hashed-token-2"])
    mock_prisma_client.delete_data = mock_delete_data

    # Mock cache delete_cache method
    mock_user_api_key_cache.delete_cache = MagicMock()

    mock_create_many = AsyncMock()
    mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
        mock_create_many
    )

    def mock_hash_token(token):
        return token if not token.startswith("sk-") else f"hashed-{token}"

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
        mock_hash_token,
    )
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.hash_token",
        mock_hash_token,
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.prisma_client",
        mock_prisma_client,
    )

    result, deleted_keys = await delete_verification_tokens(
        tokens=["sk-token-1", "sk-token-2"],
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by="admin-user",
    )

    mock_create_many.assert_called_once()
    call_args = mock_create_many.call_args
    assert "data" in call_args.kwargs
    records = call_args.kwargs["data"]
    assert len(records) == 2
    assert all(record["deleted_by"] == "admin-user" for record in records)
    assert all(record["litellm_changed_by"] == "admin-user" for record in records)
    # delete_data returns the list directly, which gets wrapped in {"deleted_keys": ...}
    assert isinstance(result["deleted_keys"], list)
    assert set(result["deleted_keys"]) == {"hashed-token-1", "hashed-token-2"}
    assert len(deleted_keys) == 2


@pytest.mark.asyncio
async def test_delete_key_fn_persists_deleted_keys(monkeypatch):
    from litellm.proxy._types import KeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        delete_key_fn,
        delete_verification_tokens,
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    user_api_key_dict = UserAPIKeyAuth(
        user_id="admin-user",
        api_key="sk-admin",
        user_role=LitellmUserRoles.PROXY_ADMIN.value,
    )

    key1 = LiteLLM_VerificationToken(
        token="hashed-token-1",
        user_id="user-123",
        team_id="team-456",
        key_alias="test-key-1",
        spend=100.0,
        max_budget=1000.0,
        models=["gpt-4"],
        aliases={},
        config={},
        permissions={},
        metadata={},
        model_max_budget={},
        model_spend={},
        soft_budget_cooldown=False,
        allowed_routes=[],
    )

    async def mock_delete_verification_tokens(*args, **kwargs):
        return ({"deleted_keys": ["sk-token-1"]}, [key1])

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.delete_verification_tokens",
        mock_delete_verification_tokens,
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.prisma_client",
        mock_prisma_client,
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache",
        mock_user_api_key_cache,
    )
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_deleted_hook",
        AsyncMock(),
    )

    data = KeyRequest(keys=["sk-token-1"])

    result = await delete_key_fn(
        data=data,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by="admin-user",
    )

    assert result["deleted_keys"] == ["sk-token-1"]


@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch):
    """Test that team admin can delete team keys from their own team."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="other-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="team-admin-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="team-admin-user", role="admin"),
            Member(user_id="other-user", role="user"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_delete_verification_token_team_admin_different_team(monkeypatch):
    """Test that team admin cannot delete team keys from a different team."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="other-user",
        team_id="test-team-456",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="team-admin-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-456",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="different-admin", role="admin"),
            Member(user_id="other-user", role="user"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_delete_verification_token_key_owner_team_key(monkeypatch):
    """Test that key owner can delete their own team key."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="key-owner-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="key-owner-user", role="user"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_delete_verification_token_key_owner_personal_key(monkeypatch):
    """Test that key owner can delete their own personal key."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="key-owner-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_delete_verification_token_other_user_team_key(monkeypatch):
    """Test that other user cannot delete team keys they don't own and aren't admin for."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="other-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="key-owner-user", role="user"),
            Member(user_id="other-user", role="user"),
            Member(user_id="team-admin-user", role="admin"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_delete_verification_token_other_user_personal_key(monkeypatch):
    """Test that other user cannot delete personal keys they don't own."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="other-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_delete_verification_token_team_key_no_team_found(monkeypatch):
    """Test that deletion fails when team is not found in database."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id="non-existent-team",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="key-owner-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return None

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_delete_verification_token_personal_key_no_user_id(monkeypatch):
    """Test that deletion fails for personal key when key has no user_id."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id=None,
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="some-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_modify_verification_token_proxy_admin_team_key(monkeypatch):
    """Test that proxy admin can modify any team key."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="other-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        user_id="admin-user",
        api_key="sk-admin",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_modify_verification_token_proxy_admin_personal_key(monkeypatch):
    """Test that proxy admin can modify any personal key."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="other-user",
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        user_id="admin-user",
        api_key="sk-admin",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_modify_verification_token_team_admin_own_team(monkeypatch):
    """Test that team admin can modify team keys from their own team."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="other-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="team-admin-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="team-admin-user", role="admin"),
            Member(user_id="other-user", role="user"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_modify_verification_token_team_admin_different_team(monkeypatch):
    """Test that team admin cannot modify team keys from a different team."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="other-user",
        team_id="test-team-456",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="team-admin-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-456",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="different-admin", role="admin"),
            Member(user_id="other-user", role="user"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_modify_verification_token_key_owner_team_key(monkeypatch):
    """Test that key owner can modify their own team key."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="key-owner-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="key-owner-user", role="user"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_modify_verification_token_key_owner_personal_key(monkeypatch):
    """Test that key owner can modify their own personal key."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="key-owner-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is True


@pytest.mark.asyncio
async def test_can_modify_verification_token_other_user_team_key(monkeypatch):
    """Test that other user cannot modify team keys they don't own and aren't admin for."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id="test-team-123",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="other-user",
        api_key="sk-user",
    )

    team_table = LiteLLM_TeamTableCachedObj(
        team_id="test-team-123",
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="key-owner-user", role="user"),
            Member(user_id="other-user", role="user"),
            Member(user_id="team-admin-user", role="admin"),
        ],
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_modify_verification_token_other_user_personal_key(monkeypatch):
    """Test that other user cannot modify personal keys they don't own."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="other-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_modify_verification_token_team_key_no_team_found(monkeypatch):
    """Test that modification fails when team is not found in database."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id="key-owner-user",
        team_id="non-existent-team",
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="key-owner-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    async def mock_get_team_object(*args, **kwargs):
        return None

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch):
    """Test that modification fails for personal key when key has no user_id."""
    key_info = LiteLLM_VerificationToken(
        token="test-token",
        user_id=None,
        team_id=None,
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="some-user",
        api_key="sk-user",
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()

    result = await can_modify_verification_token(
        key_info=key_info,
        user_api_key_cache=mock_user_api_key_cache,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    assert result is False


@pytest.mark.asyncio
async def test_list_keys_with_expand_user():
    """
    Test that expand=user parameter correctly includes user information in the response.
    """
    mock_prisma_client = AsyncMock()

    # Create mock keys with user_ids
    key1_dict = {
        "token": "token1",
        "user_id": "user123",
        "key_alias": "key1",
        "models": ["gpt-4"],
    }
    mock_key1 = MagicMock()
    mock_key1.token = "token1"
    mock_key1.user_id = "user123"
    mock_key1.created_by = None
    # Set up model_dump() to raise AttributeError so it falls back to dict()
    mock_key1.model_dump = MagicMock(
        side_effect=AttributeError("model_dump not available")
    )
    mock_key1.dict = MagicMock(return_value=key1_dict)

    key2_dict = {
        "token": "token2",
        "user_id": "user456",
        "key_alias": "key2",
        "models": ["gpt-3.5-turbo"],
    }
    mock_key2 = MagicMock()
    mock_key2.token = "token2"
    mock_key2.user_id = "user456"
    mock_key2.created_by = None
    # Set up model_dump() to raise AttributeError so it falls back to dict()
    mock_key2.model_dump = MagicMock(
        side_effect=AttributeError("model_dump not available")
    )
    mock_key2.dict = MagicMock(return_value=key2_dict)

    mock_find_many_keys = AsyncMock(return_value=[mock_key1, mock_key2])
    mock_count_keys = AsyncMock(return_value=2)

    # Create mock users
    user1_dict = {
        "user_id": "user123",
        "user_email": "user1@example.com",
        "user_alias": "User One",
    }
    mock_user1 = MagicMock()
    # Set user_id as a real attribute (not a MagicMock)
    mock_user1.user_id = "user123"
    mock_user1.user_email = "user1@example.com"
    # Set up both model_dump() and dict() to return the same dict
    mock_user1.model_dump = MagicMock(return_value=user1_dict)
    mock_user1.dict = MagicMock(return_value=user1_dict)

    user2_dict = {
        "user_id": "user456",
        "user_email": "user2@example.com",
        "user_alias": "User Two",
    }
    mock_user2 = MagicMock()
    # Set user_id as a real attribute (not a MagicMock)
    mock_user2.user_id = "user456"
    mock_user2.user_email = "user2@example.com"
    # Set up both model_dump() and dict() to return the same dict
    mock_user2.model_dump = MagicMock(return_value=user2_dict)
    mock_user2.dict = MagicMock(return_value=user2_dict)

    mock_find_many_users = AsyncMock(return_value=[mock_user1, mock_user2])

    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
    mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys
    mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users

    # Patch attach_object_permission_to_dict to just return the dict unchanged
    async def mock_attach_object_permission(d, _):
        return d

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.attach_object_permission_to_dict",
        side_effect=mock_attach_object_permission,
    ):
        args = {
            "prisma_client": mock_prisma_client,
            "page": 1,
            "size": 50,
            "user_id": None,
            "team_id": None,
            "organization_id": None,
            "key_alias": None,
            "key_hash": None,
            "exclude_team_id": None,
            "return_full_object": False,  # This should be overridden by expand=user
            "admin_team_ids": None,
            "include_created_by_keys": False,
            "expand": ["user"],  # Test the expand parameter
        }

        result = await _list_key_helper(**args)

        # Verify that keys were fetched
        mock_find_many_keys.assert_called_once()
        mock_count_keys.assert_called_once()

        # Verify that users were fetched
        # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present
        call_args = mock_find_many_users.call_args
        assert call_args is not None
        where_clause = call_args.kwargs["where"]
        assert "user_id" in where_clause
        assert "in" in where_clause["user_id"]
        user_ids_in_query = set(where_clause["user_id"]["in"])
        assert user_ids_in_query == {"user123", "user456"}

        # Verify response structure
        assert len(result["keys"]) == 2
        assert result["total_count"] == 2
        assert result["current_page"] == 1
        assert result["total_pages"] == 1

        # Verify that user data is included in the response
        # Since expand=user is specified, keys should be full objects
        assert isinstance(result["keys"][0], UserAPIKeyAuth)
        assert isinstance(result["keys"][1], UserAPIKeyAuth)

        # Verify user data is attached to keys
        assert result["keys"][0].user == {
            "user_id": "user123",
            "user_email": "user1@example.com",
            "user_alias": "User One",
        }
        assert result["keys"][1].user == {
            "user_id": "user456",
            "user_email": "user2@example.com",
            "user_alias": "User Two",
        }


@pytest.mark.asyncio
async def test_list_keys_with_expand_user_includes_created_by_user():
    """
    Test that expand=user also resolves created_by to a user object.
    """
    mock_prisma_client = AsyncMock()

    # Key created by user789 but owned by user123
    key1_dict = {
        "token": "token1",
        "user_id": "user123",
        "created_by": "user789",
        "key_alias": "key1",
        "models": ["gpt-4"],
    }
    mock_key1 = MagicMock()
    mock_key1.token = "token1"
    mock_key1.user_id = "user123"
    mock_key1.created_by = "user789"
    mock_key1.model_dump = MagicMock(return_value=key1_dict)

    mock_find_many_keys = AsyncMock(return_value=[mock_key1])
    mock_count_keys = AsyncMock(return_value=1)

    # Create mock users for both user_id and created_by
    mock_user_owner = MagicMock()
    mock_user_owner.user_id = "user123"
    mock_user_owner.user_email = "owner@example.com"
    mock_user_owner.user_alias = "Owner"
    mock_user_owner.model_dump = MagicMock(
        return_value={
            "user_id": "user123",
            "user_email": "owner@example.com",
            "user_alias": "Owner",
        }
    )

    mock_user_creator = MagicMock()
    mock_user_creator.user_id = "user789"
    mock_user_creator.user_email = "creator@example.com"
    mock_user_creator.user_alias = "Creator"

    mock_find_many_users = AsyncMock(return_value=[mock_user_owner, mock_user_creator])

    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
    mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys
    mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users

    async def mock_attach_object_permission(d, _):
        return d

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.attach_object_permission_to_dict",
        side_effect=mock_attach_object_permission,
    ):
        args = {
            "prisma_client": mock_prisma_client,
            "page": 1,
            "size": 50,
            "user_id": None,
            "team_id": None,
            "organization_id": None,
            "key_alias": None,
            "key_hash": None,
            "exclude_team_id": None,
            "return_full_object": False,
            "admin_team_ids": None,
            "include_created_by_keys": False,
            "expand": ["user"],
        }

        result = await _list_key_helper(**args)

        # Verify that the user lookup included both user_id and created_by
        call_args = mock_find_many_users.call_args
        user_ids_in_query = set(call_args.kwargs["where"]["user_id"]["in"])
        assert user_ids_in_query == {"user123", "user789"}

        # Verify created_by_user is attached
        key_result = result["keys"][0]
        assert key_result.created_by_user == {
            "user_id": "user789",
            "user_email": "creator@example.com",
            "user_alias": "Creator",
        }

        # Verify user (owner) is also still attached
        assert key_result.user == {
            "user_id": "user123",
            "user_email": "owner@example.com",
            "user_alias": "Owner",
        }


@pytest.mark.asyncio
async def test_list_keys_with_status_deleted():
    """
    Test that status="deleted" parameter correctly queries the deleted keys table.
    """
    mock_prisma_client = AsyncMock()

    # Mock deleted keys table
    mock_deleted_key1 = MagicMock()
    mock_deleted_key1.token = "deleted_token1"
    mock_deleted_key1.user_id = "user123"
    mock_deleted_key1.dict.return_value = {
        "token": "deleted_token1",
        "user_id": "user123",
        "key_alias": "deleted_key1",
    }

    mock_deleted_key2 = MagicMock()
    mock_deleted_key2.token = "deleted_token2"
    mock_deleted_key2.user_id = "user456"
    mock_deleted_key2.dict.return_value = {
        "token": "deleted_token2",
        "user_id": "user456",
        "key_alias": "deleted_key2",
    }

    mock_find_many_deleted = AsyncMock(
        return_value=[mock_deleted_key1, mock_deleted_key2]
    )
    mock_count_deleted = AsyncMock(return_value=2)

    # Mock regular keys table (should not be called)
    mock_find_many_regular = AsyncMock(return_value=[])
    mock_count_regular = AsyncMock(return_value=0)

    mock_prisma_client.db.litellm_deletedverificationtoken.find_many = (
        mock_find_many_deleted
    )
    mock_prisma_client.db.litellm_deletedverificationtoken.count = mock_count_deleted
    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_regular
    mock_prisma_client.db.litellm_verificationtoken.count = mock_count_regular

    args = {
        "prisma_client": mock_prisma_client,
        "page": 1,
        "size": 50,
        "user_id": None,
        "team_id": None,
        "organization_id": None,
        "key_alias": None,
        "key_hash": None,
        "exclude_team_id": None,
        "return_full_object": False,
        "admin_team_ids": None,
        "include_created_by_keys": False,
        "status": "deleted",  # Test the status parameter
    }

    result = await _list_key_helper(**args)

    # Verify that deleted table was queried
    mock_find_many_deleted.assert_called_once()
    mock_count_deleted.assert_called_once()

    # Verify that regular table was NOT queried
    mock_find_many_regular.assert_not_called()
    mock_count_regular.assert_not_called()

    # Verify response structure
    assert len(result["keys"]) == 2
    assert result["total_count"] == 2
    assert result["current_page"] == 1
    assert result["total_pages"] == 1


@pytest.mark.asyncio
async def test_list_keys_with_invalid_status():
    """
    Test that invalid status parameter raises ProxyException.
    """
    from unittest.mock import Mock, patch

    mock_prisma_client = AsyncMock()

    # Mock the endpoint function directly to test validation
    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import list_keys
    from litellm.proxy.utils import ProxyException

    mock_request = Mock()
    mock_user_api_key_dict = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)

    # Mock prisma_client to be non-None
    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        # Should raise ProxyException for invalid status (HTTPException is caught and re-raised as ProxyException)
        with pytest.raises(ProxyException) as exc_info:
            await list_keys(
                request=mock_request,
                user_api_key_dict=mock_user_api_key_dict,
                status="invalid_status",  # Invalid status value
            )

        # Verify ProxyException properties
        assert exc_info.value.code == "400"
        assert "Invalid status value" in str(exc_info.value.message)
        assert "deleted" in str(exc_info.value.message)


@pytest.mark.asyncio
async def test_list_keys_non_admin_user_id_auto_set():
    """
    Test that when a non-admin user calls list_keys with user_id=None,
    the user_id is automatically set to the authenticated user's user_id.
    """
    from unittest.mock import Mock, patch

    mock_prisma_client = AsyncMock()

    # Create a non-admin user with a user_id
    test_user_id = "test-user-123"
    mock_user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id=test_user_id,
    )

    # Mock user info returned by validate_key_list_check
    mock_user_info = LiteLLM_UserTable(
        user_id=test_user_id,
        user_email="test@example.com",
        teams=[],
        organization_memberships=[],
    )

    # Mock _list_key_helper to capture the user_id argument
    mock_list_key_helper = AsyncMock(
        return_value={
            "keys": [],
            "total_count": 0,
            "current_page": 1,
            "total_pages": 0,
        }
    )

    # Mock prisma_client to be non-None
    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check",
            return_value=mock_user_info,
        ):
            with patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_admin_team_ids",
                return_value=[],
            ):
                with patch(
                    "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper",
                    mock_list_key_helper,
                ):
                    mock_request = Mock()

                    # Call list_keys with user_id=None
                    await list_keys(
                        request=mock_request,
                        user_api_key_dict=mock_user_api_key_dict,
                        user_id=None,  # This should be auto-set to test_user_id
                        status=None,  # Explicitly set status to None to avoid validation errors
                    )

                    # Verify that _list_key_helper was called with user_id set to the authenticated user's user_id
                    mock_list_key_helper.assert_called_once()
                    call_kwargs = mock_list_key_helper.call_args.kwargs
                    assert call_kwargs["user_id"] == test_user_id, (
                        f"Expected user_id to be set to {test_user_id}, "
                        f"but got {call_kwargs.get('user_id')}"
                    )


def _make_member_team_table(
    team_id: str,
    member_user_id: str,
    member_role: str = "user",
    team_member_permissions=None,
):
    """Build a LiteLLM_TeamTable with a single member, suitable for list_keys tests."""
    from litellm.proxy._types import LiteLLM_TeamTable, Member

    return LiteLLM_TeamTable(
        team_id=team_id,
        members_with_roles=[Member(user_id=member_user_id, role=member_role)],
        team_member_permissions=team_member_permissions,
    )


async def _invoke_list_keys_and_capture_helper_kwargs(
    user_api_key_dict,
    team_objects,
    *,
    include_team_keys: bool = True,
):
    """
    Invoke list_keys with mocked dependencies and return the kwargs that
    list_keys passes to _list_key_helper (so tests can assert on
    admin_team_ids / member_team_ids classification).
    """
    from unittest.mock import Mock, patch

    from litellm.proxy._types import LiteLLM_UserTable

    mock_prisma_client = AsyncMock()
    mock_user_info = LiteLLM_UserTable(
        user_id=user_api_key_dict.user_id,
        user_email="member@example.com",
        teams=[t.team_id for t in team_objects],
        organization_memberships=[],
    )
    mock_list_key_helper = AsyncMock(
        return_value={
            "keys": [],
            "total_count": 0,
            "current_page": 1,
            "total_pages": 0,
        }
    )
    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check",
            return_value=mock_user_info,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._fetch_user_team_objects",
            AsyncMock(return_value=team_objects),
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper",
            mock_list_key_helper,
        ),
    ):
        await list_keys(
            request=Mock(),
            user_api_key_dict=user_api_key_dict,
            include_team_keys=include_team_keys,
            status=None,
        )
    mock_list_key_helper.assert_called_once()
    return mock_list_key_helper.call_args.kwargs


@pytest.mark.asyncio
async def test_list_keys_team_member_with_key_list_permission_sees_all_team_keys():
    """
    Bug fix: when a team has /key/list in team_member_permissions, regular
    team members must get full key visibility for that team — same as a team
    admin would. This means other members' keys AND service account keys
    (user_id=NULL) must be returned, not only the caller's own keys.

    This test pins down list_keys' classification: the team must be passed
    to _list_key_helper as a full-visibility team (admin_team_ids), not as
    a service-account-only team (member_team_ids).
    """
    member_user_id = "member-user-1"
    team_id = "team-with-permission"
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id=member_user_id,
    )
    team_objects = [
        _make_member_team_table(
            team_id=team_id,
            member_user_id=member_user_id,
            member_role="user",
            team_member_permissions=["/key/list"],
        )
    ]

    helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs(
        user_api_key_dict=user_api_key_dict,
        team_objects=team_objects,
    )

    admin_team_ids = helper_kwargs.get("admin_team_ids") or []
    member_team_ids = helper_kwargs.get("member_team_ids") or []
    assert team_id in admin_team_ids, (
        "team granting /key/list permission must be classified as full-visibility "
        "(admin_team_ids), but got "
        f"admin_team_ids={admin_team_ids}, member_team_ids={member_team_ids}"
    )


@pytest.mark.asyncio
async def test_list_keys_team_member_without_key_list_permission_only_service_accounts():
    """
    Without the /key/list permission, the existing scoping must hold: the
    team is classified as member-only, so only service account keys
    (user_id=NULL) for that team are visible to the caller.
    """
    member_user_id = "member-user-2"
    team_id = "team-no-permission"
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id=member_user_id,
    )
    team_objects = [
        _make_member_team_table(
            team_id=team_id,
            member_user_id=member_user_id,
            member_role="user",
            team_member_permissions=None,
        )
    ]

    helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs(
        user_api_key_dict=user_api_key_dict,
        team_objects=team_objects,
    )

    admin_team_ids = helper_kwargs.get("admin_team_ids") or []
    member_team_ids = helper_kwargs.get("member_team_ids") or []
    assert team_id not in admin_team_ids
    assert team_id in member_team_ids


@pytest.mark.asyncio
async def test_list_keys_team_member_with_permission_in_one_team_only():
    """
    Granular: a user is a member of two teams. Only one team grants
    /key/list — the other does not. The classification must respect the
    per-team permission, not leak full visibility across both teams.
    """
    member_user_id = "member-user-3"
    team_with_permission = "team-A"
    team_without_permission = "team-B"
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id=member_user_id,
    )
    team_objects = [
        _make_member_team_table(
            team_id=team_with_permission,
            member_user_id=member_user_id,
            member_role="user",
            team_member_permissions=["/key/list"],
        ),
        _make_member_team_table(
            team_id=team_without_permission,
            member_user_id=member_user_id,
            member_role="user",
            team_member_permissions=[],
        ),
    ]

    helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs(
        user_api_key_dict=user_api_key_dict,
        team_objects=team_objects,
    )

    admin_team_ids = helper_kwargs.get("admin_team_ids") or []
    member_team_ids = helper_kwargs.get("member_team_ids") or []
    assert team_with_permission in admin_team_ids
    assert team_without_permission not in admin_team_ids
    assert team_without_permission in member_team_ids


@pytest.mark.asyncio
async def test_list_keys_team_admin_unaffected_by_member_permission_logic():
    """
    Sanity: a team admin's classification is unchanged by the new
    permission-aware path. They still appear in admin_team_ids (full
    visibility) regardless of the team_member_permissions value.
    """
    admin_user_id = "team-admin-user"
    team_id = "team-with-admin"
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id=admin_user_id,
    )
    team_objects = [
        _make_member_team_table(
            team_id=team_id,
            member_user_id=admin_user_id,
            member_role="admin",
            team_member_permissions=None,
        )
    ]

    helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs(
        user_api_key_dict=user_api_key_dict,
        team_objects=team_objects,
    )

    admin_team_ids = helper_kwargs.get("admin_team_ids") or []
    assert team_id in admin_team_ids


def test_build_key_filter_conditions_full_visibility_team_includes_service_accounts():
    """
    Direct check on the SQL filter: when a team is in the full-visibility
    set (admin_team_ids), the filter clause for that team is
    {"team_id": {"in": [...]}} with NO user_id constraint — so service
    account keys (user_id=NULL) AND other members' keys are returned.

    This is the SQL-level proof that a member with /key/list permission
    will see service account keys for the team.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    full_visibility_team = "team-full-vis"
    where = _build_key_filter_conditions(
        user_id="member-user-x",
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=[full_visibility_team],
        member_team_ids=[full_visibility_team],
        include_created_by_keys=False,
    )

    serialized = json.dumps(where)
    # Full-visibility team filter: no user_id restriction
    assert (
        json.dumps({"team_id": {"in": [full_visibility_team]}}) in serialized
    ), f"expected unrestricted team_id IN clause, got: {serialized}"
    # No service-account-only AND clause for this team (it would be redundant
    # and would erroneously narrow the visibility back to user_id=NULL).
    sa_only_clause = json.dumps(
        {"AND": [{"team_id": {"in": [full_visibility_team]}}, {"user_id": None}]}
    )
    assert (
        sa_only_clause not in serialized
    ), f"team in admin_team_ids must not also be filtered to user_id=NULL: {serialized}"


def test_build_key_filter_conditions_member_only_team_restricts_to_service_accounts():
    """
    Existing-behavior pin: when a team is ONLY in member_team_ids (not
    admin_team_ids), the filter for that team must be
    {"AND": [{"team_id": {"in": [...]}}, {"user_id": None}]} — i.e. only
    service accounts visible. This is the "no permission" baseline that
    must keep working.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    member_only_team = "team-member-only"
    where = _build_key_filter_conditions(
        user_id="member-user-y",
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=[],
        member_team_ids=[member_only_team],
        include_created_by_keys=False,
    )

    serialized = json.dumps(where)
    expected = json.dumps(
        {"AND": [{"team_id": {"in": [member_only_team]}}, {"user_id": None}]}
    )
    assert (
        expected in serialized
    ), f"member-only team must be restricted to user_id=NULL keys, got: {serialized}"


@pytest.mark.asyncio
async def test_generate_key_negative_max_budget():
    """
    Test that GenerateKeyRequest model allows negative max_budget values.
    Validation is done at API level, not model level.

    This prevents GET requests from breaking when they receive data with negative budgets.
    """
    # Should not raise any errors at model level
    request = GenerateKeyRequest(max_budget=-7.0)
    assert request.max_budget == -7.0


@pytest.mark.asyncio
async def test_generate_key_negative_soft_budget():
    """
    Test that GenerateKeyRequest model allows negative soft_budget values.
    Validation is done at API level, not model level.
    """
    # Should not raise any errors at model level
    request = GenerateKeyRequest(soft_budget=-10.0)
    assert request.soft_budget == -10.0


@pytest.mark.asyncio
async def test_generate_key_positive_budgets_accepted():
    """
    Test that GenerateKeyRequest accepts positive budget values.
    """
    # Should not raise any errors
    request = GenerateKeyRequest(max_budget=100.0, soft_budget=50.0)
    assert request.max_budget == 100.0
    assert request.soft_budget == 50.0


@pytest.mark.asyncio
async def test_update_key_negative_max_budget():
    """
    Test that UpdateKeyRequest model allows negative max_budget values.
    Validation is done at API level, not model level.
    """
    # Should not raise any errors at model level
    request = UpdateKeyRequest(key="test-key", max_budget=-5.0)
    assert request.max_budget == -5.0


@pytest.mark.asyncio
async def test_generate_key_with_router_settings(monkeypatch):
    """
    Test that /key/generate correctly handles router_settings by:
    1. Accepting router_settings as a dict parameter
    2. Serializing router_settings to JSON when saving to database
    3. Storing router_settings in the key record
    """
    mock_prisma_client = AsyncMock()
    mock_prisma_client.jsonify_object = lambda data: data

    # Mock prisma_client.insert_data for both user and key tables
    async def _insert_data_side_effect(*args, **kwargs):
        table_name = kwargs.get("table_name")
        if table_name == "user":
            return MagicMock(models=[], spend=0)
        elif table_name == "key":
            return MagicMock(
                token="hashed_token_router",
                litellm_budget_table=None,
                object_permission=None,
            )
        return MagicMock()

    mock_prisma_client.insert_data = AsyncMock(side_effect=_insert_data_side_effect)
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )
    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
    from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        generate_key_fn,
    )

    # Test router_settings with sample data
    # Using valid UpdateRouterConfig fields (retry_policy is not a valid field,
    # but model_group_retry_policy is, which also tests nested dict serialization)
    router_settings_data = {
        "routing_strategy": "usage-based",
        "num_retries": 3,
        "model_group_retry_policy": {"max_retries": 5},
    }

    request_data = GenerateKeyRequest(
        models=["gpt-4"],
        router_settings=router_settings_data,
    )

    await generate_key_fn(
        data=request_data,
        user_api_key_dict=UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-1234",
            user_id="user-router-1",
        ),
    )

    # Verify key insertion was called
    assert mock_prisma_client.insert_data.call_count >= 1
    key_insert_calls = [
        call.kwargs
        for call in mock_prisma_client.insert_data.call_args_list
        if call.kwargs.get("table_name") == "key"
    ]
    assert len(key_insert_calls) >= 1
    key_data = key_insert_calls[0]["data"]

    # Verify router_settings is present
    assert "router_settings" in key_data

    # router_settings should be present in the data passed to insert_data
    # The code uses safe_dumps to serialize router_settings, so it will be a JSON string
    router_settings_value = key_data["router_settings"]

    # Get the actual settings value for comparison
    # The code uses safe_dumps to serialize and yaml.safe_load to deserialize
    if isinstance(router_settings_value, str):
        # If it's a JSON string (from safe_dumps), deserialize it using json.loads
        # (safe_dumps produces JSON, and json.loads is the correct way to deserialize it)
        actual_settings = json.loads(router_settings_value)
    elif isinstance(router_settings_value, dict):
        # If it's still a dict, use it directly
        actual_settings = router_settings_value
    else:
        raise AssertionError(
            f"router_settings should be str or dict, got {type(router_settings_value)}"
        )

    # Verify router_settings matches input (regardless of serialization state)
    assert actual_settings == router_settings_data


@pytest.mark.asyncio
async def test_update_key_with_router_settings(monkeypatch):
    """
    Test that /key/update correctly handles router_settings by:
    1. Accepting router_settings as a dict parameter
    2. Serializing router_settings to JSON when updating database
    3. Updating router_settings in the key record
    """
    from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        prepare_key_update_data,
    )

    # Mock existing key
    existing_key = LiteLLM_VerificationToken(
        token="test-token-router",
        key_alias="test-key",
        models=["gpt-3.5-turbo"],
        user_id="test-user",
        team_id=None,
        auto_rotate=False,
        rotation_interval=None,
        metadata={},
    )

    # Test updating router_settings
    router_settings_data = {
        "routing_strategy": "latency-based",
        "num_retries": 2,
    }

    update_request = UpdateKeyRequest(
        key="test-token-router", router_settings=router_settings_data
    )

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    # Verify router_settings is serialized to JSON string
    assert "router_settings" in result
    assert isinstance(result["router_settings"], str)

    # Verify router_settings can be deserialized and matches input
    deserialized_settings = json.loads(result["router_settings"])
    assert deserialized_settings == router_settings_data


@pytest.mark.asyncio
async def test_validate_max_budget():
    """
    Test _validate_max_budget helper function.

    Tests:
    1. Positive max_budget should pass
    2. Zero max_budget should pass
    3. Negative max_budget should raise HTTPException
    4. None max_budget should pass
    """
    from fastapi import HTTPException

    # Test Case 1: Positive max_budget should pass
    try:
        _validate_max_budget(100.0)
        _validate_max_budget(0.0)
    except HTTPException:
        pytest.fail("_validate_max_budget raised HTTPException for valid values")

    # Test Case 2: None max_budget should pass
    try:
        _validate_max_budget(None)
    except HTTPException:
        pytest.fail("_validate_max_budget raised HTTPException for None")

    # Test Case 3: Negative max_budget should raise HTTPException
    with pytest.raises(HTTPException) as exc_info:
        _validate_max_budget(-10.0)

    assert exc_info.value.status_code == 400
    assert "max_budget must be a non-negative finite number" in str(
        exc_info.value.detail
    )
    assert "negative" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_get_and_validate_existing_key():
    """
    Test _get_and_validate_existing_key helper function.

    Tests:
    1. Successfully retrieve existing key
    2. Key not found raises ProxyException
    3. Database not connected raises HTTPException
    """
    from fastapi import HTTPException

    from litellm.proxy._types import ProxyException

    # Test Case 1: Successfully retrieve existing key
    mock_prisma_client = AsyncMock()
    mock_key = LiteLLM_VerificationToken(
        token="test-key-123",
        user_id="user-123",
        models=["gpt-4"],
        team_id=None,
    )
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_key
    )

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
        return_value="hashed-test-key-123",
    ):
        result = await _get_and_validate_existing_key(
            token="test-key-123",
            prisma_client=mock_prisma_client,
        )

    assert result == mock_key
    mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with(
        where={"token": "hashed-test-key-123"}
    )

    # Test Case 2: Key not found raises ProxyException
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
        return_value="hashed-non-existent-key",
    ):
        with pytest.raises(ProxyException) as exc_info:
            await _get_and_validate_existing_key(
                token="non-existent-key",
                prisma_client=mock_prisma_client,
            )

    assert str(exc_info.value.code) == "404"
    assert "Key not found" in exc_info.value.message

    # Test Case 3: Database not connected raises HTTPException
    with pytest.raises(HTTPException) as exc_info:
        await _get_and_validate_existing_key(
            token="test-key-123",
            prisma_client=None,
        )

    assert exc_info.value.status_code == 500
    assert "Database not connected" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_process_single_key_update():
    """
    Test _process_single_key_update helper function.

    Tests successful key update with all validations passing.
    """
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateKeyRequestItem,
    )

    # Setup mocks
    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()
    mock_llm_router = MagicMock()

    # Mock existing key
    existing_key = LiteLLM_VerificationToken(
        token="test-key-123",
        user_id="user-123",
        models=["gpt-4"],
        team_id=None,
        max_budget=None,
        tags=None,
    )

    # Mock updated key response
    updated_key_data = {
        "user_id": "user-123",
        "models": ["gpt-4"],
        "team_id": None,
        "max_budget": 100.0,
        "tags": ["production"],
    }

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=existing_key
    )
    mock_updated_key_obj = MagicMock()
    mock_updated_key_obj.model_dump.return_value = updated_key_data
    mock_prisma_client.update_data = AsyncMock(
        return_value={"data": mock_updated_key_obj}
    )

    # Mock prepare_key_update_data
    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
    ) as mock_prepare:
        mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]}

        # Mock TeamMemberPermissionChecks
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
        ) as mock_permission_check:
            mock_permission_check.return_value = None

            # Mock _delete_cache_key_object
            with patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
            ) as mock_delete_cache:
                mock_delete_cache.return_value = None

                # Mock hash_token (imported from litellm.proxy._types)
                with patch("litellm.proxy._types.hash_token") as mock_hash:
                    mock_hash.return_value = "hashed-test-key-123"

                    # Mock _hash_token_if_needed
                    with patch(
                        "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
                        return_value="hashed-test-key-123",
                    ):
                        # Mock KeyManagementEventHooks
                        with patch(
                            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
                        ):
                            # Create update request
                            update_key_request = UpdateKeyRequest(
                                key="test-key-123",
                                max_budget=100.0,
                                tags=["production"],
                            )

                            user_api_key_dict = UserAPIKeyAuth(
                                user_role=LitellmUserRoles.PROXY_ADMIN,
                                api_key="sk-admin",
                                user_id="admin-user",
                            )

                            # Call the function
                            result = await _process_single_key_update(
                                update_key_request=update_key_request,
                                user_api_key_dict=user_api_key_dict,
                                litellm_changed_by=None,
                                prisma_client=mock_prisma_client,
                                user_api_key_cache=mock_user_api_key_cache,
                                proxy_logging_obj=mock_proxy_logging_obj,
                                llm_router=mock_llm_router,
                            )

                            # Verify results
                            assert result is not None
                            assert "token" not in result  # Token should be removed
                            assert result.get("max_budget") == 100.0
                            assert result.get("tags") == ["production"]

                            # Verify mocks were called
                            mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once()
                            mock_prisma_client.update_data.assert_called_once()
                            mock_delete_cache.assert_called_once()


@pytest.mark.asyncio
async def test_bulk_update_keys_success(monkeypatch):
    """
    Test /key/bulk_update endpoint with successful updates.

    Tests:
    1. Multiple keys updated successfully
    2. Response contains correct counts and data
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_keys,
    )
    from litellm.proxy.proxy_server import (
        llm_router,
        prisma_client,
        proxy_logging_obj,
        user_api_key_cache,
    )
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateKeyRequest,
        BulkUpdateKeyRequestItem,
    )

    # Setup mocks
    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()
    mock_llm_router = MagicMock()

    # Mock existing keys
    existing_key_1 = LiteLLM_VerificationToken(
        token="test-key-1",
        user_id="user-123",
        models=["gpt-4"],
        team_id=None,
        max_budget=None,
    )
    existing_key_2 = LiteLLM_VerificationToken(
        token="test-key-2",
        user_id="user-123",
        models=["gpt-3.5-turbo"],
        team_id=None,
        max_budget=50.0,
    )

    # Mock updated key responses
    updated_key_1_data = {
        "user_id": "user-123",
        "models": ["gpt-4"],
        "max_budget": 100.0,
        "tags": ["production"],
    }
    updated_key_2_data = {
        "user_id": "user-123",
        "models": ["gpt-3.5-turbo"],
        "max_budget": 200.0,
        "tags": ["staging"],
    }

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        side_effect=[existing_key_1, existing_key_2]
    )
    mock_updated_key_1_obj = MagicMock()
    mock_updated_key_1_obj.model_dump.return_value = updated_key_1_data
    mock_updated_key_2_obj = MagicMock()
    mock_updated_key_2_obj.model_dump.return_value = updated_key_2_data
    mock_prisma_client.update_data = AsyncMock(
        side_effect=[
            {"data": mock_updated_key_1_obj},
            {"data": mock_updated_key_2_obj},
        ]
    )

    # Patch dependencies
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router)

    # Mock helper functions
    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
    ) as mock_prepare:
        mock_prepare.side_effect = [
            {"max_budget": 100.0, "tags": ["production"]},
            {"max_budget": 200.0, "tags": ["staging"]},
        ]

        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
        ):
            with patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
            ):
                with patch("litellm.proxy._types.hash_token") as mock_hash:
                    mock_hash.side_effect = ["hashed-key-1", "hashed-key-2"]

                    def _hash_for_bulk_success(token: str) -> str:
                        return {
                            "test-key-1": "hashed-key-1",
                            "test-key-2": "hashed-key-2",
                        }[token]

                    with patch(
                        "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
                        side_effect=_hash_for_bulk_success,
                    ):
                        with patch(
                            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
                        ):
                            # Create request
                            request_data = BulkUpdateKeyRequest(
                                keys=[
                                    BulkUpdateKeyRequestItem(
                                        key="test-key-1",
                                        max_budget=100.0,
                                        tags=["production"],
                                    ),
                                    BulkUpdateKeyRequestItem(
                                        key="test-key-2",
                                        max_budget=200.0,
                                        tags=["staging"],
                                    ),
                                ]
                            )

                            user_api_key_dict = UserAPIKeyAuth(
                                user_role=LitellmUserRoles.PROXY_ADMIN,
                                api_key="sk-admin",
                                user_id="admin-user",
                            )

                            # Call endpoint
                            response = await bulk_update_keys(
                                data=request_data,
                                user_api_key_dict=user_api_key_dict,
                                litellm_changed_by=None,
                            )

                            # Verify response
                            assert response.total_requested == 2
                            assert len(response.successful_updates) == 2
                            assert len(response.failed_updates) == 0
                            assert response.successful_updates[0].key == "test-key-1"
                            assert response.successful_updates[1].key == "test-key-2"


@pytest.mark.asyncio
async def test_bulk_update_keys_partial_failures(monkeypatch):
    """
    Test /key/bulk_update endpoint with partial failures.

    Tests:
    1. Some keys update successfully, others fail
    2. Response contains both successful and failed updates
    3. Failed updates include error messages
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_keys,
    )
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateKeyRequest,
        BulkUpdateKeyRequestItem,
    )

    # Setup mocks
    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()
    mock_llm_router = MagicMock()

    # Mock existing keys
    existing_key_1 = LiteLLM_VerificationToken(
        token="test-key-1",
        user_id="user-123",
        models=["gpt-4"],
        team_id=None,
        max_budget=None,
    )

    # Mock updated key response for successful update
    updated_key_1_data = {
        "user_id": "user-123",
        "models": ["gpt-4"],
        "max_budget": 100.0,
        "tags": ["production"],
    }

    # First key exists, second key doesn't exist
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        side_effect=[existing_key_1, None]  # Second key not found
    )
    mock_updated_key_1_obj = MagicMock()
    mock_updated_key_1_obj.model_dump.return_value = updated_key_1_data
    mock_prisma_client.update_data = AsyncMock(
        return_value={"data": mock_updated_key_1_obj}
    )
    # Mock get_data for the error handler path (used to fetch key_info on failure)
    mock_prisma_client.get_data = AsyncMock(return_value=None)

    # Patch dependencies
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router)

    # Mock helper functions
    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
    ) as mock_prepare:
        mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]}

        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
        ):
            with patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
            ):
                with patch("litellm.proxy._types.hash_token") as mock_hash:
                    mock_hash.return_value = "hashed-key-1"

                    def _hash_for_bulk_partial(token: str) -> str:
                        return {
                            "test-key-1": "hashed-key-1",
                            "non-existent-key": "hashed-non-existent-key",
                        }[token]

                    with patch(
                        "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
                        side_effect=_hash_for_bulk_partial,
                    ):
                        with patch(
                            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
                        ):
                            # Create request with one valid and one invalid key
                            request_data = BulkUpdateKeyRequest(
                                keys=[
                                    BulkUpdateKeyRequestItem(
                                        key="test-key-1",
                                        max_budget=100.0,
                                        tags=["production"],
                                    ),
                                    BulkUpdateKeyRequestItem(
                                        key="non-existent-key",
                                        max_budget=200.0,
                                        tags=["staging"],
                                    ),
                                ]
                            )

                            user_api_key_dict = UserAPIKeyAuth(
                                user_role=LitellmUserRoles.PROXY_ADMIN,
                                api_key="sk-admin",
                                user_id="admin-user",
                            )

                            # Call endpoint
                            response = await bulk_update_keys(
                                data=request_data,
                                user_api_key_dict=user_api_key_dict,
                                litellm_changed_by=None,
                            )

                            # Verify response
                            assert response.total_requested == 2
                            assert len(response.successful_updates) == 1
                            assert len(response.failed_updates) == 1
                            assert response.successful_updates[0].key == "test-key-1"
                            assert response.failed_updates[0].key == "non-existent-key"
                            assert (
                                "Key not found"
                                in response.failed_updates[0].failed_reason
                            )


@pytest.mark.parametrize(
    "reset_to,key_spend,key_max_budget,budget_max_budget,expected_error",
    [
        ("not_a_number", 100.0, None, None, "reset_to must be a float"),
        (None, 100.0, None, None, "reset_to must be a float"),
        ([], 100.0, None, None, "reset_to must be a float"),
        ({}, 100.0, None, None, "reset_to must be a float"),
        (-1.0, 100.0, None, None, "reset_to must be >= 0"),
        (-0.1, 100.0, None, None, "reset_to must be >= 0"),
        (101.0, 100.0, None, None, "reset_to (101.0) must be <= current spend (100.0)"),
        (150.0, 100.0, None, None, "reset_to (150.0) must be <= current spend (100.0)"),
        (50.0, 100.0, 30.0, None, "reset_to (50.0) must be <= budget (30.0)"),
    ],
)
def test_validate_reset_spend_value_invalid(
    reset_to, key_spend, key_max_budget, budget_max_budget, expected_error
):
    key_in_db = LiteLLM_VerificationToken(
        token="test-token",
        user_id="test-user",
        spend=key_spend,
        max_budget=key_max_budget,
        litellm_budget_table=(
            LiteLLM_BudgetTable(
                budget_id="test-budget", max_budget=budget_max_budget
            ).dict()
            if budget_max_budget is not None
            else None
        ),
    )

    with pytest.raises(HTTPException) as exc_info:
        _validate_reset_spend_value(reset_to, key_in_db)

    assert exc_info.value.status_code == 400
    assert expected_error in str(exc_info.value.detail)


@pytest.mark.parametrize(
    "reset_to,key_spend,key_max_budget,budget_max_budget",
    [
        (0.0, 100.0, None, None),
        (0, 100.0, None, None),
        (50.0, 100.0, None, None),
        (100.0, 100.0, None, None),
        (25.0, 100.0, 50.0, None),
        (0.0, 0.0, None, None),
        (10.5, 50.0, 20.0, None),
    ],
)
def test_validate_reset_spend_value_valid(
    reset_to, key_spend, key_max_budget, budget_max_budget
):
    key_in_db = LiteLLM_VerificationToken(
        token="test-token",
        user_id="test-user",
        spend=key_spend,
        max_budget=key_max_budget,
        litellm_budget_table=(
            LiteLLM_BudgetTable(
                budget_id="test-budget", max_budget=budget_max_budget
            ).dict()
            if budget_max_budget is not None
            else None
        ),
    )

    result = _validate_reset_spend_value(reset_to, key_in_db)
    assert result == float(reset_to)


def test_validate_reset_spend_value_no_budget_table():
    key_in_db = LiteLLM_VerificationToken(
        token="test-token",
        user_id="test-user",
        spend=100.0,
        max_budget=50.0,
        litellm_budget_table=None,
    )

    result = _validate_reset_spend_value(25.0, key_in_db)
    assert result == 25.0


def test_validate_reset_spend_value_none_spend():
    key_in_db = LiteLLM_VerificationToken(
        token="test-token",
        user_id="test-user",
        spend=0.0,
        max_budget=None,
        litellm_budget_table=None,
    )

    result = _validate_reset_spend_value(0.0, key_in_db)
    assert result == 0.0

    with pytest.raises(HTTPException) as exc_info:
        _validate_reset_spend_value(1.0, key_in_db)
    assert exc_info.value.status_code == 400
    assert "must be <= current spend" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_reset_key_spend_success(monkeypatch):
    mock_prisma_client = MagicMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    hashed_key = "hashed-test-key"
    key_in_db = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        spend=100.0,
        max_budget=200.0,
        litellm_budget_table=None,
    )

    updated_key = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        spend=50.0,
        max_budget=200.0,
        budget_reset_at=None,
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_in_db
    )
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=updated_key
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )

    with (
        patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_proxy_or_team_admin_for_key"
        ) as mock_check_admin,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
        ) as mock_delete_cache,
        patch(
            "litellm.proxy.proxy_server._invalidate_spend_counter"
        ) as mock_invalidate,
    ):
        mock_hash_token.return_value = hashed_key
        mock_check_admin.return_value = None
        mock_delete_cache.return_value = None

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-admin",
            user_id="admin-user",
        )

        response = await reset_key_spend_fn(
            key="sk-test-key",
            data=ResetSpendRequest(reset_to=50.0),
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

        assert response["spend"] == 50.0
        assert response["previous_spend"] == 100.0
        assert response["key_hash"] == hashed_key
        assert response["max_budget"] == 200.0
        mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
        mock_delete_cache.assert_awaited_once()
        mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}")


@pytest.mark.asyncio
async def test_update_key_spend_invalidates_counter(monkeypatch):
    """
    Test that updating a key's spend via update_key_fn immediately invalidates the spend counter.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = AsyncMock()
    mock_proxy_logging_obj = MagicMock()

    hashed_key = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b"
    key_in_db = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        spend=10.0,
        max_budget=200.0,
        litellm_budget_table=None,
    )

    mock_prisma_client.get_data = AsyncMock(return_value=key_in_db)
    mock_prisma_client.update_data = AsyncMock(return_value={"data": {"spend": 0.0}})
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_in_db
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
    monkeypatch.setattr("litellm.store_audit_logs", False)

    with (
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
        ) as mock_delete_cache,
        patch(
            "litellm.proxy.proxy_server._invalidate_spend_counter"
        ) as mock_invalidate,
    ):
        mock_delete_cache.return_value = None

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-admin",
            user_id="admin-user",
        )

        mock_request = MagicMock()
        mock_request.query_params = {}

        await update_key_fn(
            request=mock_request,
            data=UpdateKeyRequest(key="sk-test-key", spend=0.0),
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

        mock_delete_cache.assert_awaited_once()
        mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}")


@pytest.mark.asyncio
async def test_reset_key_spend_success_team_admin(monkeypatch):
    """Test that team admin can reset key spend for keys in their team."""
    mock_prisma_client = MagicMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    hashed_key = "hashed-test-key"
    team_id = "test-team-123"
    key_in_db = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        team_id=team_id,
        spend=100.0,
        max_budget=200.0,
        litellm_budget_table=None,
    )

    updated_key = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        team_id=team_id,
        spend=50.0,
        max_budget=200.0,
        budget_reset_at=None,
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_in_db
    )
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=updated_key
    )

    # Set up team table with user as admin
    team_table = LiteLLM_TeamTableCachedObj(
        team_id=team_id,
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="team-admin-user", role="admin"),
            Member(user_id="test-user", role="user"),
        ],
    )

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    with (
        patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
        ) as mock_delete_cache,
    ):
        mock_hash_token.return_value = hashed_key
        mock_delete_cache.return_value = None

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.INTERNAL_USER,
            api_key="sk-team-admin",
            user_id="team-admin-user",
        )

        response = await reset_key_spend_fn(
            key="sk-test-key",
            data=ResetSpendRequest(reset_to=50.0),
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

        assert response["spend"] == 50.0
        assert response["previous_spend"] == 100.0
        assert response["key_hash"] == hashed_key
        assert response["max_budget"] == 200.0
        mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once()
        mock_delete_cache.assert_awaited_once()


@pytest.mark.asyncio
async def test_reset_key_spend_key_not_found(monkeypatch):
    mock_prisma_client = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    with patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token:
        mock_hash_token.return_value = "hashed-key"

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-admin",
            user_id="admin-user",
        )

        with pytest.raises(HTTPException) as exc_info:
            await reset_key_spend_fn(
                key="sk-test-key",
                data=ResetSpendRequest(reset_to=50.0),
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )

        assert exc_info.value.status_code == 404
        assert "Key not found" in str(
            exc_info.value.detail
        ) or "Key sk-test-key not found" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_reset_key_spend_db_not_connected(monkeypatch):
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-admin",
        user_id="admin-user",
    )

    with pytest.raises(HTTPException) as exc_info:
        await reset_key_spend_fn(
            key="sk-test-key",
            data=ResetSpendRequest(reset_to=50.0),
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc_info.value.status_code == 500
    assert "DB not connected" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_reset_key_spend_validation_error(monkeypatch):
    mock_prisma_client = MagicMock()
    key_in_db = LiteLLM_VerificationToken(
        token="hashed-key",
        user_id="test-user",
        spend=100.0,
        max_budget=None,
        litellm_budget_table=None,
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_in_db
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    with patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token:
        mock_hash_token.return_value = "hashed-key"

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-admin",
            user_id="admin-user",
        )

        with pytest.raises(HTTPException) as exc_info:
            await reset_key_spend_fn(
                key="sk-test-key",
                data=ResetSpendRequest(reset_to=150.0),
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )

        assert exc_info.value.status_code == 400
        assert "must be <= current spend" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_reset_key_spend_authorization_failure(monkeypatch):
    mock_prisma_client = MagicMock()
    mock_user_api_key_cache = MagicMock()

    hashed_key = "hashed-test-key"
    key_in_db = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        team_id="team-1",
        spend=100.0,
        max_budget=None,
        litellm_budget_table=None,
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_in_db
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )

    with (
        patch("litellm.proxy.proxy_server.hash_token") as mock_hash_token,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_proxy_or_team_admin_for_key"
        ) as mock_check_admin,
    ):
        mock_hash_token.return_value = hashed_key
        mock_check_admin.side_effect = HTTPException(
            status_code=403, detail={"error": "Not authorized"}
        )

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.INTERNAL_USER,
            api_key="sk-user",
            user_id="user-1",
        )

        with pytest.raises(HTTPException) as exc_info:
            await reset_key_spend_fn(
                key="sk-test-key",
                data=ResetSpendRequest(reset_to=50.0),
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )

        assert exc_info.value.status_code == 403


@pytest.mark.asyncio
async def test_reset_key_spend_hashed_key(monkeypatch):
    mock_prisma_client = MagicMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    hashed_key = "already-hashed-key"
    key_in_db = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        spend=100.0,
        max_budget=None,
        litellm_budget_table=None,
    )

    updated_key = LiteLLM_VerificationToken(
        token=hashed_key,
        user_id="test-user",
        spend=50.0,
        max_budget=None,
        budget_reset_at=None,
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_in_db
    )
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=updated_key
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )

    with (
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_proxy_or_team_admin_for_key"
        ) as mock_check_admin,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
        ) as mock_delete_cache,
    ):
        mock_check_admin.return_value = None
        mock_delete_cache.return_value = None

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-admin",
            user_id="admin-user",
        )

        response = await reset_key_spend_fn(
            key=hashed_key,
            data=ResetSpendRequest(reset_to=50.0),
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

        assert response["spend"] == 50.0
        mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with(
            where={"token": hashed_key}, include={"litellm_budget_table": True}
        )


@pytest.mark.asyncio
async def test_validate_key_list_check_proxy_admin():
    mock_prisma_client = AsyncMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        user_id="admin-user",
    )

    result = await validate_key_list_check(
        user_api_key_dict=user_api_key_dict,
        user_id=None,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        prisma_client=mock_prisma_client,
    )

    assert result is None


@pytest.mark.asyncio
async def test_validate_key_list_check_team_admin_success():
    mock_prisma_client = AsyncMock()
    user_info = LiteLLM_UserTable(
        user_id="test-user",
        user_email="test@example.com",
        teams=["team-1"],
        organization_memberships=[],
    )

    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
        return_value=user_info
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="test-user",
    )

    result = await validate_key_list_check(
        user_api_key_dict=user_api_key_dict,
        user_id=None,
        team_id="team-1",
        organization_id=None,
        key_alias=None,
        key_hash=None,
        prisma_client=mock_prisma_client,
    )

    assert result is not None
    assert result.user_id == "test-user"


@pytest.mark.asyncio
async def test_validate_key_list_check_team_admin_fail():
    mock_prisma_client = AsyncMock()
    user_info = LiteLLM_UserTable(
        user_id="test-user",
        user_email="test@example.com",
        teams=["team-1"],
        organization_memberships=[],
    )

    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
        return_value=user_info
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="test-user",
    )

    with pytest.raises(ProxyException) as exc_info:
        await validate_key_list_check(
            user_api_key_dict=user_api_key_dict,
            user_id=None,
            team_id="team-2",
            organization_id=None,
            key_alias=None,
            key_hash=None,
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.code == "403" or exc_info.value.code == 403
    assert "not authorized to check this team's keys" in exc_info.value.message


@pytest.mark.asyncio
async def test_validate_key_list_check_key_hash_authorized():
    mock_prisma_client = AsyncMock()
    user_info = LiteLLM_UserTable(
        user_id="test-user",
        user_email="test@example.com",
        teams=[],
        organization_memberships=[],
    )

    key_info = LiteLLM_VerificationToken(
        token="hashed-key",
        user_id="test-user",
    )

    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
        return_value=user_info
    )
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_info
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="test-user",
    )

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints._can_user_query_key_info"
    ) as mock_can_query:
        mock_can_query.return_value = True

        result = await validate_key_list_check(
            user_api_key_dict=user_api_key_dict,
            user_id=None,
            team_id=None,
            organization_id=None,
            key_alias=None,
            key_hash="hashed-key",
            prisma_client=mock_prisma_client,
        )

        assert result is not None
        assert result.user_id == "test-user"


@pytest.mark.asyncio
async def test_validate_key_list_check_key_hash_unauthorized():
    mock_prisma_client = AsyncMock()
    user_info = LiteLLM_UserTable(
        user_id="test-user",
        user_email="test@example.com",
        teams=[],
        organization_memberships=[],
    )

    key_info = LiteLLM_VerificationToken(
        token="hashed-key",
        user_id="other-user",
    )

    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
        return_value=user_info
    )
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=key_info
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="test-user",
    )

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints._can_user_query_key_info"
    ) as mock_can_query:
        mock_can_query.return_value = False

        with pytest.raises(HTTPException) as exc_info:
            await validate_key_list_check(
                user_api_key_dict=user_api_key_dict,
                user_id=None,
                team_id=None,
                organization_id=None,
                key_alias=None,
                key_hash="hashed-key",
                prisma_client=mock_prisma_client,
            )

        assert exc_info.value.status_code == 403
        assert "not allowed to access this key's info" in str(exc_info.value.detail)


@pytest.mark.asyncio
async def test_validate_key_list_check_key_hash_not_found():
    mock_prisma_client = AsyncMock()
    user_info = LiteLLM_UserTable(
        user_id="test-user",
        user_email="test@example.com",
        teams=[],
        organization_memberships=[],
    )

    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
        return_value=user_info
    )
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        side_effect=Exception("Key not found")
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        user_id="test-user",
    )

    with pytest.raises(ProxyException) as exc_info:
        await validate_key_list_check(
            user_api_key_dict=user_api_key_dict,
            user_id=None,
            team_id=None,
            organization_id=None,
            key_alias=None,
            key_hash="non-existent-key",
            prisma_client=mock_prisma_client,
        )

    assert exc_info.value.code == "403" or exc_info.value.code == 403
    assert "Key Hash not found" in exc_info.value.message


@pytest.mark.asyncio
async def test_key_with_budget_id_does_not_store_budget_duration():
    """
    Test that when a key is created with budget_id but without explicit
    budget_duration, the key does NOT get budget_duration stored on it.

    Keys with budget_id follow their linked budget tier's reset schedule;
    reset_budget_for_keys_linked_to_budgets() resets them when the tier resets.
    This avoids duplicating budget_duration on keys so tier updates apply
    automatically to all linked keys.
    """
    from unittest.mock import AsyncMock, MagicMock, patch

    mock_prisma = MagicMock()

    mock_generate_key = AsyncMock(
        return_value={
            "key": "sk-test-key",
            "expires": None,
            "user_id": "test-user",
            "team_id": None,
        }
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            mock_generate_key,
        ),
    ):
        await _common_key_generation_helper(
            data=GenerateKeyRequest(
                budget_id="7d-budget-tier",
                max_budget=10.0,
                # NOTE: budget_duration is intentionally NOT set here
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN,
                api_key="sk-1234",
                user_id="admin-user",
            ),
            litellm_changed_by=None,
            team_table=None,
        )

    mock_generate_key.assert_awaited_once()
    call_kwargs = mock_generate_key.call_args.kwargs

    # Key should NOT have key_budget_duration - it follows the budget tier's schedule
    assert call_kwargs.get("key_budget_duration") is None, (
        "key_budget_duration should be None for budget-linked keys without explicit "
        f"budget_duration; got: {call_kwargs.get('key_budget_duration')}"
    )

    # No budget tier lookup - we don't copy budget_duration onto the key
    mock_prisma.db.litellm_budgettable.find_unique.assert_not_called()


@pytest.mark.asyncio
async def test_key_does_not_override_explicit_budget_duration():
    """
    Test that when a key is created with both budget_id and an explicit
    budget_duration, the explicit budget_duration takes precedence over
    the budget tier's budget_duration.
    """
    from unittest.mock import AsyncMock, MagicMock, patch

    mock_prisma = MagicMock()
    # The budget tier has budget_duration="7d"
    mock_budget_row = MagicMock()
    mock_budget_row.budget_duration = "7d"
    mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(
        return_value=mock_budget_row
    )

    mock_generate_key = AsyncMock(
        return_value={
            "key": "sk-test-key",
            "expires": None,
            "user_id": "test-user",
            "team_id": None,
        }
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            mock_generate_key,
        ),
    ):
        await _common_key_generation_helper(
            data=GenerateKeyRequest(
                budget_id="7d-budget-tier",
                max_budget=10.0,
                budget_duration="30d",  # explicit budget_duration should take precedence
            ),
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN,
                api_key="sk-1234",
                user_id="admin-user",
            ),
            litellm_changed_by=None,
            team_table=None,
        )

    mock_generate_key.assert_awaited_once()
    call_kwargs = mock_generate_key.call_args.kwargs

    # The explicit budget_duration should take precedence
    assert call_kwargs.get("key_budget_duration") == "30d", (
        "Explicit budget_duration should take precedence over the budget tier's value "
        f"but got: {call_kwargs.get('key_budget_duration')}"
    )

    # The budget tier should NOT have been looked up since budget_duration was explicit
    mock_prisma.db.litellm_budgettable.find_unique.assert_not_called()


@pytest.mark.asyncio
@patch(
    "litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key"
)
async def test_rotate_master_key_model_data_valid_for_prisma(
    mock_rotate_mcp,
):
    """
    Test that _rotate_master_key produces valid data for Prisma create_many().

    Regression test for: master key rotation fails with Prisma validation error
    because created_at/updated_at are None (non-nullable DateTime) and
    litellm_params/model_info are JSON strings (create_many expects dicts).
    """
    from unittest.mock import AsyncMock, MagicMock

    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _rotate_master_key,
    )

    # Setup mock prisma client
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db = MagicMock()

    # Mock model table — return one model
    mock_model = MagicMock()
    mock_model.model_id = "model-1"
    mock_model.model_name = "test-model"
    mock_model.litellm_params = (
        '{"model": "openai/gpt-4", "api_key": "sk-encrypted-old"}'
    )
    mock_model.model_info = '{"id": "model-1"}'
    mock_model.created_by = "admin"
    mock_model.updated_by = "admin"
    mock_prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(
        return_value=[mock_model]
    )

    # Mock transaction context manager
    mock_tx = AsyncMock()
    mock_tx.litellm_proxymodeltable = MagicMock()
    mock_tx.litellm_proxymodeltable.delete_many = AsyncMock()
    mock_tx.litellm_proxymodeltable.create_many = AsyncMock()
    mock_prisma_client.db.tx = MagicMock(
        return_value=AsyncMock(
            __aenter__=AsyncMock(return_value=mock_tx),
            __aexit__=AsyncMock(return_value=False),
        )
    )

    # Mock config table — no env vars
    mock_prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])

    # Mock credentials table — no credentials
    mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(
        return_value=[]
    )

    # Mock MCP rotation
    mock_rotate_mcp.return_value = None

    # Mock proxy_config
    mock_proxy_config = MagicMock()
    mock_proxy_config.decrypt_model_list_from_db.return_value = [
        {
            "model_name": "test-model",
            "litellm_params": {
                "model": "openai/gpt-4",
                "api_key": "sk-decrypted-key",
            },
            "model_info": {"id": "model-1"},
        }
    ]

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-1234",
        user_id="test-user",
    )

    with patch(
        "litellm.proxy.proxy_server.proxy_config",
        mock_proxy_config,
    ):
        await _rotate_master_key(
            prisma_client=mock_prisma_client,
            user_api_key_dict=user_api_key_dict,
            current_master_key="sk-old-master-key",
            new_master_key="sk-new-master-key",
        )

    # Verify create_many was called
    mock_tx.litellm_proxymodeltable.create_many.assert_called_once()

    # Get the data passed to create_many
    call_args = mock_tx.litellm_proxymodeltable.create_many.call_args
    created_models = call_args.kwargs.get("data") or call_args[1].get("data")

    assert len(created_models) == 1
    model_data = created_models[0]

    # Verify timestamps are NOT present (Prisma @default(now()) should apply)
    assert (
        "created_at" not in model_data
    ), "created_at should be excluded so Prisma @default(now()) applies"
    assert (
        "updated_at" not in model_data
    ), "updated_at should be excluded so Prisma @default(now()) applies"

    # Verify litellm_params and model_info are prisma.Json wrappers, NOT JSON strings
    import prisma

    assert isinstance(
        model_data["litellm_params"], prisma.Json
    ), f"litellm_params should be prisma.Json for create_many(), got {type(model_data['litellm_params'])}"
    assert isinstance(
        model_data["model_info"], prisma.Json
    ), f"model_info should be prisma.Json for create_many(), got {type(model_data['model_info'])}"

    # Verify delete_many was called inside the transaction (before create_many)
    mock_tx.litellm_proxymodeltable.delete_many.assert_called_once()


async def test_default_key_generate_params_duration(monkeypatch):
    """
    Test that default_key_generate_params with 'duration' is applied
    when no duration is provided in the key generation request.

    Regression test for bug where 'duration' was missing from the list
    of fields populated from default_key_generate_params.
    """
    import litellm

    mock_prisma_client = AsyncMock()
    mock_insert_data = AsyncMock(
        return_value=MagicMock(
            token="hashed_token_123", litellm_budget_table=None, object_permission=None
        )
    )
    mock_prisma_client.insert_data = mock_insert_data
    mock_prisma_client.db = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken = MagicMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=None
    )
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )
    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=MagicMock(
            token="hashed_token_123", litellm_budget_table=None, object_permission=None
        )
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)

    # Set default_key_generate_params with duration
    original_value = litellm.default_key_generate_params
    litellm.default_key_generate_params = {"duration": "180d"}

    try:
        request = GenerateKeyRequest()  # No duration specified
        response = await _common_key_generation_helper(
            data=request,
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN,
                api_key="sk-1234",
                user_id="1234",
            ),
            litellm_changed_by=None,
            team_table=None,
        )

        # Verify duration was applied from defaults
        assert request.duration == "180d"
    finally:
        litellm.default_key_generate_params = original_value


@pytest.mark.asyncio
async def test_build_key_filter_member_team_service_accounts():
    """
    Test that regular team members can see service accounts (user_id=NULL)
    for their teams, but NOT other members' personal keys.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "regular-member-123"
    member_team_ids = ["team-A", "team-B"]

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=member_team_ids,
        include_created_by_keys=False,
    )

    # Should have AND with OR conditions
    assert "AND" in where
    or_conditions = where["AND"][1]["OR"]

    # Should have 2 conditions: user's own keys + member team service accounts
    assert len(or_conditions) == 2

    # First: user's own keys (exact match — non-admin callers use exact matching)
    user_cond = or_conditions[0]
    assert user_cond["user_id"] == user_id

    # Second: service accounts for member teams (user_id=None AND team_id in member teams)
    service_account_cond = or_conditions[1]
    assert "AND" in service_account_cond
    and_parts = service_account_cond["AND"]
    assert {"team_id": {"in": member_team_ids}} in and_parts
    assert {"user_id": None} in and_parts


@pytest.mark.asyncio
async def test_build_key_filter_admin_sees_all_team_keys():
    """
    Test that team admins see ALL keys for their teams (not just service accounts),
    and that member_team_ids doesn't duplicate admin teams.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "admin-user-123"
    admin_team_ids = ["team-A"]
    member_team_ids = ["team-A", "team-B"]

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=admin_team_ids,
        member_team_ids=member_team_ids,
        include_created_by_keys=False,
    )

    assert "AND" in where
    or_conditions = where["AND"][1]["OR"]

    # Should have 3 conditions:
    # 1. user's own keys
    # 2. admin team keys (all keys for team-A)
    # 3. member-only service accounts (only service accounts for team-B, since team-A is already covered by admin)
    assert len(or_conditions) == 3

    # Find admin condition
    admin_cond = None
    service_account_cond = None
    for cond in or_conditions:
        if isinstance(cond.get("team_id"), dict) and "in" in cond.get("team_id", {}):
            admin_cond = cond
        elif "AND" in cond:
            service_account_cond = cond

    assert admin_cond is not None, "Admin team condition should be present"
    assert admin_cond["team_id"]["in"] == admin_team_ids

    # member-only condition should only include team-B (team-A is covered by admin)
    assert (
        service_account_cond is not None
    ), "Service account condition should be present"
    and_parts = service_account_cond["AND"]
    assert {"team_id": {"in": ["team-B"]}} in and_parts
    assert {"user_id": None} in and_parts


@pytest.mark.asyncio
async def test_build_key_filter_created_by_scoped_to_current_teams():
    """
    Test that created_by filter is scoped to teams user currently belongs to.
    A former team member should NOT see service accounts they created for
    a team they've left.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-456"
    # User is currently only a member of team-A (left team-B)
    member_team_ids = ["team-A"]

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=member_team_ids,
        include_created_by_keys=True,
    )

    assert "AND" in where
    or_conditions = where["AND"][1]["OR"]

    # Find the created_by condition
    created_by_cond = None
    for cond in or_conditions:
        if "AND" in cond:
            and_parts = cond["AND"]
            for part in and_parts:
                if isinstance(part, dict) and "created_by" in part:
                    created_by_cond = cond
                    break

    assert created_by_cond is not None, "Created by condition should be present"

    # created_by should be scoped: created_by=user AND (team_id in [team-A] OR team_id=None)
    and_parts = created_by_cond["AND"]
    assert {"created_by": user_id} in and_parts

    # Find the OR part that scopes to current teams
    team_scope = None
    for part in and_parts:
        if isinstance(part, dict) and "OR" in part:
            team_scope = part["OR"]

    assert team_scope is not None, "Team scope OR condition should be present"
    assert {"team_id": {"in": member_team_ids}} in team_scope
    assert {"team_id": None} in team_scope


@pytest.mark.asyncio
async def test_build_key_filter_created_by_no_teams():
    """
    Test that when user has no team memberships (empty list), created_by
    only returns non-team keys (personal keys).
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-no-teams"
    member_team_ids = []  # User has no team memberships

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=member_team_ids,
        include_created_by_keys=True,
    )

    assert "AND" in where
    or_conditions = where["AND"][1]["OR"]

    # Find the created_by condition
    created_by_cond = None
    for cond in or_conditions:
        if "AND" in cond:
            and_parts = cond["AND"]
            for part in and_parts:
                if isinstance(part, dict) and "created_by" in part:
                    created_by_cond = cond
                    break

    assert created_by_cond is not None
    and_parts = created_by_cond["AND"]
    assert {"created_by": user_id} in and_parts
    assert {"team_id": None} in and_parts
    # Should NOT have an OR with team_id in [] - just a simple team_id=None
    for part in and_parts:
        if isinstance(part, dict) and "OR" in part:
            pytest.fail("Should not have OR condition when member_team_ids is empty")


@pytest.mark.asyncio
async def test_build_key_filter_backward_compat_no_member_team_ids():
    """
    Test backward compatibility: when member_team_ids is None (not provided),
    created_by filter should use the old unrestricted behavior.
    This ensures direct callers of _list_key_helper (like Prometheus) still work.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-789"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,  # Not provided
        include_created_by_keys=True,
    )

    assert "AND" in where
    or_conditions = where["AND"][1]["OR"]

    # Find the created_by condition - should be simple {"created_by": user_id}
    created_by_cond = None
    for cond in or_conditions:
        if "created_by" in cond:
            created_by_cond = cond

    assert created_by_cond is not None
    assert created_by_cond == {"created_by": user_id}
    assert len(created_by_cond) == 1, "Should be simple created_by without team scoping"


@pytest.mark.asyncio
async def test_build_key_filter_admin_all_member_overlap():
    """
    Test that when user is admin of ALL teams they belong to,
    no member-only service account condition is added (would be redundant).
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "admin-all"
    admin_team_ids = ["team-A", "team-B"]
    member_team_ids = ["team-A", "team-B"]

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=admin_team_ids,
        member_team_ids=member_team_ids,
        include_created_by_keys=False,
    )

    assert "AND" in where
    or_conditions = where["AND"][1]["OR"]

    # Should only have 2 conditions: user's own keys + admin team keys
    # No member-only service account condition since all teams are admin
    assert len(or_conditions) == 2

    # Verify no AND condition with user_id=None exists (that's the member-only pattern)
    for cond in or_conditions:
        if "AND" in cond:
            and_parts = cond["AND"]
            if {"user_id": None} in and_parts:
                pytest.fail(
                    "Should not have member-only service account condition "
                    "when user is admin of all teams"
                )


@pytest.mark.asyncio
async def test_build_key_filter_project_id():
    """
    Test that project_id is applied as a global AND condition, narrowing all results
    to keys that belong to the specified project.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-123"
    project_id = "proj-abc"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,
        include_created_by_keys=False,
        project_id=project_id,
    )

    # Should be wrapped in a top-level AND for the project_id filter
    assert "AND" in where
    and_parts = where["AND"]
    assert len(and_parts) == 2

    # Second part of AND should be the project_id filter
    assert {"project_id": project_id} in and_parts


@pytest.mark.asyncio
async def test_build_key_filter_access_group_id():
    """
    Test that access_group_id is applied as a global AND condition using hasSome,
    narrowing results to keys whose access_group_ids array contains the given ID.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-123"
    access_group_id = "ag-xyz"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,
        include_created_by_keys=False,
        access_group_id=access_group_id,
    )

    # Should be wrapped in a top-level AND for the access_group_id filter
    assert "AND" in where
    and_parts = where["AND"]
    assert len(and_parts) == 2

    # Second part of AND should use hasSome for the array field
    assert {"access_group_ids": {"hasSome": [access_group_id]}} in and_parts


@pytest.mark.asyncio
async def test_build_key_filter_project_id_and_access_group_id():
    """
    Test that project_id and access_group_id stack correctly when both are provided.
    Both should be applied as AND conditions, narrowing results to keys that match both.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-123"
    project_id = "proj-abc"
    access_group_id = "ag-xyz"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,
        include_created_by_keys=False,
        project_id=project_id,
        access_group_id=access_group_id,
    )

    # After project_id: {"AND": [visibility_where, {"project_id": ...}]}
    # After access_group_id: {"AND": [above, {"access_group_ids": ...}]}
    assert "AND" in where
    outer_and = where["AND"]
    assert len(outer_and) == 2

    # The access_group_ids filter is the outermost AND
    access_group_filter = outer_and[1]
    assert access_group_filter == {"access_group_ids": {"hasSome": [access_group_id]}}

    # The project_id filter is nested one level in
    inner = outer_and[0]
    assert "AND" in inner
    inner_and = inner["AND"]
    assert {"project_id": project_id} in inner_and


@pytest.mark.asyncio
async def test_build_key_filter_team_id_scoped():
    """
    When team_id is provided, it should act as a global AND filter so keys
    from other teams are excluded — even when the user is admin of multiple teams.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    where = _build_key_filter_conditions(
        user_id="multi-team-user",
        team_id="team-A",
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=["team-A", "team-B"],
        member_team_ids=["team-A", "team-B"],
        include_created_by_keys=True,
    )

    # The team_id filter must be a direct child of the outermost AND,
    # not buried inside an OR branch (which was the bug).
    assert "AND" in where, f"Expected top-level AND, got: {where}"
    outer_and = where["AND"]
    assert {
        "team_id": "team-A"
    } in outer_and, (
        f"Expected {{'team_id': 'team-A'}} as a direct AND condition, got: {outer_and}"
    )


@pytest.mark.asyncio
async def test_build_key_filter_admin_substring_matching():
    """
    Admin callers get substring (contains + insensitive) matching for user_id
    and key_alias when use_substring_matching=True.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "alice"
    key_alias = "prod"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=key_alias,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,
        include_created_by_keys=False,
        use_substring_matching=True,
    )

    # Single OR condition is flattened into the top-level where dict
    assert where["user_id"] == {"contains": user_id, "mode": "insensitive"}
    assert where["key_alias"] == {"contains": key_alias, "mode": "insensitive"}


@pytest.mark.asyncio
async def test_build_key_filter_non_admin_exact_matching():
    """
    Non-admin callers get exact matching for user_id and key_alias when
    use_substring_matching=False (the default).  This prevents a user whose
    ID is a substring of another user's ID from seeing that user's keys.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "alice@example.com"
    key_alias = "my-key"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=key_alias,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,
        include_created_by_keys=False,
        use_substring_matching=False,
    )

    # Single OR condition is flattened into the top-level where dict
    # Exact match — no contains/insensitive wrapping
    assert where["user_id"] == user_id
    assert where["key_alias"] == key_alias


@pytest.mark.asyncio
async def test_build_key_filter_default_is_exact_matching():
    """
    The default for use_substring_matching is False, ensuring backward
    compatibility — callers that don't pass the flag get exact matching.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _build_key_filter_conditions,
    )

    user_id = "user-123"

    where = _build_key_filter_conditions(
        user_id=user_id,
        team_id=None,
        organization_id=None,
        key_alias=None,
        key_hash=None,
        exclude_team_id=None,
        admin_team_ids=None,
        member_team_ids=None,
        include_created_by_keys=False,
    )

    # Single OR condition is flattened into the top-level where dict
    assert where["user_id"] == user_id


@pytest.mark.asyncio
async def test_get_member_team_ids():
    """
    Test that get_member_team_ids returns all teams where user is a member
    (any role), not just admin teams.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        get_member_team_ids,
    )

    user_id = "member-user-123"

    # Create mock user info with teams
    user_info = LiteLLM_UserTable(
        user_id=user_id,
        teams=["team-A", "team-B", "team-C"],
    )

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-test",
        user_id=user_id,
    )

    # Mock prisma client
    mock_prisma_client = AsyncMock()

    # Create mock team objects - user is admin of team-A, member of team-B, not in team-C's members list
    mock_team_a = MagicMock()
    mock_team_a.model_dump.return_value = {
        "team_id": "team-A",
        "team_alias": "Team A",
        "members_with_roles": [
            {"user_id": user_id, "role": "admin", "user_email": None}
        ],
        "max_budget": None,
        "budget_duration": None,
        "budget_reset_at": None,
        "tpm_limit": None,
        "rpm_limit": None,
        "models": [],
        "blocked": False,
    }

    mock_team_b = MagicMock()
    mock_team_b.model_dump.return_value = {
        "team_id": "team-B",
        "team_alias": "Team B",
        "members_with_roles": [
            {"user_id": user_id, "role": "user", "user_email": None}
        ],
        "max_budget": None,
        "budget_duration": None,
        "budget_reset_at": None,
        "tpm_limit": None,
        "rpm_limit": None,
        "models": [],
        "blocked": False,
    }

    mock_team_c = MagicMock()
    mock_team_c.model_dump.return_value = {
        "team_id": "team-C",
        "team_alias": "Team C",
        "members_with_roles": [
            {"user_id": "other-user", "role": "admin", "user_email": None}
        ],
        "max_budget": None,
        "budget_duration": None,
        "budget_reset_at": None,
        "tpm_limit": None,
        "rpm_limit": None,
        "models": [],
        "blocked": False,
    }

    mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(
        return_value=[mock_team_a, mock_team_b, mock_team_c]
    )

    result = await get_member_team_ids(
        complete_user_info=user_info,
        user_api_key_dict=user_api_key_dict,
        prisma_client=mock_prisma_client,
    )

    # Should return team-A and team-B (user is a member of both)
    # Should NOT return team-C (user is not in members list)
    assert sorted(result) == ["team-A", "team-B"]


@pytest.mark.asyncio
async def test_generate_key_with_agent_id():
    """Test that agent_id is accepted in GenerateKeyRequest and passed to generate_key_helper_fn."""
    from litellm.proxy._types import GenerateKeyRequest

    # Verify GenerateKeyRequest accepts agent_id
    request = GenerateKeyRequest(
        key_alias="agent-test-key",
        agent_id="test-agent-123",
        models=[],
    )
    assert request.agent_id == "test-agent-123"
    data_json = request.model_dump(exclude_unset=True, exclude_none=True)
    assert data_json["agent_id"] == "test-agent-123"


@pytest.mark.asyncio
async def test_generate_key_helper_fn_agent_id():
    """Test that generate_key_helper_fn passes agent_id into the insert_data call."""
    from unittest.mock import AsyncMock, MagicMock, call, patch

    import litellm.proxy.management_endpoints.key_management_endpoints as km

    mock_prisma_client = AsyncMock()
    mock_insert = AsyncMock(
        return_value=MagicMock(
            token="sk-test",
            created_at=None,
            updated_at=None,
            litellm_budget_table=None,
        )
    )
    mock_prisma_client.insert_data = mock_insert

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        await generate_key_helper_fn(
            request_type="key",
            agent_id="test-agent-456",
            key_alias="test-agent-key",
            models=[],
            table_name="key",
        )

    assert mock_insert.called, "insert_data was never called"
    # insert_data is called as insert_data(data=key_data, ...)
    call_kwargs = mock_insert.call_args.kwargs
    key_data = call_kwargs.get("data", {})
    assert (
        key_data.get("agent_id") == "test-agent-456"
    ), f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}"


def _make_admin_key_dict() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN.value,
        user_id="admin-user",
    )


@pytest.mark.asyncio
async def test_key_aliases_response_shape():
    """Test that key_aliases returns the correct paginated response shape."""
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.query_raw = AsyncMock(
        side_effect=[
            [{"count": 2}],
            [{"key_alias": "alias-alpha"}, {"key_alias": "alias-beta"}],
        ]
    )

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        result = await key_aliases(
            user_api_key_dict=_make_admin_key_dict(),
            page=1,
            size=50,
            search=None,
        )

    assert result["aliases"] == ["alias-alpha", "alias-beta"]
    assert result["total_count"] == 2
    assert result["current_page"] == 1
    assert result["total_pages"] == 1
    assert result["size"] == 50

    # Both SQL calls must filter out null/empty aliases
    count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0]
    aliases_sql = mock_prisma_client.db.query_raw.call_args_list[1].args[0]
    assert "key_alias IS NOT NULL" in count_sql
    assert "key_alias IS NOT NULL" in aliases_sql


@pytest.mark.asyncio
async def test_key_aliases_pagination_skip_take():
    """Test that LIMIT and OFFSET are correctly derived from page and size."""
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.query_raw = AsyncMock(
        side_effect=[
            [{"count": 120}],
            [],
        ]
    )

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        result = await key_aliases(
            user_api_key_dict=_make_admin_key_dict(),
            page=3,
            size=25,
            search=None,
        )

    assert result["current_page"] == 3
    assert result["size"] == 25
    assert result["total_count"] == 120
    assert result["total_pages"] == 5  # ceil(120 / 25)

    # aliases query params: [UI_SESSION_TOKEN_TEAM_ID, size=25, offset=50]
    aliases_call_args = mock_prisma_client.db.query_raw.call_args_list[1].args
    assert aliases_call_args[-2] == 25  # LIMIT = size
    assert aliases_call_args[-1] == 50  # OFFSET = (3 - 1) * 25


@pytest.mark.asyncio
async def test_key_aliases_search_filter():
    """Test that the search param adds a case-insensitive ILIKE condition."""
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.query_raw = AsyncMock(
        side_effect=[
            [{"count": 0}],
            [],
        ]
    )

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        await key_aliases(
            user_api_key_dict=_make_admin_key_dict(),
            page=1,
            size=50,
            search="my-key",
        )

    count_call = mock_prisma_client.db.query_raw.call_args_list[0]
    count_sql = count_call.args[0]
    count_params = count_call.args[1:]

    assert "ILIKE" in count_sql
    assert "%my-key%" in count_params


@pytest.mark.asyncio
async def test_key_aliases_no_search_omits_ilike_filter():
    """Test that without a search term no ILIKE condition is added."""
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.query_raw = AsyncMock(
        side_effect=[
            [{"count": 0}],
            [],
        ]
    )

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        await key_aliases(
            user_api_key_dict=_make_admin_key_dict(),
            page=1,
            size=50,
            search=None,
        )

    count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0]
    assert "ILIKE" not in count_sql


@pytest.mark.asyncio
async def test_key_aliases_internal_user_scoped_to_own_keys_and_teams():
    """Test that internal users only see aliases for their own keys and team keys."""
    mock_prisma_client = AsyncMock()

    # Mock user table lookup to return teams
    mock_user_row = MagicMock()
    mock_user_row.teams = ["team-abc", "team-xyz"]
    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
        return_value=mock_user_row
    )

    mock_prisma_client.db.query_raw = AsyncMock(
        side_effect=[
            [{"count": 1}],
            [{"key_alias": "my-alias"}],
        ]
    )

    internal_user = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER.value,
        user_id="user-123",
    )

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        result = await key_aliases(
            user_api_key_dict=internal_user,
            page=1,
            size=50,
            search=None,
        )

    assert result["aliases"] == ["my-alias"]

    # Verify the SQL includes user_id and team_id scoping
    count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0]
    assert "user_id = " in count_sql
    assert "team_id IN " in count_sql

    # Verify the params include user_id and team IDs
    count_params = mock_prisma_client.db.query_raw.call_args_list[0].args[1:]
    assert "user-123" in count_params
    assert "team-abc" in count_params
    assert "team-xyz" in count_params


@pytest.mark.asyncio
async def test_key_aliases_admin_sees_all():
    """Test that proxy admins see all aliases without user/team scoping."""
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.query_raw = AsyncMock(
        side_effect=[
            [{"count": 3}],
            [
                {"key_alias": "alias-a"},
                {"key_alias": "alias-b"},
                {"key_alias": "alias-c"},
            ],
        ]
    )

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
        result = await key_aliases(
            user_api_key_dict=_make_admin_key_dict(),
            page=1,
            size=50,
            search=None,
        )

    assert len(result["aliases"]) == 3

    # Admin queries should NOT have user_id or team_id scoping
    count_sql = mock_prisma_client.db.query_raw.call_args_list[0].args[0]
    assert "user_id = " not in count_sql
    assert "team_id IN " not in count_sql


class TestValidateKeyAliasFormat:
    @pytest.fixture(autouse=True)
    def reset_key_alias_flag(self):
        litellm.enable_key_alias_format_validation = False
        yield
        litellm.enable_key_alias_format_validation = False

    def test_validation_skipped_when_flag_disabled(self):
        """When enable_key_alias_format_validation is False (default), no validation occurs."""
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )

        # Even invalid aliases should pass silently when the flag is off
        _validate_key_alias_format(None)
        _validate_key_alias_format("")
        _validate_key_alias_format("!invalid!")
        _validate_key_alias_format("a" * 256)

    def test_validate_key_alias_format_valid(self):
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )

        litellm.enable_key_alias_format_validation = True
        # Valid cases
        _validate_key_alias_format(None)  # OK
        _validate_key_alias_format("valid-alias")
        _validate_key_alias_format("valid_alias")
        _validate_key_alias_format("valid.alias")
        _validate_key_alias_format("valid/alias")
        _validate_key_alias_format("a" * 255)
        _validate_key_alias_format("my-key-123")
        _validate_key_alias_format("user/user@example.com")
        _validate_key_alias_format("team/user@example.com")

    def test_validate_key_alias_format_invalid(self):
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )
        from litellm.proxy._types import ProxyException

        litellm.enable_key_alias_format_validation = True
        invalid_aliases = [
            "",  # empty
            " ",  # whitespace
            "a",  # too short (min 2)
            "!",  # special char
            "-start",  # non-alphanumeric start
            "end-",  # non-alphanumeric end
            "invalid#char",  # invalid char
            "a" * 256,  # too long
            "  leading",
            "trailing  ",
        ]

        for alias in invalid_aliases:
            with pytest.raises(ProxyException) as exc:
                _validate_key_alias_format(alias)
            assert str(exc.value.code) == "400"
            assert "Invalid key_alias format" in str(exc.value.message)


@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_within_bounds():
    """
    Test that _check_org_key_limits works with UpdateKeyRequest when updating
    a key's TPM/RPM limits within organization bounds.
    """
    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[]
    )

    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-update-1",
        organization_alias="test-org",
        budget_id="budget-123",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-123",
            tpm_limit=20000,
            rpm_limit=2000,
        ),
    )

    data = UpdateKeyRequest(
        key="sk-test-key",
        tpm_limit=10000,
        rpm_limit=1000,
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
        organization_id="test-org-update-1",
    )

    # Should not raise any exception
    await _check_org_key_limits(
        org_table=org_table,
        data=data,
        prisma_client=mock_prisma_client,
    )

    mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
        where={"organization_id": "test-org-update-1"}
    )


@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_overallocation():
    """
    Test that _check_org_key_limits raises HTTPException when updating a key
    would exceed organization TPM limits.
    """
    from litellm.proxy._types import hash_token as _hash_token

    existing_key = MagicMock()
    existing_key.token = _hash_token("sk-other-key")
    existing_key.tpm_limit = 15000
    existing_key.rpm_limit = 1500
    existing_key.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[existing_key]
    )

    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-update-2",
        organization_alias="test-org",
        budget_id="budget-456",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-456",
            tpm_limit=20000,
            rpm_limit=2000,
        ),
    )

    data = UpdateKeyRequest(
        key="sk-test-key",
        tpm_limit=10000,  # 15000 + 10000 = 25000 > 20000
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
        organization_id="test-org-update-2",
    )

    with pytest.raises(HTTPException) as exc:
        await _check_org_key_limits(
            org_table=org_table,
            data=data,
            prisma_client=mock_prisma_client,
        )
    assert exc.value.status_code == 400
    assert "TPM limit" in str(exc.value.detail)


@pytest.mark.asyncio
async def test_check_org_key_limits_on_update_excludes_self():
    """
    Test that _check_org_key_limits excludes the key being updated from the
    allocated totals. Without this, the key's current limits would be
    double-counted: once from find_many and once from data.tpm_limit/rpm_limit.
    """
    from litellm.proxy._types import hash_token

    # The key being updated is returned by find_many with its current limits.
    # In the DB, token is stored as a SHA-256 hash of the raw key.
    self_key = MagicMock()
    self_key.token = hash_token("sk-test-key")
    self_key.tpm_limit = 10000
    self_key.rpm_limit = 1000
    self_key.metadata = {}

    # Another key in the org
    other_key = MagicMock()
    other_key.token = hash_token("sk-other-key")
    other_key.tpm_limit = 5000
    other_key.rpm_limit = 500
    other_key.metadata = {}

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[self_key, other_key]
    )

    org_table = LiteLLM_OrganizationTable(
        organization_id="test-org-self",
        organization_alias="test-org",
        budget_id="budget-789",
        models=["gpt-4"],
        created_by="admin",
        updated_by="admin",
        litellm_budget_table=LiteLLM_BudgetTable(
            budget_id="budget-789",
            tpm_limit=20000,
            rpm_limit=2000,
        ),
    )

    # Updating the key to 12000 TPM. Other key uses 5000, so total = 17000 < 20000.
    # Without the fix, this would be 10000 (self) + 5000 (other) + 12000 = 27000 > 20000.
    data = UpdateKeyRequest(
        key="sk-test-key",
        tpm_limit=12000,
        rpm_limit=1200,
        tpm_limit_type="guaranteed_throughput",
        rpm_limit_type="guaranteed_throughput",
        organization_id="test-org-self",
    )

    # Should not raise - the key's own limits should be excluded from the sum
    await _check_org_key_limits(
        org_table=org_table,
        data=data,
        prisma_client=mock_prisma_client,
    )


def test_update_key_skips_org_check_when_no_throughput_fields_changed():
    """
    Test that the org limit check guard condition correctly skips validation
    when only non-throughput fields change on a key that belongs to an org.
    This prevents blocking updates when the org has been deleted.
    """

    def _check_throughput_changed(data: UpdateKeyRequest) -> bool:
        return (
            data.organization_id is not None
            or data.tpm_limit is not None
            or data.rpm_limit is not None
            or data.tpm_limit_type is not None
            or data.rpm_limit_type is not None
        )

    # Updating only key_alias — no throughput fields changed
    data = UpdateKeyRequest(key="sk-test-key", key_alias="new-alias")
    assert _check_throughput_changed(data) is False

    # Updating tpm_limit — throughput field changed
    data_with_tpm = UpdateKeyRequest(key="sk-test-key", tpm_limit=5000)
    assert _check_throughput_changed(data_with_tpm) is True

    # Updating organization_id — org change triggers check
    data_with_org = UpdateKeyRequest(key="sk-test-key", organization_id="new-org")
    assert _check_throughput_changed(data_with_org) is True

    # Updating tpm_limit_type — limit type change triggers check
    data_with_tpm_type = UpdateKeyRequest(
        key="sk-test-key", tpm_limit_type="guaranteed_throughput"
    )
    assert _check_throughput_changed(data_with_tpm_type) is True

    # Updating rpm_limit_type — limit type change triggers check
    data_with_rpm_type = UpdateKeyRequest(
        key="sk-test-key", rpm_limit_type="guaranteed_throughput"
    )
    assert _check_throughput_changed(data_with_rpm_type) is True


def test_update_key_request_has_organization_id():
    """
    Test that UpdateKeyRequest accepts organization_id field.
    """
    data = UpdateKeyRequest(
        key="sk-test-key",
        organization_id="test-org-123",
    )
    assert data.organization_id == "test-org-123"

    # Also verify it defaults to None
    data_no_org = UpdateKeyRequest(key="sk-test-key")
    assert data_no_org.organization_id is None


# ============================================================================
# Tests for admin-only access on /key/block, /key/unblock, /key/update max_budget
# ============================================================================


def _setup_block_unblock_mocks(monkeypatch, mock_key_team_id=None):
    """Helper to set up common mocks for block/unblock tests."""
    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    test_hashed_token = (
        "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
    )

    mock_key_record = MagicMock()
    mock_key_record.token = test_hashed_token
    mock_key_record.blocked = False
    mock_key_record.team_id = mock_key_team_id
    mock_key_record.model_dump_json.return_value = (
        f'{{"token": "{test_hashed_token}", "blocked": false}}'
    )

    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_key_record
    )
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=mock_key_record
    )

    mock_key_object = MagicMock()
    mock_key_object.blocked = True

    def mock_hash_token(token):
        if token.startswith("sk-"):
            return test_hashed_token
        return token

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
    monkeypatch.setattr("litellm.store_audit_logs", False)

    async def mock_delete_cache_key_object(**kwargs):
        pass

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
        mock_delete_cache_key_object,
    )

    return mock_prisma_client, test_hashed_token


@pytest.mark.asyncio
async def test_block_key_rejected_for_internal_user(monkeypatch):
    """Internal users should not be able to block keys."""
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import block_key

    _setup_block_unblock_mocks(monkeypatch)

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="internal_user",
    )

    with pytest.raises(HTTPException) as exc:
        await block_key(
            data=BlockKeyRequest(key="sk-test123456789"),
            http_request=mock_request,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc.value.status_code == 403
    assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail)


@pytest.mark.asyncio
async def test_unblock_key_rejected_for_internal_user(monkeypatch):
    """Internal users should not be able to unblock keys."""
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key

    _setup_block_unblock_mocks(monkeypatch)

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="internal_user",
    )

    with pytest.raises(HTTPException) as exc:
        await unblock_key(
            data=BlockKeyRequest(key="sk-test123456789"),
            http_request=mock_request,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert exc.value.status_code == 403
    assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail)


@pytest.mark.asyncio
async def test_block_key_allowed_for_proxy_admin(monkeypatch):
    """Proxy admins should be able to block keys."""
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import block_key

    _setup_block_unblock_mocks(monkeypatch)

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-admin",
        user_id="admin_user",
    )

    result = await block_key(
        data=BlockKeyRequest(key="sk-test123456789"),
        http_request=mock_request,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by=None,
    )
    assert result is not None


@pytest.mark.asyncio
async def test_block_key_allowed_for_team_admin(monkeypatch):
    """Team admins should be able to block keys belonging to their team."""
    from litellm.proxy._types import BlockKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import block_key

    team_id = "team-123"
    _setup_block_unblock_mocks(monkeypatch, mock_key_team_id=team_id)

    # Mock get_team_object to return a team where the user is admin
    team_obj = LiteLLM_TeamTableCachedObj(
        team_id=team_id,
        members_with_roles=[
            Member(user_id="team_admin_user", role="admin"),
        ],
    )

    async def mock_get_team_object(**kwargs):
        return team_obj

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )

    mock_request = MagicMock()
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-teamadmin",
        user_id="team_admin_user",
    )

    result = await block_key(
        data=BlockKeyRequest(key="sk-test123456789"),
        http_request=mock_request,
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by=None,
    )
    assert result is not None


@pytest.mark.asyncio
async def test_update_key_max_budget_rejected_for_internal_user(monkeypatch):
    """Internal users should not be able to modify max_budget on keys."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = AsyncMock()
    mock_proxy_logging_obj = MagicMock()

    test_hashed_token = (
        "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
    )

    # Mock existing key row
    mock_existing_key = MagicMock()
    mock_existing_key.token = test_hashed_token
    mock_existing_key.user_id = "internal_user"
    mock_existing_key.team_id = None
    mock_existing_key.project_id = None
    mock_existing_key.max_budget = 10.0
    mock_existing_key.models = []
    mock_existing_key.model_dump.return_value = {
        "token": test_hashed_token,
        "user_id": "internal_user",
        "team_id": None,
        "max_budget": 10.0,
    }

    mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_existing_key
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)

    mock_request = MagicMock()
    mock_request.query_params = {}
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="internal_user",
    )

    with pytest.raises(ProxyException) as exc:
        await update_key_fn(
            request=mock_request,
            data=UpdateKeyRequest(key=test_hashed_token, max_budget=999999),
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )

    assert str(exc.value.code) == "403"
    assert "Only proxy admins, team admins, or org admins" in str(exc.value.message)


@pytest.mark.asyncio
async def test_update_key_non_budget_fields_allowed_for_internal_user(monkeypatch):
    """Internal users should still be able to update non-budget fields on their own keys."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    mock_prisma_client = AsyncMock()
    mock_user_api_key_cache = AsyncMock()
    mock_proxy_logging_obj = MagicMock()

    test_hashed_token = (
        "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
    )

    # Mock existing key row
    mock_existing_key = MagicMock()
    mock_existing_key.token = test_hashed_token
    mock_existing_key.user_id = "internal_user"
    mock_existing_key.created_by = "internal_user"
    mock_existing_key.team_id = None
    mock_existing_key.project_id = None
    mock_existing_key.max_budget = 10.0
    mock_existing_key.key_alias = None
    mock_existing_key.models = []
    mock_existing_key.metadata = {}
    mock_existing_key.model_dump.return_value = {
        "token": test_hashed_token,
        "user_id": "internal_user",
        "team_id": None,
        "max_budget": 10.0,
    }

    mock_updated_key = MagicMock()
    mock_updated_key.token = test_hashed_token
    mock_updated_key.key_alias = "my-alias"

    mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
    mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key)
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_existing_key
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
    )
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
    monkeypatch.setattr("litellm.store_audit_logs", False)

    def mock_hash_token(token):
        return test_hashed_token

    monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)

    async def mock_delete_cache_key_object(**kwargs):
        pass

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
        mock_delete_cache_key_object,
    )

    # Mock _enforce_unique_key_alias to avoid DB call
    async def mock_enforce_unique_key_alias(**kwargs):
        pass

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias",
        mock_enforce_unique_key_alias,
    )

    mock_request = MagicMock()
    mock_request.query_params = {}
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="internal_user",
    )

    # Updating key_alias (non-budget field) should succeed
    result = await update_key_fn(
        request=mock_request,
        data=UpdateKeyRequest(key=test_hashed_token, key_alias="my-alias"),
        user_api_key_dict=user_api_key_dict,
        litellm_changed_by=None,
    )

    assert result is not None


@pytest.mark.asyncio
async def test_update_key_non_budget_rejects_cross_user_modification(monkeypatch):
    """Regression: previously _check_key_admin_access was gated on
    max_budget/spend changes only, so an internal user could rewrite any
    OTHER field (alias, models, tpm_limit, blocked, metadata, …) on any
    key they weren't admin of as long as they avoided budget/spend. This
    confirms that a non-admin user updating a key that belongs to another
    user fails with 403 even for non-budget fields."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    mock_prisma_client = AsyncMock()
    test_hashed_token = "cafebabe" * 8

    mock_existing_key = MagicMock()
    mock_existing_key.token = test_hashed_token
    mock_existing_key.user_id = "victim_user"  # owned by someone else
    mock_existing_key.team_id = None
    mock_existing_key.project_id = None
    mock_existing_key.max_budget = 10.0
    mock_existing_key.key_alias = "original"
    mock_existing_key.models = []
    mock_existing_key.model_dump.return_value = {
        "token": test_hashed_token,
        "user_id": "victim_user",
        "team_id": None,
        "max_budget": 10.0,
    }

    mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_existing_key
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token
    )

    mock_request = MagicMock()
    mock_request.query_params = {}
    attacker = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-attacker",
        user_id="attacker_user",  # NOT the owner
    )

    # Trying to blanket-rewrite a non-budget field on someone else's key
    # must now fail.
    with pytest.raises(ProxyException) as exc:
        await update_key_fn(
            request=mock_request,
            data=UpdateKeyRequest(
                key=test_hashed_token, key_alias="pwned", blocked=True
            ),
            user_api_key_dict=attacker,
            litellm_changed_by=None,
        )
    assert str(exc.value.code) == "403"


@pytest.mark.asyncio
async def test_update_key_creator_reassigned_key_blocked(monkeypatch):
    """Regression: creator who no longer owns the key (user_id ≠ caller) must
    not bypass _check_key_admin_access via the caller_is_creator shortcut."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    test_hashed_token = "aabbccdd" * 8
    mock_prisma_client = AsyncMock()

    mock_existing_key = MagicMock()
    mock_existing_key.token = test_hashed_token
    mock_existing_key.created_by = "demoted-admin"  # creator
    mock_existing_key.user_id = "victim-user"  # reassigned to someone else
    mock_existing_key.team_id = None
    mock_existing_key.project_id = None
    mock_existing_key.max_budget = 10.0
    mock_existing_key.key_alias = "original"
    mock_existing_key.models = []
    mock_existing_key.model_dump.return_value = {
        "token": test_hashed_token,
        "user_id": "victim-user",
        "created_by": "demoted-admin",
        "team_id": None,
        "max_budget": 10.0,
    }

    mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_existing_key
    )

    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
    monkeypatch.setattr("litellm.store_audit_logs", False)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token
    )

    demoted_admin = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-demoted",
        user_id="demoted-admin",
    )

    mock_request = MagicMock()
    mock_request.query_params = {}

    with pytest.raises(Exception) as exc:
        await update_key_fn(
            request=mock_request,
            data=UpdateKeyRequest(key=test_hashed_token, key_alias="hijacked"),
            user_api_key_dict=demoted_admin,
            litellm_changed_by=None,
        )
    assert str(exc.value.code) == "403"


@pytest.mark.asyncio
async def test_update_key_team_member_with_permission_can_update_non_budget(
    monkeypatch,
):
    """A team member whose team grants /key/update in member_permissions can
    update non-budget fields on a team key even though they are not a team
    admin. Regression: the cross-key admin check was over-broad and rejected
    this documented path."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    test_hashed_token = "deadbeef" * 8
    team_id = "team-with-update-grant"
    member_user_id = "team-member-user"

    mock_existing_key = MagicMock()
    mock_existing_key.token = test_hashed_token
    mock_existing_key.user_id = None  # team-scoped key (no owning user)
    mock_existing_key.team_id = team_id
    mock_existing_key.project_id = None
    mock_existing_key.max_budget = 10.0
    mock_existing_key.key_alias = "original"
    mock_existing_key.models = []
    mock_existing_key.model_dump.return_value = {
        "token": test_hashed_token,
        "user_id": None,
        "team_id": team_id,
        "max_budget": 10.0,
    }

    team_table = LiteLLM_TeamTableCachedObj(
        team_id=team_id,
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="some-team-admin", role="admin"),
            Member(user_id=member_user_id, role="user"),
        ],
        team_member_permissions=["/key/update", "/key/info"],
    )

    mock_updated_key = MagicMock()
    mock_updated_key.token = test_hashed_token
    mock_updated_key.key_alias = "renamed-by-member"

    mock_prisma_client = AsyncMock()
    mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
    mock_prisma_client.update_data = AsyncMock(return_value=mock_updated_key)
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_existing_key
    )

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    async def mock_enforce_unique_key_alias(**kwargs):
        pass

    async def mock_delete_cache_key_object(**kwargs):
        pass

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )
    monkeypatch.setattr(
        "litellm.proxy.management_helpers.team_member_permission_checks.get_team_object",
        mock_get_team_object,
    )
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias",
        mock_enforce_unique_key_alias,
    )
    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
        mock_delete_cache_key_object,
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
    monkeypatch.setattr("litellm.store_audit_logs", False)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token
    )

    mock_request = MagicMock()
    mock_request.query_params = {}
    team_member = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-team-member",
        user_id=member_user_id,
        team_id=team_id,
    )

    # Non-budget update on a team key by a team member with /key/update
    # permission should succeed.
    result = await update_key_fn(
        request=mock_request,
        data=UpdateKeyRequest(key=test_hashed_token, key_alias="renamed-by-member"),
        user_api_key_dict=team_member,
        litellm_changed_by=None,
    )

    assert result is not None


@pytest.mark.asyncio
async def test_update_key_team_member_cannot_change_budget(monkeypatch):
    """A team member with /key/update in member_permissions still cannot
    change max_budget — budget/spend changes require team/org admin. The
    member_permissions bypass only applies to non-budget fields."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        update_key_fn,
    )

    test_hashed_token = "feedface" * 8
    team_id = "team-with-update-grant"
    member_user_id = "team-member-user"

    mock_existing_key = MagicMock()
    mock_existing_key.token = test_hashed_token
    mock_existing_key.user_id = None  # team-scoped key (no owning user)
    mock_existing_key.team_id = team_id
    mock_existing_key.project_id = None
    mock_existing_key.max_budget = 10.0
    mock_existing_key.key_alias = "original"
    mock_existing_key.models = []
    mock_existing_key.model_dump.return_value = {
        "token": test_hashed_token,
        "user_id": None,
        "team_id": team_id,
        "max_budget": 10.0,
    }

    team_table = LiteLLM_TeamTableCachedObj(
        team_id=team_id,
        team_alias="test-team",
        tpm_limit=None,
        rpm_limit=None,
        max_budget=None,
        spend=0.0,
        models=[],
        blocked=False,
        members_with_roles=[
            Member(user_id="some-team-admin", role="admin"),
            Member(user_id=member_user_id, role="user"),
        ],
        team_member_permissions=["/key/update", "/key/info"],
    )

    mock_prisma_client = AsyncMock()
    mock_prisma_client.get_data = AsyncMock(return_value=mock_existing_key)
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=mock_existing_key
    )

    async def mock_get_team_object(*args, **kwargs):
        return team_table

    monkeypatch.setattr(
        "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
        mock_get_team_object,
    )
    monkeypatch.setattr(
        "litellm.proxy.management_helpers.team_member_permission_checks.get_team_object",
        mock_get_team_object,
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
    monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.hash_token", lambda t: test_hashed_token
    )

    mock_request = MagicMock()
    mock_request.query_params = {}
    team_member = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-team-member",
        user_id=member_user_id,
        team_id=team_id,
    )

    with pytest.raises(ProxyException) as exc:
        await update_key_fn(
            request=mock_request,
            data=UpdateKeyRequest(key=test_hashed_token, max_budget=500.0),
            user_api_key_dict=team_member,
            litellm_changed_by=None,
        )
    assert str(exc.value.code) == "403"


# ============================================================================
# LIT-1884: Internal users cannot create invalid keys
# ============================================================================


class TestLIT1884KeyGenerateValidation:
    """Tests for LIT-1884: internal users should not be able to generate invalid keys."""

    @pytest.mark.asyncio
    async def test_internal_user_generate_key_no_user_id_auto_assigns(self):
        """
        When an internal_user calls /key/generate without user_id,
        the caller's user_id should be auto-assigned before reaching
        _common_key_generation_helper.
        """
        mock_prisma_client = AsyncMock()

        data = GenerateKeyRequest(key_alias="test-alias")
        assert data.user_id is None

        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )

        # Patch _common_key_generation_helper to avoid needing full DB mocks.
        # We just want to verify user_id is set before we reach this point.
        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
                new_callable=AsyncMock,
                return_value=MagicMock(),
            ),
        ):
            await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )

        # The data object should have been mutated to include the caller's user_id
        assert data.user_id == "internal-user-123"

    @pytest.mark.asyncio
    async def test_internal_user_generate_key_invalid_team_id_rejected(self):
        """
        When an internal_user provides a non-existent team_id,
        key/generate should raise ProxyException with status 400.
        """
        mock_prisma_client = AsyncMock()

        data = GenerateKeyRequest(
            key_alias="test-alias",
            team_id="nonexistent-team-id",
        )

        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
                AsyncMock(side_effect=Exception("Team not found")),
            ),
        ):
            with pytest.raises(ProxyException) as exc_info:
                await generate_key_fn(
                    data=data,
                    user_api_key_dict=user_api_key_dict,
                    litellm_changed_by=None,
                )
            assert str(exc_info.value.code) == "400"
            assert "Team not found" in str(exc_info.value.message)

    @pytest.mark.asyncio
    async def test_admin_generate_key_invalid_team_id_allowed(self):
        """
        Admin callers should be allowed to create keys with any team_id,
        even if the team doesn't exist (team_table=None is OK for admins).
        """
        data = GenerateKeyRequest(
            key_alias="admin-key",
            team_id="nonexistent-team-id",
            user_id="admin-user",
        )

        user_api_key_dict = UserAPIKeyAuth(
            user_id="admin-user",
            user_role=LitellmUserRoles.PROXY_ADMIN,
        )

        mock_prisma_client = AsyncMock()

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
                AsyncMock(side_effect=Exception("Team not found")),
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
                new_callable=AsyncMock,
                return_value=MagicMock(),
            ),
        ):
            # Should NOT raise — admin bypasses team validation
            result = await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )
            assert result is not None

    @pytest.mark.asyncio
    async def test_admin_generate_key_no_user_id_not_auto_assigned(self):
        """
        Admin callers should NOT have user_id auto-assigned — they may
        intentionally create keys without a user_id.
        """
        data = GenerateKeyRequest(key_alias="admin-unbound-key")
        assert data.user_id is None

        user_api_key_dict = UserAPIKeyAuth(
            user_id="admin-user",
            user_role=LitellmUserRoles.PROXY_ADMIN,
        )

        mock_prisma_client = AsyncMock()

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
                new_callable=AsyncMock,
                return_value=MagicMock(),
            ),
        ):
            await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )

        # user_id should remain None for admin
        assert data.user_id is None

    def test_key_generation_check_non_admin_no_team_table_raises(self):
        """
        key_generation_check should raise 400 for non-admin when team_table is None
        and key_generation_settings is not set.
        """
        data = GenerateKeyRequest(team_id="some-team-id")
        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )

        with patch.object(litellm, "key_generation_settings", None):
            with pytest.raises(HTTPException) as exc_info:
                key_generation_check(
                    team_table=None,
                    user_api_key_dict=user_api_key_dict,
                    data=data,
                    route="key_generate",
                )
            assert exc_info.value.status_code == 400
            assert "Unable to find team object" in str(exc_info.value.detail)

    def test_key_generation_check_admin_no_team_table_allowed(self):
        """
        key_generation_check should allow admin to proceed even when team_table is None.
        """
        data = GenerateKeyRequest(team_id="some-team-id")
        user_api_key_dict = UserAPIKeyAuth(
            user_id="admin-user",
            user_role=LitellmUserRoles.PROXY_ADMIN,
        )

        with patch.object(litellm, "key_generation_settings", None):
            result = key_generation_check(
                team_table=None,
                user_api_key_dict=user_api_key_dict,
                data=data,
                route="key_generate",
            )
            assert result is True


class TestLIT1884KeyUpdateValidation:
    """Tests for LIT-1884: internal users should not be able to update keys to remove user_id or set invalid team."""

    @pytest.mark.asyncio
    async def test_internal_user_cannot_remove_user_id(self):
        """
        Non-admin users should not be able to set user_id to empty string (remove it).
        """
        data = UpdateKeyRequest(key="sk-test-key", user_id="")
        existing_key_row = MagicMock()
        existing_key_row.user_id = "internal-user-123"
        existing_key_row.token = "hashed_token"
        existing_key_row.team_id = None

        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )

        with pytest.raises(HTTPException) as exc_info:
            await _validate_update_key_data(
                data=data,
                existing_key_row=existing_key_row,
                user_api_key_dict=user_api_key_dict,
                llm_router=None,
                premium_user=False,
                prisma_client=AsyncMock(),
                user_api_key_cache=MagicMock(),
            )
        assert exc_info.value.status_code == 403
        assert "cannot remove the user_id" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_internal_user_cannot_set_invalid_team_id(self):
        """
        Non-admin users should not be able to update a key to a non-existent team.
        get_team_object raises HTTPException(404) when team doesn't exist in DB.
        """
        data = UpdateKeyRequest(key="sk-test-key", team_id="nonexistent-team")
        existing_key_row = MagicMock()
        existing_key_row.user_id = "internal-user-123"
        existing_key_row.token = "hashed_token"
        existing_key_row.team_id = None
        existing_key_row.organization_id = None
        existing_key_row.project_id = None

        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )

        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
            AsyncMock(
                side_effect=HTTPException(
                    status_code=404,
                    detail="Team doesn't exist in db. Team=nonexistent-team.",
                )
            ),
        ):
            with pytest.raises(HTTPException) as exc_info:
                await _validate_update_key_data(
                    data=data,
                    existing_key_row=existing_key_row,
                    user_api_key_dict=user_api_key_dict,
                    llm_router=None,
                    premium_user=False,
                    prisma_client=AsyncMock(),
                    user_api_key_cache=MagicMock(),
                )
            assert exc_info.value.status_code == 404
            assert "Team doesn't exist" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_admin_can_remove_user_id(self):
        """
        Admin users should be allowed to set user_id to empty string.
        """
        data = UpdateKeyRequest(key="sk-test-key", user_id="")
        existing_key_row = MagicMock()
        existing_key_row.user_id = "some-user"
        existing_key_row.token = "hashed_token"
        existing_key_row.team_id = None
        existing_key_row.organization_id = None
        existing_key_row.project_id = None

        user_api_key_dict = UserAPIKeyAuth(
            user_id="admin-user",
            user_role=LitellmUserRoles.PROXY_ADMIN,
        )

        mock_prisma_client = AsyncMock()

        # Should NOT raise
        await _validate_update_key_data(
            data=data,
            existing_key_row=existing_key_row,
            user_api_key_dict=user_api_key_dict,
            llm_router=None,
            premium_user=False,
            prisma_client=mock_prisma_client,
            user_api_key_cache=MagicMock(),
        )


class TestKeyOwnerPrivilegeEscalation:
    """
    Policy:
    - created_by == caller → can edit any non-budget field without admin
    - created_by != caller (assigned user) → must pass admin check for any edit
    - budget changes (max_budget/spend) → always require admin
    - PROXY_ADMIN → unrestricted
    """

    def _make_existing_key(self, user_id="creator-123", created_by="creator-123"):
        row = MagicMock()
        # user_id must be set explicitly — MagicMock auto-attrs are not None and
        # trip the _is_allowed_to_make_key_request assert before our check runs.
        row.user_id = user_id
        row.created_by = created_by
        row.token = "hashed_token"
        row.team_id = None
        row.max_budget = None
        row.spend = 0.0
        row.organization_id = None
        row.project_id = None
        return row

    def _make_auth(self, user_id="creator-123"):
        return UserAPIKeyAuth(
            user_id=user_id,
            user_role=LitellmUserRoles.INTERNAL_USER,
        )

    @pytest.mark.asyncio
    async def test_assigned_user_blocked_from_any_edit(self):
        """User who did not create the key cannot edit it at all."""
        data = UpdateKeyRequest(key="sk-test", key_alias="hacked")
        # user_id matches caller so _is_allowed_to_make_key_request passes,
        # but created_by != caller so our creator check requires admin.
        existing = self._make_existing_key(
            user_id="assigned-user", created_by="admin-456"
        )
        auth = self._make_auth(user_id="assigned-user")

        mock_check = AsyncMock(
            side_effect=HTTPException(status_code=403, detail="Not authorized")
        )
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access",
            mock_check,
        ):
            with pytest.raises(HTTPException) as exc_info:
                await _validate_update_key_data(
                    data=data,
                    existing_key_row=existing,
                    user_api_key_dict=auth,
                    llm_router=None,
                    premium_user=False,
                    prisma_client=AsyncMock(),
                    user_api_key_cache=MagicMock(),
                )
        assert exc_info.value.status_code == 403
        mock_check.assert_called_once()

    @pytest.mark.asyncio
    async def test_assigned_user_blocked_from_model_escalation(self):
        data = UpdateKeyRequest(key="sk-test", models=["gpt-4", "claude-opus"])
        existing = self._make_existing_key(
            user_id="assigned-user", created_by="admin-456"
        )
        auth = self._make_auth(user_id="assigned-user")

        mock_check = AsyncMock(
            side_effect=HTTPException(status_code=403, detail="Not authorized")
        )
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access",
            mock_check,
        ):
            with pytest.raises(HTTPException):
                await _validate_update_key_data(
                    data=data,
                    existing_key_row=existing,
                    user_api_key_dict=auth,
                    llm_router=None,
                    premium_user=False,
                    prisma_client=AsyncMock(),
                    user_api_key_cache=MagicMock(),
                )
        mock_check.assert_called_once()

    @pytest.mark.asyncio
    async def test_creator_can_edit_own_key(self):
        """Key creator can update any non-budget field without admin."""
        data = UpdateKeyRequest(
            key="sk-test", models=["gpt-4"], rpm_limit=500, key_alias="my-key"
        )
        existing = self._make_existing_key(created_by="creator-123")
        auth = self._make_auth(user_id="creator-123")

        mock_check = AsyncMock()
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access",
            mock_check,
        ):
            await _validate_update_key_data(
                data=data,
                existing_key_row=existing,
                user_api_key_dict=auth,
                llm_router=None,
                premium_user=False,
                prisma_client=AsyncMock(),
                user_api_key_cache=MagicMock(),
            )
        mock_check.assert_not_called()

    @pytest.mark.asyncio
    async def test_creator_cannot_change_own_budget(self):
        """Budget changes require admin even for the key creator."""
        data = UpdateKeyRequest(key="sk-test", max_budget=9999.0)
        existing = self._make_existing_key(created_by="creator-123")
        existing.max_budget = 10.0
        auth = self._make_auth(user_id="creator-123")

        mock_check = AsyncMock(
            side_effect=HTTPException(status_code=403, detail="Not authorized")
        )
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access",
            mock_check,
        ):
            with pytest.raises(HTTPException):
                await _validate_update_key_data(
                    data=data,
                    existing_key_row=existing,
                    user_api_key_dict=auth,
                    llm_router=None,
                    premium_user=False,
                    prisma_client=AsyncMock(),
                    user_api_key_cache=MagicMock(),
                )
        mock_check.assert_called_once()

    @pytest.mark.asyncio
    async def test_admin_can_update_any_field(self):
        data = UpdateKeyRequest(key="sk-test", models=["gpt-4"], max_budget=999.0)
        existing = self._make_existing_key(created_by="someone-else")
        existing.max_budget = 1.0
        auth = UserAPIKeyAuth(
            user_id="admin-user",
            user_role=LitellmUserRoles.PROXY_ADMIN,
        )

        mock_check = AsyncMock()
        with patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access",
            mock_check,
        ):
            await _validate_update_key_data(
                data=data,
                existing_key_row=existing,
                user_api_key_dict=auth,
                llm_router=None,
                premium_user=False,
                prisma_client=AsyncMock(),
                user_api_key_cache=MagicMock(),
            )
        mock_check.assert_not_called()


class TestKeyAliasSkipValidationOnUnchanged:
    """
    Test that updating/regenerating a key without changing its key_alias
    does NOT re-validate the alias. This prevents legacy aliases (created
    before stricter validation rules) from blocking edits to other fields.
    """

    @pytest.fixture(autouse=True)
    def enable_validation(self):
        litellm.enable_key_alias_format_validation = True
        yield
        litellm.enable_key_alias_format_validation = False

    @pytest.fixture
    def mock_prisma(self):
        prisma = MagicMock()
        prisma.db = MagicMock()
        prisma.db.litellm_verificationtoken = MagicMock()
        prisma.get_data = AsyncMock(return_value=None)  # no duplicate alias
        prisma.update_data = AsyncMock(return_value=None)
        prisma.jsonify_object = MagicMock(side_effect=lambda data: data)
        return prisma

    @pytest.fixture
    def existing_key_with_legacy_alias(self):
        """A key whose alias contains '@' — valid now, but simulates a legacy alias."""
        return LiteLLM_VerificationToken(
            token="hashed_token_123",
            key_alias="user@domain.com",
            team_id="team-1",
            models=[],
            max_budget=100.0,
        )

    @pytest.mark.asyncio
    async def test_update_key_unchanged_legacy_alias_passes(
        self, mock_prisma, existing_key_with_legacy_alias
    ):
        """
        Updating a key without changing its key_alias should skip format
        validation — even if the alias wouldn't pass current rules.
        """
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )

        # Temporarily make the regex reject '@' to simulate stricter rules
        import re
        from litellm.proxy.management_endpoints import key_management_endpoints as mod

        original_pattern = mod._KEY_ALIAS_PATTERN
        mod._KEY_ALIAS_PATTERN = re.compile(
            r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.]{0,253}[a-zA-Z0-9]$"
        )
        try:
            # Confirm the alias WOULD fail validation directly
            with pytest.raises(ProxyException):
                _validate_key_alias_format("user@domain.com")

            # But prepare_key_update_data + the skip logic should allow it
            # Simulate what update_key_fn does: alias is in non_default_values
            # but matches existing_key_row.key_alias => skip validation
            existing_alias = existing_key_with_legacy_alias.key_alias
            new_alias = "user@domain.com"  # same as existing
            assert new_alias == existing_alias  # unchanged

            # This is the core logic from update_key_fn:
            if new_alias != existing_alias:
                _validate_key_alias_format(new_alias)
            # No exception raised — test passes
        finally:
            mod._KEY_ALIAS_PATTERN = original_pattern

    @pytest.mark.asyncio
    async def test_update_key_changed_alias_still_validated(
        self, mock_prisma, existing_key_with_legacy_alias
    ):
        """
        When the alias IS being changed, validation should still run.
        """
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )

        existing_alias = existing_key_with_legacy_alias.key_alias
        new_alias = "!invalid!"

        assert new_alias != existing_alias
        with pytest.raises(ProxyException):
            if new_alias != existing_alias:
                _validate_key_alias_format(new_alias)

    @pytest.mark.asyncio
    async def test_update_key_changed_to_valid_alias_passes(
        self, mock_prisma, existing_key_with_legacy_alias
    ):
        """
        Changing the alias to a new valid value should pass validation.
        """
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )

        existing_alias = existing_key_with_legacy_alias.key_alias
        new_alias = "new-valid-alias"

        assert new_alias != existing_alias
        # Should not raise
        if new_alias != existing_alias:
            _validate_key_alias_format(new_alias)

    @pytest.mark.asyncio
    async def test_update_key_alias_none_skips_validation(self):
        """
        When key_alias is not in the update payload (None), validation
        should be skipped regardless.
        """
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            _validate_key_alias_format,
        )

        # None alias should always pass
        _validate_key_alias_format(None)


# --- Tests: _enforce_upperbound_key_params ---


def test_enforce_upperbound_rejects_over_limit_on_generate():
    """Test that key generation is rejected when values exceed upperbound."""
    import litellm
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            tpm_limit=1000, rpm_limit=100, max_budget=10.0
        )
        data = GenerateKeyRequest(tpm_limit=5000)
        with pytest.raises(HTTPException) as exc_info:
            _enforce_upperbound_key_params(data, fill_defaults=True)
        assert exc_info.value.status_code == 400
        assert "tpm_limit" in str(exc_info.value.detail)
    finally:
        litellm.upperbound_key_generate_params = original


def test_enforce_upperbound_fills_defaults_on_generate():
    """Test that None values are filled with upperbound defaults during generation."""
    import litellm
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            tpm_limit=1000, rpm_limit=100
        )
        data = GenerateKeyRequest()  # tpm_limit=None, rpm_limit=None
        _enforce_upperbound_key_params(data, fill_defaults=True)
        assert data.tpm_limit == 1000
        assert data.rpm_limit == 100
    finally:
        litellm.upperbound_key_generate_params = original


def test_enforce_upperbound_skips_none_on_update():
    """Test that None values are NOT filled during update (fill_defaults=False)."""
    import litellm
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            tpm_limit=1000, rpm_limit=100
        )
        data = UpdateKeyRequest(key="sk-test")  # tpm_limit=None, rpm_limit=None
        _enforce_upperbound_key_params(data, fill_defaults=False)
        assert data.tpm_limit is None  # should NOT be filled
        assert data.rpm_limit is None  # should NOT be filled
    finally:
        litellm.upperbound_key_generate_params = original


def test_enforce_upperbound_rejects_over_limit_on_update():
    """Test that key update is rejected when values exceed upperbound."""
    import litellm
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            tpm_limit=1000, rpm_limit=100, max_budget=10.0
        )
        data = UpdateKeyRequest(key="sk-test", tpm_limit=5000)
        with pytest.raises(HTTPException) as exc_info:
            _enforce_upperbound_key_params(data, fill_defaults=False)
        assert exc_info.value.status_code == 400
        assert "tpm_limit" in str(exc_info.value.detail)
    finally:
        litellm.upperbound_key_generate_params = original


def test_enforce_upperbound_allows_within_limit_on_update():
    """Test that key update passes when values are within upperbound."""
    import litellm
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            tpm_limit=1000, rpm_limit=100, max_budget=10.0
        )
        data = UpdateKeyRequest(
            key="sk-test", tpm_limit=500, rpm_limit=50, max_budget=5.0
        )
        _enforce_upperbound_key_params(data, fill_defaults=False)
        # Should not raise
        assert data.tpm_limit == 500
        assert data.rpm_limit == 50
        assert data.max_budget == 5.0
    finally:
        litellm.upperbound_key_generate_params = original


def test_enforce_upperbound_duration_over_limit():
    """Test that duration exceeding upperbound is rejected."""
    import litellm
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            duration="7d"
        )
        data = UpdateKeyRequest(key="sk-test", duration="30d")
        with pytest.raises(HTTPException) as exc_info:
            _enforce_upperbound_key_params(data, fill_defaults=False)
        assert exc_info.value.status_code == 400
        assert "duration" in str(exc_info.value.detail)
    finally:
        litellm.upperbound_key_generate_params = original


def test_enforce_upperbound_no_config_is_noop():
    """Test that no enforcement happens when upperbound params are not configured."""
    import litellm

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = None
        data = UpdateKeyRequest(key="sk-test", tpm_limit=999999)
        _enforce_upperbound_key_params(data, fill_defaults=False)
        # Should not raise — no enforcement configured
        assert data.tpm_limit == 999999
    finally:
        litellm.upperbound_key_generate_params = original


# --- Tests: _execute_virtual_key_regeneration enforces upperbound ---


def _make_regenerate_mock_prisma():
    """Mock prisma client shaped for _execute_virtual_key_regeneration."""

    class DictLikeResult:
        def __init__(self, data):
            self._data = data

        def __iter__(self):
            return iter(self._data.items())

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=DictLikeResult(
            {
                "token": "new-hashed-token",
                "key_name": "sk-...ab12",
                "user_id": "user-1",
            }
        )
    )
    mock_prisma_client.db.litellm_verificationtoken.create = AsyncMock(
        return_value=None
    )
    mock_prisma_client.jsonify_object = MagicMock(side_effect=lambda data: data)
    return mock_prisma_client


def _make_regenerate_user_api_key_dict():
    return UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-admin",
        user_id="admin-user",
    )


def _make_regenerate_existing_key():
    return LiteLLM_VerificationToken(
        token="abc123",
        user_id="user-1",
        models=["gpt-4"],
        team_id=None,
        max_budget=None,
        tags=None,
    )


@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_rejects_over_limit_duration():
    """Regenerate must reject durations exceeding upperbound_key_generate_params.duration."""
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            duration="1h"
        )
        existing_key = _make_regenerate_existing_key()
        data = RegenerateKeyRequest(duration="2h")
        user_api_key_dict = _make_regenerate_user_api_key_dict()
        mock_prisma_client = _make_regenerate_mock_prisma()

        with (
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
                new_callable=AsyncMock,
                return_value="sk-newtoken1234ab12",
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
                new_callable=AsyncMock,
            ),
        ):
            with pytest.raises(HTTPException) as exc_info:
                await _execute_virtual_key_regeneration(
                    prisma_client=mock_prisma_client,
                    key_in_db=existing_key,
                    hashed_api_key="abc123",
                    key="abc123",
                    data=data,
                    user_api_key_dict=user_api_key_dict,
                    litellm_changed_by=None,
                    user_api_key_cache=MagicMock(),
                    proxy_logging_obj=MagicMock(),
                )
        assert exc_info.value.status_code == 400
        assert "duration" in str(exc_info.value.detail)
        # Rejected regenerate must not reach the DB update.
        assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0
    finally:
        litellm.upperbound_key_generate_params = original


@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_allows_within_limit_duration():
    """Regenerate must accept durations within upperbound_key_generate_params.duration."""
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            duration="1h"
        )
        existing_key = _make_regenerate_existing_key()
        data = RegenerateKeyRequest(duration="30m")
        user_api_key_dict = _make_regenerate_user_api_key_dict()
        mock_prisma_client = _make_regenerate_mock_prisma()

        with (
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
                new_callable=AsyncMock,
                return_value="sk-newtoken1234ab12",
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
                new_callable=AsyncMock,
            ),
        ):
            await _execute_virtual_key_regeneration(
                prisma_client=mock_prisma_client,
                key_in_db=existing_key,
                hashed_api_key="abc123",
                key="abc123",
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
                user_api_key_cache=MagicMock(),
                proxy_logging_obj=MagicMock(),
            )
        assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
    finally:
        litellm.upperbound_key_generate_params = original


@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_rejects_over_limit_max_budget():
    """Regenerate must reject max_budget exceeding upperbound — proves the fix covers non-duration fields."""
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            max_budget=10.0
        )
        existing_key = _make_regenerate_existing_key()
        data = RegenerateKeyRequest(max_budget=500.0)
        user_api_key_dict = _make_regenerate_user_api_key_dict()
        mock_prisma_client = _make_regenerate_mock_prisma()

        with (
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
                new_callable=AsyncMock,
                return_value="sk-newtoken1234ab12",
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
                new_callable=AsyncMock,
            ),
        ):
            with pytest.raises(HTTPException) as exc_info:
                await _execute_virtual_key_regeneration(
                    prisma_client=mock_prisma_client,
                    key_in_db=existing_key,
                    hashed_api_key="abc123",
                    key="abc123",
                    data=data,
                    user_api_key_dict=user_api_key_dict,
                    litellm_changed_by=None,
                    user_api_key_cache=MagicMock(),
                    proxy_logging_obj=MagicMock(),
                )
        assert exc_info.value.status_code == 400
        assert "max_budget" in str(exc_info.value.detail)
        assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0
    finally:
        litellm.upperbound_key_generate_params = original


@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_skips_none_values():
    """Regenerate with data.duration=None must not raise, even when upperbound is set
    (fill_defaults=False semantic — None means 'inherit from existing key')."""
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )
    from litellm.types.proxy.management_endpoints.ui_sso import (
        LiteLLM_UpperboundKeyGenerateParams,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = LiteLLM_UpperboundKeyGenerateParams(
            duration="1h"
        )
        existing_key = _make_regenerate_existing_key()
        data = RegenerateKeyRequest()  # all fields None
        user_api_key_dict = _make_regenerate_user_api_key_dict()
        mock_prisma_client = _make_regenerate_mock_prisma()

        with (
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
                new_callable=AsyncMock,
                return_value="sk-newtoken1234ab12",
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
                new_callable=AsyncMock,
            ),
        ):
            await _execute_virtual_key_regeneration(
                prisma_client=mock_prisma_client,
                key_in_db=existing_key,
                hashed_api_key="abc123",
                key="abc123",
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
                user_api_key_cache=MagicMock(),
                proxy_logging_obj=MagicMock(),
            )
        assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
    finally:
        litellm.upperbound_key_generate_params = original


@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_no_upperbound_config_is_noop():
    """Regenerate with no upperbound config set must accept any duration."""
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )

    original = litellm.upperbound_key_generate_params
    try:
        litellm.upperbound_key_generate_params = None
        existing_key = _make_regenerate_existing_key()
        data = RegenerateKeyRequest(duration="30d")
        user_api_key_dict = _make_regenerate_user_api_key_dict()
        mock_prisma_client = _make_regenerate_mock_prisma()

        with (
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
                new_callable=AsyncMock,
                return_value="sk-newtoken1234ab12",
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
                new_callable=AsyncMock,
            ),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
                new_callable=AsyncMock,
            ),
        ):
            await _execute_virtual_key_regeneration(
                prisma_client=mock_prisma_client,
                key_in_db=existing_key,
                hashed_api_key="abc123",
                key="abc123",
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
                user_api_key_cache=MagicMock(),
                proxy_logging_obj=MagicMock(),
            )
        assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1
    finally:
        litellm.upperbound_key_generate_params = original


class TestAllowedRoutesCallerPermission:
    """
    Non-admins must not be able to set `allowed_routes` on a key. The field
    bypasses the role-based route gate in
    RouteChecks.non_proxy_admin_allowed_routes_check, so allowing a non-admin
    to populate it grants them arbitrary endpoint access.
    """

    @pytest.mark.asyncio
    async def test_non_admin_generate_key_with_allowed_routes_rejected(self):
        data = GenerateKeyRequest(
            key_alias="escalate",
            allowed_routes=["/*"],
        )
        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )
        mock_prisma_client = AsyncMock()

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
                new_callable=AsyncMock,
                return_value=MagicMock(),
            ),
        ):
            with pytest.raises(ProxyException) as exc_info:
                await generate_key_fn(
                    data=data,
                    user_api_key_dict=user_api_key_dict,
                    litellm_changed_by=None,
                )
        assert str(exc_info.value.code) == "403"
        assert "allowed_routes" in str(exc_info.value.message)

    @pytest.mark.asyncio
    async def test_admin_generate_key_with_allowed_routes_allowed(self):
        data = GenerateKeyRequest(
            key_alias="admin-key",
            allowed_routes=["/chat/completions"],
            user_id="admin-user",
        )
        user_api_key_dict = UserAPIKeyAuth(
            user_id="admin-user",
            user_role=LitellmUserRoles.PROXY_ADMIN,
        )
        mock_prisma_client = AsyncMock()
        stub_response = MagicMock()

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
                new_callable=AsyncMock,
                return_value=stub_response,
            ),
        ):
            result = await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )
        assert result is stub_response

    @pytest.mark.asyncio
    async def test_non_admin_generate_key_default_empty_allowed_routes_ok(self):
        """
        Regression guard: GenerateKeyRequest.allowed_routes defaults to [], so
        the helper must treat empty-list as "not set" or every non-admin key
        creation breaks.
        """
        data = GenerateKeyRequest(key_alias="plain-key")
        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )
        mock_prisma_client = AsyncMock()
        stub_response = MagicMock()

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
                new_callable=AsyncMock,
                return_value=stub_response,
            ),
        ):
            result = await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )
        assert result is stub_response

    @pytest.mark.asyncio
    async def test_non_admin_update_key_with_allowed_routes_rejected(self):
        from litellm.proxy.management_endpoints.key_management_endpoints import (
            update_key_fn,
        )

        data = UpdateKeyRequest(key="sk-test", allowed_routes=["/*"])
        user_api_key_dict = UserAPIKeyAuth(
            user_id="internal-user-123",
            user_role=LitellmUserRoles.INTERNAL_USER,
        )
        mock_prisma_client = AsyncMock()

        with (
            patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
            patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
            patch("litellm.proxy.proxy_server.user_custom_key_update", None),
            patch("litellm.proxy.proxy_server.llm_router", None),
            patch("litellm.proxy.proxy_server.premium_user", True),
            patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
            patch(
                "litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key",
                new_callable=AsyncMock,
                return_value=MagicMock(),
            ),
        ):
            with pytest.raises(ProxyException) as exc_info:
                await update_key_fn(
                    request=MagicMock(),
                    data=data,
                    user_api_key_dict=user_api_key_dict,
                    litellm_changed_by=None,
                )
        assert str(exc_info.value.code) == "403"
        assert "allowed_routes" in str(exc_info.value.message)


def test_jinja_prompt_manager_is_sandboxed():
    """
    PromptManager renders user-supplied templates via /prompts/test, so its
    jinja env must reject access to unsafe Python attributes like
    ``__class__`` and ``__mro__``.
    """
    from jinja2.exceptions import SecurityError

    from litellm.integrations.dotprompt.prompt_manager import PromptManager

    pm = PromptManager()
    template = pm.jinja_env.from_string("{{ ''.__class__.__mro__ }}")
    with pytest.raises(SecurityError):
        template.render()


def test_validate_public_image_url_rejects_local_paths():
    from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
        _validate_public_image_url,
    )

    for bad in ("/etc/passwd", "file:///etc/passwd", "../../etc/passwd"):
        with pytest.raises(HTTPException) as exc_info:
            _validate_public_image_url(bad, "logo_url")
        assert exc_info.value.status_code == 400


def test_validate_public_image_url_accepts_http_and_noop_empty():
    from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
        _validate_public_image_url,
    )

    _validate_public_image_url("https://example.com/logo.png", "logo_url")
    _validate_public_image_url("http://cdn.internal/logo.svg", "logo_url")
    _validate_public_image_url(None, "logo_url")
    _validate_public_image_url("", "logo_url")
    _validate_public_image_url("   ", "logo_url")


@pytest.mark.asyncio
async def test_process_single_key_update_cache_invalidation_with_token_hash():
    """
    _process_single_key_update must pass the token hash as-is (not
    double-hashed) to _delete_cache_key_object when the key is already a
    pre-hashed token ID rather than an sk- prefixed key.

    Without this, cache invalidation silently fails: the wrong cache entry
    is deleted while the stale entry (with outdated fields) persists and
    gets refreshed indefinitely by update_cache on every successful request.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _process_single_key_update,
    )

    token_hash = "abc123def456"

    existing_key = LiteLLM_VerificationToken(
        token=token_hash,
        user_id="user-1",
        models=["gpt-4"],
        team_id=None,
        max_budget=None,
        tags=None,
    )

    mock_prisma_client = AsyncMock()
    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
        return_value=existing_key
    )
    mock_updated = MagicMock()
    mock_updated.model_dump.return_value = {"max_budget": 100.0}
    mock_prisma_client.update_data = AsyncMock(return_value={"data": mock_updated})

    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()
    mock_llm_router = MagicMock()

    with (
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data",
            return_value={"max_budget": 100.0},
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
            return_value=None,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
            new_callable=AsyncMock,
        ) as mock_delete_cache,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook",
            new_callable=AsyncMock,
        ),
    ):
        update_key_request = UpdateKeyRequest(
            key=token_hash,
            max_budget=100.0,
        )

        user_api_key_dict = UserAPIKeyAuth(
            user_role=LitellmUserRoles.PROXY_ADMIN,
            api_key="sk-admin",
            user_id="admin-user",
        )

        await _process_single_key_update(
            update_key_request=update_key_request,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
            prisma_client=mock_prisma_client,
            user_api_key_cache=mock_user_api_key_cache,
            proxy_logging_obj=mock_proxy_logging_obj,
            llm_router=mock_llm_router,
        )

        mock_delete_cache.assert_called_once()
        call_kwargs = mock_delete_cache.call_args.kwargs
        # The token hash should be passed as-is, NOT double-hashed
        assert call_kwargs["hashed_token"] == token_hash


@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_hash():
    """
    _execute_virtual_key_regeneration must pass the token hash as-is (not
    double-hashed) to _delete_cache_key_object when the key is a
    pre-hashed token ID.
    """
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )

    token_hash = "abc123def456"

    existing_key = LiteLLM_VerificationToken(
        token=token_hash,
        user_id="user-1",
        models=["gpt-4"],
        team_id=None,
        max_budget=None,
        tags=None,
    )

    mock_prisma_client = AsyncMock()

    # _execute_virtual_key_regeneration calls dict(updated_token) which
    # needs the return value to be iterable as key-value pairs.
    class DictLikeResult:
        def __init__(self, data):
            self._data = data

        def __iter__(self):
            return iter(self._data.items())

    mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
        return_value=DictLikeResult(
            {"token": "new-hashed-token", "key_name": "sk-...ab12", "user_id": "user-1"}
        )
    )
    mock_prisma_client.db.litellm_verificationtoken.create = AsyncMock(
        return_value=None
    )
    mock_prisma_client.jsonify_object = MagicMock(side_effect=lambda data: data)

    mock_user_api_key_cache = MagicMock()
    mock_proxy_logging_obj = MagicMock()

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-admin",
        user_id="admin-user",
    )

    with (
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
            new_callable=AsyncMock,
            return_value="sk-newtoken1234ab12",
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
            new_callable=AsyncMock,
        ) as mock_delete_cache,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data",
            new_callable=AsyncMock,
            return_value={},
        ),
    ):
        await _execute_virtual_key_regeneration(
            prisma_client=mock_prisma_client,
            key_in_db=existing_key,
            hashed_api_key=token_hash,
            key=token_hash,
            data=None,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
            user_api_key_cache=mock_user_api_key_cache,
            proxy_logging_obj=mock_proxy_logging_obj,
        )

        mock_delete_cache.assert_called_once()
        call_kwargs = mock_delete_cache.call_args.kwargs
        # The token hash should be passed as-is, NOT double-hashed
        assert call_kwargs["hashed_token"] == token_hash


# ---------------------------------------------------------------------------
# /team/key/bulk_update tests
# ---------------------------------------------------------------------------


_BULK_PKG = "litellm.proxy.management_endpoints.key_management_endpoints"


def _make_team_key(token: str, team_id: str = "team-abc") -> LiteLLM_VerificationToken:
    return LiteLLM_VerificationToken(
        token=token,
        user_id="user-123",
        models=[],
        team_id=team_id,
        max_budget=None,
    )


def _admin() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin"
    )


def _internal_user() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-iu", user_id="iu"
    )


def _updated(payload):
    m = MagicMock()
    m.model_dump.return_value = payload
    return m


def _setup_team_keys_mocks(
    monkeypatch,
    *,
    find_many=None,
    find_unique=None,
    update_data=None,
    hash_identity=True,
):
    """Set up mocks for bulk_update_team_keys; returns mock_prisma."""
    mock_prisma = AsyncMock()
    mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[] if find_many is None else find_many
    )
    if find_unique is not None:
        mock_prisma.db.litellm_verificationtoken.find_unique = find_unique
    if update_data is not None:
        mock_prisma.update_data = update_data
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
    monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.user_custom_key_update", None)
    monkeypatch.setattr(
        f"{_BULK_PKG}.prepare_key_update_data",
        AsyncMock(return_value={"max_budget": 50.0}),
    )
    monkeypatch.setattr(f"{_BULK_PKG}._delete_cache_key_object", AsyncMock())
    monkeypatch.setattr(
        f"{_BULK_PKG}.KeyManagementEventHooks.async_key_updated_hook", AsyncMock()
    )
    monkeypatch.setattr(f"{_BULK_PKG}.get_team_object", AsyncMock(return_value=None))
    monkeypatch.setattr(f"{_BULK_PKG}._check_team_key_limits", AsyncMock())
    monkeypatch.setattr(
        f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
        AsyncMock(),
    )
    if hash_identity:
        # Tests use already-hashed tokens; the raw-sk regression opts out.
        monkeypatch.setattr(f"{_BULK_PKG}._hash_token_if_needed", lambda token: token)
    return mock_prisma


async def _call_as_admin(data):
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_team_keys,
    )

    return await bulk_update_team_keys(
        data=data, user_api_key_dict=_admin(), litellm_changed_by=None
    )


# ---- happy paths ----------------------------------------------------------


@pytest.mark.asyncio
async def test_bulk_update_team_keys_success_with_key_ids(monkeypatch):
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    keys = [_make_team_key("tok-a"), _make_team_key("tok-b")]
    find_unique = AsyncMock(side_effect=keys)
    mock = _setup_team_keys_mocks(
        monkeypatch,
        find_many=keys,
        find_unique=find_unique,
        update_data=AsyncMock(
            side_effect=[{"data": _updated({"max_budget": 50.0})}] * 2
        ),
    )

    response = await _call_as_admin(
        BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            key_ids=["tok-a", "tok-b"],
            update_fields=KeyUpdateFields(max_budget=50.0),
        )
    )

    assert len(response.successful_updates) == 2
    assert len(response.failed_updates) == 0
    where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
    assert where["team_id"] == "team-abc"
    assert where["token"] == {"in": ["tok-a", "tok-b"]}
    find_unique.assert_not_called()


@pytest.mark.asyncio
async def test_bulk_update_team_keys_success_all_keys_in_team(monkeypatch):
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    keys = [_make_team_key(f"tok-{i}") for i in range(3)]
    find_unique = AsyncMock(side_effect=keys)
    mock = _setup_team_keys_mocks(
        monkeypatch,
        find_many=keys,
        find_unique=find_unique,
        update_data=AsyncMock(
            side_effect=[{"data": _updated({"max_budget": 50.0})}] * 3
        ),
    )

    response = await _call_as_admin(
        BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            all_keys_in_team=True,
            update_fields=KeyUpdateFields(max_budget=50.0),
        )
    )

    assert len(response.successful_updates) == 3
    where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
    # `blocked` is Boolean? with no default → /key/generate writes NULL. Prisma's
    # NOT excludes NULLs, so the filter has to OR `false` with `null` explicitly.
    blocked_or, expires_or = where["AND"][0]["OR"], where["AND"][1]["OR"]
    assert {"blocked": False} in blocked_or and {"blocked": None} in blocked_or
    assert {"expires": None} in expires_or
    assert any(
        "gt" in c.get("expires", {})
        for c in expires_or
        if isinstance(c.get("expires"), dict)
    )
    find_unique.assert_not_called()


@pytest.mark.asyncio
async def test_bulk_update_team_keys_key_not_in_team(monkeypatch):
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    in_team = _make_team_key("tok-a")
    _setup_team_keys_mocks(
        monkeypatch,
        find_many=[in_team],
        find_unique=AsyncMock(return_value=in_team),
        update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
    )

    response = await _call_as_admin(
        BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            key_ids=["tok-a", "tok-foreign"],
            update_fields=KeyUpdateFields(max_budget=50.0),
        )
    )
    assert [u.key for u in response.successful_updates] == ["tok-a"]
    assert [u.key for u in response.failed_updates] == ["tok-foreign"]
    assert "not found in team" in response.failed_updates[0].failed_reason


# ---- error paths ----------------------------------------------------------


@pytest.mark.asyncio
async def test_bulk_update_team_keys_batch_size_cap(monkeypatch):
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    _setup_team_keys_mocks(
        monkeypatch,
        find_many=[_make_team_key(f"tok-{i}") for i in range(501)],
    )

    with pytest.raises(HTTPException) as exc:
        await _call_as_admin(
            BulkUpdateTeamKeysRequest(
                team_id="team-abc",
                all_keys_in_team=True,
                update_fields=KeyUpdateFields(max_budget=50.0),
            )
        )
    assert exc.value.status_code == 400
    assert "more than 500" in exc.value.detail["error"]


@pytest.mark.asyncio
async def test_bulk_update_team_keys_empty_team_returns_404(monkeypatch):
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    _setup_team_keys_mocks(monkeypatch, find_many=[])
    with pytest.raises(HTTPException) as exc:
        await _call_as_admin(
            BulkUpdateTeamKeysRequest(
                team_id="team-empty",
                all_keys_in_team=True,
                update_fields=KeyUpdateFields(max_budget=50.0),
            )
        )
    assert exc.value.status_code == 404


# ---- auth -----------------------------------------------------------------


@pytest.mark.asyncio
async def test_bulk_update_team_keys_team_member_with_permission(monkeypatch):
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_team_keys,
    )
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    key_a = _make_team_key("tok-a")
    _setup_team_keys_mocks(
        monkeypatch,
        find_many=[key_a],
        find_unique=AsyncMock(return_value=key_a),
        update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
    )
    auth_check = AsyncMock()
    monkeypatch.setattr(
        f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
        auth_check,
    )

    response = await bulk_update_team_keys(
        data=BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            all_keys_in_team=True,
            update_fields=KeyUpdateFields(max_budget=50.0),
        ),
        user_api_key_dict=_internal_user(),
        litellm_changed_by=None,
    )
    assert len(response.successful_updates) == 1
    # Upfront check + per-key check inside _process_single_key_update
    assert auth_check.await_count == 2


@pytest.mark.asyncio
async def test_bulk_update_team_keys_team_member_no_permission(monkeypatch):
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_team_keys,
    )
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
    monkeypatch.setattr(
        f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
        AsyncMock(
            side_effect=ProxyException(
                message="not in team",
                type="team_member_permission_error",
                param="/key/update",
                code=401,
            )
        ),
    )

    with pytest.raises(ProxyException):
        await bulk_update_team_keys(
            data=BulkUpdateTeamKeysRequest(
                team_id="team-abc",
                all_keys_in_team=True,
                update_fields=KeyUpdateFields(max_budget=1.0),
            ),
            user_api_key_dict=_internal_user(),
            litellm_changed_by=None,
        )
    mock.update_data.assert_not_called()


# ---- pydantic-layer validation -------------------------------------------


def test_bulk_update_team_keys_request_validation():
    """Allowlist (extra='forbid'), empty-payload rejection, and selection XOR."""
    from pydantic import ValidationError

    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    forbidden = [
        "key",
        "key_alias",
        "team_id",
        "allowed_routes",
        "allowed_passthrough_routes",
        "permissions",
        "object_permission",
        "access_group_ids",
        "user_id",
        "organization_id",
        "blocked",
        "key_type",
        "models",
        "config",
        "router_settings",
        "spend",
    ]
    for f in forbidden:
        with pytest.raises(ValidationError, match=f):
            KeyUpdateFields(**{f: True})

    with pytest.raises(ValidationError, match="at least one"):
        KeyUpdateFields()

    assert KeyUpdateFields(max_budget=50.0, tags=["x"]).max_budget == 50.0

    valid = KeyUpdateFields(max_budget=10)
    with pytest.raises(ValidationError):
        BulkUpdateTeamKeysRequest(
            team_id="t", key_ids=["k"], all_keys_in_team=True, update_fields=valid
        )
    with pytest.raises(ValidationError):
        BulkUpdateTeamKeysRequest(team_id="t", update_fields=valid)


# ---- security regressions ------------------------------------------------


@pytest.mark.asyncio
async def test_bulk_update_team_keys_hashes_raw_sk_key_ids(monkeypatch):
    """Regression: raw sk-... key_ids must be hashed before the find_many lookup."""
    from litellm.proxy._types import hash_token
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    raw_sk = "sk-rawkey1234567890"
    hashed = hash_token(raw_sk)
    row = LiteLLM_VerificationToken(
        token=hashed, user_id="u", models=[], team_id="team-abc", max_budget=None
    )
    mock = _setup_team_keys_mocks(
        monkeypatch,
        find_many=[row],
        find_unique=AsyncMock(return_value=row),
        update_data=AsyncMock(return_value={"data": _updated({"max_budget": 50.0})}),
        hash_identity=False,
    )

    response = await _call_as_admin(
        BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            key_ids=[raw_sk],
            update_fields=KeyUpdateFields(max_budget=50.0),
        )
    )
    where = mock.db.litellm_verificationtoken.find_many.await_args.kwargs["where"]
    assert where["token"] == {"in": [hashed]}
    # Response reports the user-supplied form, not the hash.
    assert response.successful_updates[0].key == raw_sk


@pytest.mark.asyncio
async def test_bulk_update_team_keys_auth_check_runs_when_no_keys_match(monkeypatch):
    """Regression: non-admin with bogus key_ids must still hit the membership gate."""
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_team_keys,
    )
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    mock = _setup_team_keys_mocks(monkeypatch, find_many=[])
    auth_check = AsyncMock(
        side_effect=ProxyException(
            message="not in team",
            type="team_member_permission_error",
            param="/key/update",
            code=401,
        )
    )
    monkeypatch.setattr(
        f"{_BULK_PKG}.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
        auth_check,
    )

    with pytest.raises(ProxyException):
        await bulk_update_team_keys(
            data=BulkUpdateTeamKeysRequest(
                team_id="victim-team",
                key_ids=["bogus-1", "bogus-2"],
                update_fields=KeyUpdateFields(max_budget=1.0),
            ),
            user_api_key_dict=_internal_user(),
            litellm_changed_by=None,
        )
    # Anchored on data.team_id, not existing_keys[0].
    assert auth_check.await_args.kwargs["existing_key_row"].team_id == "victim-team"
    mock.update_data.assert_not_called()


@pytest.mark.asyncio
async def test_bulk_update_team_keys_does_not_log_raw_sk_token_on_failure(
    monkeypatch, caplog
):
    """Regression: per-key failure must not log the raw sk-... (ERROR-level logs persist)."""
    import logging

    from litellm.proxy._types import hash_token
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    raw_sk = "sk-supersecret1234567890"
    row = LiteLLM_VerificationToken(
        token=hash_token(raw_sk),
        user_id="u",
        models=[],
        team_id="team-abc",
        max_budget=None,
    )
    _setup_team_keys_mocks(
        monkeypatch,
        find_many=[row],
        update_data=AsyncMock(side_effect=RuntimeError("boom")),
        hash_identity=False,
    )

    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
        response = await _call_as_admin(
            BulkUpdateTeamKeysRequest(
                team_id="team-abc",
                key_ids=[raw_sk],
                update_fields=KeyUpdateFields(max_budget=50.0),
            )
        )
    assert len(response.failed_updates) == 1
    log_text = "\n".join(r.getMessage() for r in caplog.records)
    assert raw_sk not in log_text


@pytest.mark.asyncio
async def test_bulk_update_team_keys_propagates_team_id_to_per_key_request(monkeypatch):
    """Regression: per-key UpdateKeyRequest carries data.team_id (gates _check_team_key_limits)."""
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])
    captured = []

    async def fake_process(*, update_key_request, **kw):
        captured.append(update_key_request)
        return {"max_budget": update_key_request.max_budget}

    monkeypatch.setattr(f"{_BULK_PKG}._process_single_key_update", fake_process)

    await _call_as_admin(
        BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            key_ids=["tok-a"],
            update_fields=KeyUpdateFields(
                tpm_limit=10_000, tpm_limit_type="guaranteed_throughput"
            ),
        )
    )
    assert captured[0].team_id == "team-abc"
    assert captured[0].tpm_limit_type == "guaranteed_throughput"


@pytest.mark.asyncio
async def test_bulk_update_team_keys_dedupes_key_ids(monkeypatch):
    """Duplicate key_ids collapse to a single update (no redundant DB writes, no inflated counts)."""
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    key_a = _make_team_key("tok-a")
    update_data = AsyncMock(return_value={"data": _updated({"max_budget": 50.0})})
    _setup_team_keys_mocks(
        monkeypatch,
        find_many=[key_a],
        find_unique=AsyncMock(return_value=key_a),
        update_data=update_data,
    )

    response = await _call_as_admin(
        BulkUpdateTeamKeysRequest(
            team_id="team-abc",
            key_ids=["tok-a", "tok-a", "tok-a"],
            update_fields=KeyUpdateFields(max_budget=50.0),
        )
    )

    assert response.total_requested == 1
    assert len(response.successful_updates) == 1
    assert len(response.failed_updates) == 0
    update_data.assert_awaited_once()


@pytest.mark.asyncio
async def test_bulk_update_team_keys_blocks_metadata_allowed_passthrough_routes(
    monkeypatch,
):
    """Non-admin can't grant passthrough access by smuggling allowed_passthrough_routes through metadata."""
    from fastapi import HTTPException

    from litellm.proxy.management_endpoints.key_management_endpoints import (
        bulk_update_team_keys,
    )
    from litellm.types.proxy.management_endpoints.key_management_endpoints import (
        BulkUpdateTeamKeysRequest,
        KeyUpdateFields,
    )

    mock = _setup_team_keys_mocks(monkeypatch, find_many=[_make_team_key("tok-a")])

    request = BulkUpdateTeamKeysRequest(
        team_id="team-abc",
        all_keys_in_team=True,
        update_fields=KeyUpdateFields(
            metadata={"allowed_passthrough_routes": ["/admin/*"]}
        ),
    )

    with pytest.raises(HTTPException) as exc:
        await bulk_update_team_keys(
            data=request,
            user_api_key_dict=_internal_user(),
            litellm_changed_by=None,
        )

    assert exc.value.status_code == 403
    assert "allowed_passthrough_routes" in str(exc.value.detail)
    mock.update_data.assert_not_called()


# ---------------------------------------------------------------------------
# /key/regenerate ownership-rebind guard + premium-gate identity check
# ---------------------------------------------------------------------------

import contextlib  # noqa: E402


def _non_admin_user_api_key_dict():
    return UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="user-1",
    )


@contextlib.contextmanager
def _patch_regenerate_side_effects():
    """Mock out token creation + DB write + cache + rotation hook so
    ``_execute_virtual_key_regeneration`` runs to completion under test."""
    with (
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
            new_callable=AsyncMock,
            return_value="sk-newtoken1234ab12",
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
            new_callable=AsyncMock,
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
            new_callable=AsyncMock,
        ),
    ):
        yield


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "incoming_user_id,expected_status,expected_substring",
    [
        # Cross-user rebind: the privesc primitive.
        ("default_user_id", 403, "not allowed to rebind the key"),
        # Empty-string removal: companion guard.
        ("", 403, "remove the user_id"),
        # Explicit null: same effect as empty-string removal — survives
        # model_dump(exclude_unset=True) and writes NULL to the token row.
        (None, 403, "remove the user_id"),
        # No-op rebind (caller sends their own user_id): must succeed.
        ("user-1", None, None),
    ],
    ids=[
        "rebind_blocked",
        "empty_blocked",
        "explicit_null_blocked",
        "same_user_id_allowed",
    ],
)
async def test_regenerate_user_id_rebind_guard(
    incoming_user_id, expected_status, expected_substring
):
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        _execute_virtual_key_regeneration,
    )

    existing_key = _make_regenerate_existing_key()
    data = RegenerateKeyRequest(user_id=incoming_user_id)

    async def _run():
        await _execute_virtual_key_regeneration(
            prisma_client=_make_regenerate_mock_prisma(),
            key_in_db=existing_key,
            hashed_api_key="abc123",
            key="abc123",
            data=data,
            user_api_key_dict=_non_admin_user_api_key_dict(),
            litellm_changed_by=None,
            user_api_key_cache=MagicMock(),
            proxy_logging_obj=MagicMock(),
        )

    if expected_status is None:
        with _patch_regenerate_side_effects():
            await _run()
        return

    with pytest.raises(HTTPException) as exc:
        await _run()
    assert exc.value.status_code == expected_status
    assert expected_substring in str(exc.value.detail)


@pytest.mark.asyncio
async def test_regenerate_premium_gate_requires_actual_master_key():
    # ``regenerate_key_fn``'s decorator wraps the underlying ValueError
    # into a ProxyException with empty ``message``. The exception type
    # alone confirms the premium gate fired.
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        regenerate_key_fn,
    )

    data = RegenerateKeyRequest(key="sk-not-master", new_master_key="anything")

    with (
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch("litellm.proxy.proxy_server.master_key", "sk-the-real-master-key"),
        patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
        pytest.raises((ValueError, HTTPException, ProxyException)),
    ):
        await regenerate_key_fn(
            key="sk-not-master",
            data=data,
            user_api_key_dict=_non_admin_user_api_key_dict(),
        )


@pytest.mark.asyncio
async def test_regenerate_premium_gate_allows_actual_master_key_holder():
    from litellm.proxy._types import RegenerateKeyRequest
    from litellm.proxy.management_endpoints.key_management_endpoints import (
        regenerate_key_fn,
    )

    master = "sk-the-real-master-key"
    data = RegenerateKeyRequest(key=master, new_master_key="sk-new-master")

    with (
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch("litellm.proxy.proxy_server.master_key", master),
        patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._rotate_master_key",
            new_callable=AsyncMock,
        ),
    ):
        result = await regenerate_key_fn(
            key=master,
            data=data,
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN,
                api_key=master,
                user_id="admin",
            ),
        )

    assert result.token == "sk-new-master"


# ---------------------------------------------------------------------------
# Regression tests for GHSA-q775-qw9r-2r4g: budget escalation via key/generate
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_ghsa_q775_non_admin_unlimited_can_delegate_budget():
    """
    Non-admin caller with max_budget=None (unlimited) can legitimately create
    budget-capped keys. Any finite budget is within an unlimited ceiling.
    """
    data = GenerateKeyRequest(max_budget=999999)
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="user-1",
        max_budget=None,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
            new_callable=AsyncMock,
            return_value=MagicMock(),
        ),
    ):
        result = await generate_key_fn(
            data=data,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )
        assert result is not None


@pytest.mark.asyncio
async def test_ghsa_q775_non_admin_cannot_exceed_own_budget():
    """
    Non-admin caller with max_budget=100 must not be able to create a key
    with max_budget=500.
    """
    data = GenerateKeyRequest(max_budget=500)
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="user-1",
        max_budget=100,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
    ):
        with pytest.raises((HTTPException, ProxyException)) as exc_info:
            await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )
        err = exc_info.value
        code = getattr(err, "status_code", None) or getattr(err, "code", None)
        msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
        assert str(code) == "400"
        assert "cannot exceed" in msg.lower()


@pytest.mark.asyncio
async def test_ghsa_q775_non_admin_within_budget_allowed():
    """
    Non-admin caller with max_budget=100 can create a key with max_budget=50.
    """
    data = GenerateKeyRequest(max_budget=50)
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="user-1",
        max_budget=100,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
            new_callable=AsyncMock,
            return_value=MagicMock(),
        ),
    ):
        result = await generate_key_fn(
            data=data,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )
        assert result is not None


@pytest.mark.asyncio
async def test_ghsa_q775_upperbound_default_not_rejected():
    """
    When upperbound_key_generate_params fills max_budget as a default, the
    ceiling check must NOT fire — only explicitly requested budgets trigger it.
    """
    data = GenerateKeyRequest()
    assert data.max_budget is None

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="user-1",
        max_budget=None,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
        patch(
            "litellm.upperbound_key_generate_params",
            MagicMock(max_budget=100.0),
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
            new_callable=AsyncMock,
            return_value=MagicMock(),
        ),
    ):
        result = await generate_key_fn(
            data=data,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )
        assert result is not None


@pytest.mark.asyncio
async def test_ghsa_q775_default_key_generate_params_not_rejected():
    """
    When default_key_generate_params fills max_budget, the ceiling check must
    NOT fire — only caller-supplied budgets trigger it.
    """
    data = GenerateKeyRequest()
    assert data.max_budget is None

    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-internal",
        user_id="user-1",
        max_budget=None,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
        patch(
            "litellm.default_key_generate_params",
            {"max_budget": 50.0},
        ),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
            new_callable=AsyncMock,
            return_value=MagicMock(),
        ),
    ):
        result = await generate_key_fn(
            data=data,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )
        assert result is not None


@pytest.mark.asyncio
async def test_ghsa_q775_admin_bypasses_budget_ceiling():
    """
    Admin caller can set any max_budget regardless of own budget.
    """
    data = GenerateKeyRequest(max_budget=999999)
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.PROXY_ADMIN,
        api_key="sk-admin",
        user_id="admin-1",
        max_budget=None,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
            new_callable=AsyncMock,
            return_value=MagicMock(),
        ),
    ):
        result = await generate_key_fn(
            data=data,
            user_api_key_dict=user_api_key_dict,
            litellm_changed_by=None,
        )
        assert result is not None


@pytest.mark.asyncio
async def test_ghsa_q775_ui_session_token_team_key_exempt_from_budget_ceiling():
    """
    Regression: a UI/CLI session token (team_id=litellm-dashboard) creating a
    TEAM key (data.team_id set) is exempt from the delegated-authority ceiling.
    The session max_budget is a per-session chat spend cap (max_ui_session_budget,
    default $0.25), not a delegation authority, and the team key's spend is bounded
    by the team budget at request time. This is the team-admin key-creation flow
    blocked since v1.86.x. Calls the helper directly so the ceiling runs (mocking
    out _common_key_generation_helper would mock out the check under test).
    """
    from litellm.constants import UI_SESSION_TOKEN_TEAM_ID

    data = GenerateKeyRequest(max_budget=500, team_id="team-abc")
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-ui-session",
        user_id="user-1",
        team_id=UI_SESSION_TOKEN_TEAM_ID,
        max_budget=0.25,
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
    ):
        try:
            await _common_key_generation_helper(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
                team_table=MagicMock(),
            )
        except (HTTPException, ProxyException) as err:
            msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
            assert (
                "cannot exceed" not in msg.lower()
            ), "UI/CLI session token creating a team key must be exempt from the ceiling"


@pytest.mark.asyncio
async def test_ghsa_q775_ui_session_token_personal_key_still_capped():
    """
    Security regression for GHSA-q775: the session-token exemption must NOT extend
    to personal keys. A UI/CLI session token (team_id=litellm-dashboard) creating a
    key with no data.team_id is still bound by the ceiling; otherwise a session
    token - or a leaked one, whose blast radius is the $0.25 chat cap - could mint
    an arbitrary-budget personal key, the exact escalation GHSA-q775 closed. Unlike
    a team key, nothing else bounds a personal key's spend.
    """
    from litellm.constants import UI_SESSION_TOKEN_TEAM_ID

    data = GenerateKeyRequest(max_budget=500)
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-ui-session",
        user_id="user-1",
        team_id=UI_SESSION_TOKEN_TEAM_ID,
        max_budget=0.25,
    )

    mock_prisma_client = AsyncMock()

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.user_custom_key_generate", None),
    ):
        with pytest.raises((HTTPException, ProxyException)) as exc_info:
            await generate_key_fn(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
            )
        err = exc_info.value
        code = getattr(err, "status_code", None) or getattr(err, "code", None)
        msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
        assert str(code) == "400"
        assert "cannot exceed" in msg.lower()


@pytest.mark.asyncio
async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption():
    """
    Security regression for GHSA-q775: the team-key exemption must key off the
    team_id the CALLER supplied, not one injected by default_key_generate_params.
    With default_key_generate_params.team_id set, a UI session token's personal-key
    request (no team_id) would otherwise have team_id auto-filled before the ceiling
    check, flipping is_ui_session_team_key to True and bypassing the ceiling. The
    request must still be rejected. Mirrors how _requested_max_budget is captured
    before defaults run.
    """
    from litellm.constants import UI_SESSION_TOKEN_TEAM_ID

    data = GenerateKeyRequest(max_budget=500)
    assert data.team_id is None
    user_api_key_dict = UserAPIKeyAuth(
        user_role=LitellmUserRoles.INTERNAL_USER,
        api_key="sk-ui-session",
        user_id="user-1",
        team_id=UI_SESSION_TOKEN_TEAM_ID,
        max_budget=0.25,
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
        patch("litellm.proxy.proxy_server.llm_router", None),
        patch("litellm.proxy.proxy_server.premium_user", False),
        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id"),
        patch("litellm.default_key_generate_params", {"team_id": "injected-team"}),
    ):
        with pytest.raises((HTTPException, ProxyException)) as exc_info:
            await _common_key_generation_helper(
                data=data,
                user_api_key_dict=user_api_key_dict,
                litellm_changed_by=None,
                team_table=None,
            )
        err = exc_info.value
        code = getattr(err, "status_code", None) or getattr(err, "code", None)
        msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", ""))
        assert str(code) == "400"
        assert "cannot exceed" in msg.lower()



@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_null_clears_fields():
    """
    When budget_duration is explicitly set to null, prepare_key_update_data
    should produce budget_duration=None and budget_reset_at=None so Prisma
    clears them in the DB.
    """
    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=[],
        user_id="test-user",
        team_id=None,
        metadata={},
    )

    update_request = UpdateKeyRequest(key="test-token", budget_duration=None)

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    assert "budget_duration" in result
    assert result["budget_duration"] is None
    assert "budget_reset_at" in result
    assert result["budget_reset_at"] is None


@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_not_sent_excluded():
    """
    When budget_duration is NOT sent in the request (unset), it should not
    appear in the result dict at all — the existing DB value stays unchanged.
    """
    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=[],
        user_id="test-user",
        team_id=None,
        metadata={},
    )

    update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"])

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    assert "budget_duration" not in result
    assert "budget_reset_at" not in result


@pytest.mark.asyncio
async def test_prepare_key_update_data_budget_duration_valid_sets_reset():
    """
    When budget_duration is set to a valid duration string, both
    budget_duration and budget_reset_at should be populated.
    """
    existing_key = LiteLLM_VerificationToken(
        token="test-token",
        key_alias="test-key",
        models=[],
        user_id="test-user",
        team_id=None,
        metadata={},
    )

    update_request = UpdateKeyRequest(key="test-token", budget_duration="30d")

    result = await prepare_key_update_data(
        data=update_request, existing_key_row=existing_key
    )

    assert result["budget_duration"] == "30d"
    assert result["budget_reset_at"] is not None


