litellm/tests/proxy_unit_tests/test_jwt_key_mapping.py

import pytest
import sys
import os
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch

# Add project root to sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))

from litellm.proxy.auth.user_api_key_auth import (
    _resolve_jwt_to_virtual_key,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy._types import (
    JWTKeyMappingResponse,
    LiteLLM_JWTAuth,
    LitellmUserRoles,
    UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
    _to_response,
    create_jwt_key_mapping,
    delete_jwt_key_mapping,
    info_jwt_key_mapping,
    update_jwt_key_mapping,
)
from litellm.caching.caching import DualCache
from fastapi import HTTPException

# ──────────────────────────────────────────────
# Tests: _resolve_jwt_to_virtual_key
# ──────────────────────────────────────────────


@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_resolution():
    """
    Test that a JWT claim is correctly resolved to a virtual key token.
    """
    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email", virtual_key_mapping_cache_ttl=3600
    )

    jwt_claims = {"email": "user@example.com", "sub": "123"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock()

    # Mock finding a mapping
    mock_mapping = MagicMock()
    mock_mapping.token = "sk-1234"
    mock_mapping.is_active = True
    prisma_client.db.litellm_jwtkeymapping.find_first.return_value = mock_mapping

    # Mock getting the key object
    mock_key_obj = UserAPIKeyAuth(token="sk-1234", team_id="team1")

    user_api_key_cache = DualCache()

    # Use patch to mock get_key_object in the module where it's used
    with patch(
        "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
        new_callable=AsyncMock,
    ) as mock_get_key:
        mock_get_key.return_value = mock_key_obj

        result = await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )

        assert result == mock_key_obj
        prisma_client.db.litellm_jwtkeymapping.find_first.assert_called_once()

        # Test Cache hit
        prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock()
        result_cached = await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
        assert result_cached == mock_key_obj
        prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()


@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_no_mapping():
    """
    Test that when no mapping exists, resolve returns None.
    """
    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="email")
    jwt_claims = {"email": "unknown@example.com"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock()
    prisma_client.db.litellm_jwtkeymapping.find_first.return_value = None

    # Mock get_key_object just in case
    with patch(
        "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
        new_callable=AsyncMock,
    ):
        user_api_key_cache = DualCache()

        result = await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )

        assert result is None

        # Test Negative Cache hit
        prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock()
        result_cached = await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
        assert result_cached is None
        prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()


# ──────────────────────────────────────────────
# Tests: OIDC / JWT routing in user_api_key_auth
# ──────────────────────────────────────────────


@pytest.mark.asyncio
async def test_virtual_key_mapping_oidc_enabled_jwt_token_uses_auth_jwt():
    """
    Regression test for the is_jwt routing fix in user_api_key_auth.py.

    When oidc_userinfo_enabled=True and virtual_key_claim_field is set, but
    the token is a well-formed JWT (3-part header.payload.sig), the virtual-key
    claim lookup must call auth_jwt — not get_oidc_userinfo.
    """
    # Three-part token: is_jwt() returns True
    api_key = "eyJhbGciOiJSUzI1NiJ9.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20ifQ.sig"

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        oidc_userinfo_enabled=True,
        virtual_key_claim_field="email",
    )

    # Confirm our fixture token is treated as a JWT
    assert jwt_handler.is_jwt(token=api_key) is True

    auth_jwt_mock = AsyncMock(return_value={"email": "user@example.com", "sub": "123"})
    oidc_userinfo_mock = AsyncMock(return_value={"email": "user@example.com"})

    # Simulate the routing condition from user_api_key_auth.py
    if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(
        token=api_key
    ):
        jwt_claims = await oidc_userinfo_mock(token=api_key)
    else:
        jwt_claims = await auth_jwt_mock(token=api_key)

    auth_jwt_mock.assert_called_once_with(token=api_key)
    oidc_userinfo_mock.assert_not_called()
    assert jwt_claims["email"] == "user@example.com"


@pytest.mark.asyncio
async def test_virtual_key_mapping_oidc_enabled_opaque_token_uses_oidc_userinfo():
    """
    Complement of the test above: when oidc_userinfo_enabled=True and the token
    is an opaque access token (not a JWT), the virtual-key claim lookup must
    call get_oidc_userinfo — not auth_jwt.
    """
    # Opaque token: no dots → is_jwt() returns False
    api_key = "some_opaque_access_token_with_no_dots"

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        oidc_userinfo_enabled=True,
        virtual_key_claim_field="email",
    )

    assert jwt_handler.is_jwt(token=api_key) is False

    auth_jwt_mock = AsyncMock(return_value={"email": "user@example.com"})
    oidc_userinfo_mock = AsyncMock(
        return_value={"email": "user@example.com", "sub": "123"}
    )

    if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(
        token=api_key
    ):
        jwt_claims = await oidc_userinfo_mock(token=api_key)
    else:
        jwt_claims = await auth_jwt_mock(token=api_key)

    oidc_userinfo_mock.assert_called_once_with(token=api_key)
    auth_jwt_mock.assert_not_called()
    assert jwt_claims["sub"] == "123"


# ──────────────────────────────────────────────
# Tests: _to_response redacts hashed token
# ──────────────────────────────────────────────


def test_to_response_excludes_token():
    """_to_response should not expose the hashed token field."""
    now = datetime.now(timezone.utc)
    mock_mapping = MagicMock()
    mock_mapping.id = "mapping-1"
    mock_mapping.jwt_claim_name = "email"
    mock_mapping.jwt_claim_value = "user@example.com"
    mock_mapping.token = "hashed_secret_value"
    mock_mapping.description = "test"
    mock_mapping.is_active = True
    mock_mapping.created_at = now
    mock_mapping.updated_at = now
    mock_mapping.created_by = "admin"
    mock_mapping.updated_by = "admin"

    resp = _to_response(mock_mapping)

    assert isinstance(resp, JWTKeyMappingResponse)
    assert resp.id == "mapping-1"
    assert resp.jwt_claim_name == "email"
    assert "token" not in resp.model_fields


# ──────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────


def _make_admin_auth() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(
        token="sk-admin",
        user_role=LitellmUserRoles.PROXY_ADMIN,
    )


def _make_non_admin_auth() -> UserAPIKeyAuth:
    return UserAPIKeyAuth(
        token="sk-user",
        user_role=LitellmUserRoles.INTERNAL_USER,
    )


def _mock_prisma():
    prisma = MagicMock()
    prisma.db.litellm_jwtkeymapping.create = AsyncMock()
    prisma.db.litellm_jwtkeymapping.find_unique = AsyncMock()
    prisma.db.litellm_jwtkeymapping.find_many = AsyncMock()
    prisma.db.litellm_jwtkeymapping.update = AsyncMock()
    prisma.db.litellm_jwtkeymapping.delete = AsyncMock()
    prisma.db.litellm_jwtkeymapping.count = AsyncMock(return_value=0)
    return prisma


def _mock_mapping(
    id="mapping-1",
    claim_name="email",
    claim_value="user@example.com",
):
    now = datetime.now(timezone.utc)
    m = MagicMock()
    m.id = id
    m.jwt_claim_name = claim_name
    m.jwt_claim_value = claim_value
    m.token = "hashed_token"
    m.description = None
    m.is_active = True
    m.created_at = now
    m.updated_at = now
    m.created_by = "admin"
    m.updated_by = "admin"
    return m


# ──────────────────────────────────────────────
# Tests: CRUD endpoint error handling
# ──────────────────────────────────────────────


@pytest.mark.asyncio
async def test_create_returns_409_on_unique_violation():
    """Duplicate mapping should return 409, not 500."""
    from litellm.proxy._types import CreateJWTKeyMappingRequest

    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
        "Unique constraint failed (P2002)"
    )
    mock_cache = AsyncMock()

    data = CreateJWTKeyMappingRequest(
        jwt_claim_name="email",
        jwt_claim_value="user@example.com",
        key="sk-test-key",
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
    ):
        with pytest.raises(HTTPException) as exc_info:
            await create_jwt_key_mapping(
                data=data, user_api_key_dict=_make_admin_auth()
            )
        assert exc_info.value.status_code == 409
        assert "already exists" in exc_info.value.detail


@pytest.mark.asyncio
async def test_create_returns_400_on_foreign_key_violation():
    """Non-existent key should return 400, not 500."""
    from litellm.proxy._types import CreateJWTKeyMappingRequest

    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.create.side_effect = Exception(
        "Foreign key constraint failed on field: `token` (P2003)"
    )
    mock_cache = AsyncMock()

    data = CreateJWTKeyMappingRequest(
        jwt_claim_name="sub",
        jwt_claim_value="user-999",
        key="sk-nonexistent",
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
    ):
        with pytest.raises(HTTPException) as exc_info:
            await create_jwt_key_mapping(
                data=data, user_api_key_dict=_make_admin_auth()
            )
        assert exc_info.value.status_code == 400
        assert "does not match" in exc_info.value.detail


@pytest.mark.asyncio
async def test_create_non_admin_returns_403():
    """Non-admin users should get 403."""
    from litellm.proxy._types import CreateJWTKeyMappingRequest

    data = CreateJWTKeyMappingRequest(
        jwt_claim_name="email",
        jwt_claim_value="user@example.com",
        key="sk-test",
    )

    with pytest.raises(HTTPException) as exc_info:
        await create_jwt_key_mapping(
            data=data, user_api_key_dict=_make_non_admin_auth()
        )
    assert exc_info.value.status_code == 403


@pytest.mark.asyncio
async def test_delete_returns_404_when_not_found():
    """Deleting non-existent mapping should return 404."""
    from litellm.proxy._types import DeleteJWTKeyMappingRequest

    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
    mock_cache = AsyncMock()

    data = DeleteJWTKeyMappingRequest(id="nonexistent-id")

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
    ):
        with pytest.raises(HTTPException) as exc_info:
            await delete_jwt_key_mapping(
                data=data, user_api_key_dict=_make_admin_auth()
            )
        assert exc_info.value.status_code == 404


@pytest.mark.asyncio
async def test_update_returns_404_when_not_found():
    """Updating non-existent mapping should return 404."""
    from litellm.proxy._types import UpdateJWTKeyMappingRequest

    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None
    mock_cache = AsyncMock()

    data = UpdateJWTKeyMappingRequest(id="nonexistent-id", description="test")

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
    ):
        with pytest.raises(HTTPException) as exc_info:
            await update_jwt_key_mapping(
                data=data, user_api_key_dict=_make_admin_auth()
            )
        assert exc_info.value.status_code == 404


@pytest.mark.asyncio
async def test_update_returns_404_when_row_deleted_before_write():
    """A mapping deleted between the read and the write must 404, not 500.

    Prisma's update returns None when the row is gone, and the endpoint used to
    dereference it for the cache key.
    """
    from litellm.proxy._types import UpdateJWTKeyMappingRequest

    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = _mock_mapping()
    mock_prisma.db.litellm_jwtkeymapping.update.return_value = None
    mock_cache = AsyncMock()

    data = UpdateJWTKeyMappingRequest(id="mapping-1", description="test")

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
        patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),  # test-quality-ok: proxy_server module global is the endpoint's only injection point
    ):
        with pytest.raises(HTTPException) as exc_info:
            await update_jwt_key_mapping(
                data=data, user_api_key_dict=_make_admin_auth()
            )
        assert exc_info.value.status_code == 404
        assert exc_info.value.detail == "Mapping not found"


@pytest.mark.asyncio
async def test_info_returns_404_when_not_found():
    """Getting info for non-existent mapping should return 404."""
    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.find_unique.return_value = None

    with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
        with pytest.raises(HTTPException) as exc_info:
            await info_jwt_key_mapping(
                id="nonexistent-id", user_api_key_dict=_make_admin_auth()
            )
        assert exc_info.value.status_code == 404


@pytest.mark.asyncio
async def test_create_success_returns_response_without_token():
    """Successful create should return JWTKeyMappingResponse without hashed token."""
    from litellm.proxy._types import CreateJWTKeyMappingRequest

    mock_prisma = _mock_prisma()
    mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
    mock_cache = AsyncMock()

    data = CreateJWTKeyMappingRequest(
        jwt_claim_name="email",
        jwt_claim_value="user@example.com",
        key="sk-test-key",
    )

    with (
        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
        patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache),
    ):
        result = await create_jwt_key_mapping(
            data=data, user_api_key_dict=_make_admin_auth()
        )
        assert isinstance(result, JWTKeyMappingResponse)
        assert "token" not in result.model_fields
        assert result.jwt_claim_name == "email"


# ──────────────────────────────────────────────
# Tests: unregistered_jwt_client_behavior
# ──────────────────────────────────────────────


@pytest.mark.asyncio
async def test_reject_behavior_raises_403_on_no_mapping():
    """
    When unregistered_jwt_client_behavior='reject' and no mapping exists,
    _resolve_jwt_to_virtual_key must raise HTTP 403.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.REJECT,
    )
    jwt_claims = {"email": "unknown@example.com"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)

    user_api_key_cache = DualCache()

    with patch(
        "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
        new_callable=AsyncMock,
    ):
        with pytest.raises(HTTPException) as exc_info:
            await _resolve_jwt_to_virtual_key(
                jwt_claims=jwt_claims,
                jwt_handler=jwt_handler,
                prisma_client=prisma_client,
                user_api_key_cache=user_api_key_cache,
                parent_otel_span=None,
                proxy_logging_obj=None,
            )
        assert exc_info.value.status_code == 403
        assert "unknown@example.com" in exc_info.value.detail


@pytest.mark.asyncio
async def test_reject_behavior_caches_sentinel_after_db_miss():
    """
    On a fresh DB miss with REJECT, the __NO_MAPPING__ sentinel must be written
    to cache so that subsequent rejected requests are served from cache and do
    not re-query the DB.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.REJECT,
        virtual_key_mapping_cache_ttl=300,
    )
    jwt_claims = {"email": "unknown@example.com"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)

    user_api_key_cache = DualCache()

    with patch(
        "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
        new_callable=AsyncMock,
    ):
        # First call — DB miss, should raise 403 and write sentinel
        with pytest.raises(HTTPException) as exc_info:
            await _resolve_jwt_to_virtual_key(
                jwt_claims=jwt_claims,
                jwt_handler=jwt_handler,
                prisma_client=prisma_client,
                user_api_key_cache=user_api_key_cache,
                parent_otel_span=None,
                proxy_logging_obj=None,
            )
        assert exc_info.value.status_code == 403

        # Sentinel must now be in cache
        cached = await user_api_key_cache.async_get_cache(
            "jwt_key_mapping:email:unknown@example.com"
        )
        assert cached == "__NO_MAPPING__"

        # Second call — must raise 403 from cache, no additional DB hit
        prisma_client.db.litellm_jwtkeymapping.find_first.reset_mock()
        with pytest.raises(HTTPException) as exc_info2:
            await _resolve_jwt_to_virtual_key(
                jwt_claims=jwt_claims,
                jwt_handler=jwt_handler,
                prisma_client=prisma_client,
                user_api_key_cache=user_api_key_cache,
                parent_otel_span=None,
                proxy_logging_obj=None,
            )
        assert exc_info2.value.status_code == 403
        prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()


@pytest.mark.asyncio
async def test_reject_behavior_raises_403_on_cached_no_mapping():
    """
    When the negative-cache sentinel __NO_MAPPING__ is present and behavior is
    'reject', the function must also raise HTTP 403 (not return None silently).
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.REJECT,
    )
    jwt_claims = {"email": "unknown@example.com"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)

    # Pre-populate the negative cache so the DB is not hit
    user_api_key_cache = DualCache()
    cache_key = "jwt_key_mapping:email:unknown@example.com"
    await user_api_key_cache.async_set_cache(cache_key, "__NO_MAPPING__")

    with patch(
        "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
        new_callable=AsyncMock,
    ):
        with pytest.raises(HTTPException) as exc_info:
            await _resolve_jwt_to_virtual_key(
                jwt_claims=jwt_claims,
                jwt_handler=jwt_handler,
                prisma_client=prisma_client,
                user_api_key_cache=user_api_key_cache,
                parent_otel_span=None,
                proxy_logging_obj=None,
            )
        assert exc_info.value.status_code == 403
        # DB must NOT have been hit (sentinel served from cache)
        prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()


@pytest.mark.asyncio
async def test_auto_register_returns_pending_signal_without_creating_key():
    """
    Security: when unregistered_jwt_client_behavior='auto_register' and no
    mapping exists, _resolve_jwt_to_virtual_key must NOT create the key yet.
    It returns a _PendingAutoRegister signal so the caller can run
    JWTAuthManager.auth_builder (enforcing RBAC, scope mappings,
    custom_validate, user_allowed_email_domain) FIRST. Creating the key here
    would bypass every JWT policy beyond signature verification.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior
    from litellm.proxy.auth.user_api_key_auth import _PendingAutoRegister

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )
    jwt_claims = {"sub": "new-user-42"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()

    user_api_key_cache = DualCache()

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
        new_callable=AsyncMock,
    ) as mock_gen_key:
        result = await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )

    assert isinstance(result, _PendingAutoRegister)
    assert result.claim_field == "sub"
    assert result.claim_value == "new-user-42"
    assert result.cache_key == "jwt_key_mapping:sub:new-user-42"
    # CRITICAL: no key was created — that must wait until after auth_builder
    mock_gen_key.assert_not_called()
    prisma_client.db.litellm_jwtkeymapping.create.assert_not_called()


@pytest.mark.asyncio
async def test_auto_register_creates_key_and_mapping_when_helper_invoked():
    """
    When the caller invokes _auto_register_jwt_mapping directly (after
    auth_builder validation), the helper creates the key + mapping row and
    returns a UserAPIKeyAuth. The mapping row stores the hashed token (FK to
    LiteLLM_VerificationToken), not the plaintext key.
    """
    from litellm.proxy._types import hash_token
    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        virtual_key_mapping_cache_ttl=300,
    )

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()

    user_api_key_cache = DualCache()
    plaintext_key = "sk-auto-key"
    expected_hash = hash_token(plaintext_key)
    mock_key_obj = UserAPIKeyAuth(token=expected_hash, team_id="validated-team")

    with (
        patch(
            "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
            new_callable=AsyncMock,
        ) as mock_get_key,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            new_callable=AsyncMock,
        ) as mock_gen_key,
    ):
        mock_gen_key.return_value = {"token": plaintext_key, "key": plaintext_key}
        mock_get_key.return_value = mock_key_obj

        result = await _auto_register_jwt_mapping(
            virtual_key_claim_field="sub",
            claim_value="new-user-42",
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
            cache_key="jwt_key_mapping:sub:new-user-42",
            team_id="validated-team",
            user_id="validated-user",
        )

    assert result == mock_key_obj
    # generate_key_helper_fn was passed table_name="key" (not user-upsert path)
    # and the validated team_id + user_id from auth_builder
    assert mock_gen_key.call_args.kwargs["table_name"] == "key"
    assert mock_gen_key.call_args.kwargs["team_id"] == "validated-team"
    assert mock_gen_key.call_args.kwargs["user_id"] == "validated-user"
    # Mapping row was created with the hashed token (FK target)
    call_data = prisma_client.db.litellm_jwtkeymapping.create.call_args[1]["data"]
    assert call_data["jwt_claim_name"] == "sub"
    assert call_data["jwt_claim_value"] == "new-user-42"
    assert call_data["token"] == expected_hash
    cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:new-user-42")
    assert cached == expected_hash


@pytest.mark.asyncio
async def test_auto_register_returns_pending_signal_on_stale_no_mapping_sentinel():
    """
    If the cache holds a stale __NO_MAPPING__ sentinel (written under a prior
    fallback_team_mapping config) and behavior is now AUTO_REGISTER, the
    resolver must evict the sentinel and return _PendingAutoRegister (so the
    caller can run auth_builder before creating the key) — not silently return
    None and not create the key on the spot.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior
    from litellm.proxy.auth.user_api_key_auth import _PendingAutoRegister

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )
    jwt_claims = {"email": "alice@corp.com"}

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()

    user_api_key_cache = DualCache()
    await user_api_key_cache.async_set_cache(
        "jwt_key_mapping:email:alice@corp.com", "__NO_MAPPING__"
    )

    with patch(
        "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
        new_callable=AsyncMock,
    ) as mock_gen_key:
        result = await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )

    assert isinstance(result, _PendingAutoRegister)
    # Stale sentinel must be evicted so the deferred auto-register actually
    # runs after auth_builder validates the JWT
    cached_after = await user_api_key_cache.async_get_cache(
        "jwt_key_mapping:email:alice@corp.com"
    )
    assert cached_after is None
    mock_gen_key.assert_not_called()
    prisma_client.db.litellm_jwtkeymapping.create.assert_not_called()


@pytest.mark.asyncio
async def test_auto_register_race_condition_unique_conflict():
    """
    If two concurrent requests both call _auto_register_jwt_mapping and the
    second hits a unique-constraint violation on create, it must:
      1) delete the orphaned virtual key it just created (so orphans don't
         accumulate in LiteLLM_VerificationToken under sustained concurrency),
      2) fall back to the winner's mapping,
      3) not surface an error.
    """
    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
    from litellm.proxy._types import UnregisteredJWTClientBehavior, hash_token

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
        side_effect=Exception("Unique constraint failed (P2002)")
    )
    prisma_client.db.litellm_verificationtoken.delete = AsyncMock()
    # Simulate the winner's mapping already in DB after the conflict
    winner_mapping = MagicMock()
    winner_mapping.token = "winner_token_hash"
    winner_mapping.is_active = True
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
        return_value=winner_mapping
    )

    user_api_key_cache = DualCache()
    loser_plaintext = "sk-loser"
    loser_hash = hash_token(loser_plaintext)
    mock_key_obj = UserAPIKeyAuth(token="winner_token_hash", team_id=None)

    with (
        patch(
            "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
            new_callable=AsyncMock,
        ) as mock_get_key,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            new_callable=AsyncMock,
            return_value={"token": loser_plaintext, "key": loser_plaintext},
        ),
    ):
        mock_get_key.return_value = mock_key_obj

        result = await _auto_register_jwt_mapping(
            virtual_key_claim_field="sub",
            claim_value="user-42",
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
            cache_key="jwt_key_mapping:sub:user-42",
        )

    assert result == mock_key_obj
    # The orphaned loser key must be deleted from LiteLLM_VerificationToken
    prisma_client.db.litellm_verificationtoken.delete.assert_called_once_with(
        where={"token": loser_hash}
    )
    # Cache should hold the winner's token, not the loser's
    cached = await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:user-42")
    assert cached == "winner_token_hash"
    mock_get_key.assert_called_once_with("winner_token_hash")


# ──────────────────────────────────────────────
# Tests: prisma_client=None does not bypass no-match policy
# ──────────────────────────────────────────────


@pytest.mark.asyncio
async def test_reject_behavior_enforced_when_prisma_client_is_none():
    """
    When prisma_client is None and behavior is REJECT, a 403 must be raised —
    not silently fallen through to team auth.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.REJECT,
    )
    jwt_claims = {"email": "unknown@example.com"}

    user_api_key_cache = DualCache()

    with pytest.raises(HTTPException) as exc_info:
        await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=None,  # no DB
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
    assert exc_info.value.status_code == 403
    assert "unknown@example.com" in exc_info.value.detail


@pytest.mark.asyncio
async def test_reject_raises_403_when_claim_field_missing_from_jwt():
    """
    Security: a JWT that omits the configured virtual_key_claim_field must NOT
    bypass the REJECT policy. Previously the early `if claim_value is None:
    return None` branch ran before the policy check, letting a caller who knows
    the configured claim-field name silently fall through to team-based auth.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.REJECT,
    )
    # JWT does NOT contain "sub"
    jwt_claims = {"email": "user@example.com"}

    with pytest.raises(HTTPException) as exc_info:
        await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=MagicMock(),
            user_api_key_cache=DualCache(),
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
    assert exc_info.value.status_code == 403
    assert "'sub'" in exc_info.value.detail
    assert "missing from the JWT" in exc_info.value.detail


@pytest.mark.asyncio
async def test_auto_register_raises_403_when_claim_field_missing_from_jwt():
    """
    AUTO_REGISTER cannot create a mapping without a stable identity. When the
    configured claim field is missing from the JWT, return 403 rather than
    silently falling through (which would bypass the unregistered-client policy)
    or creating a sentinel-keyed record.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
    )
    jwt_claims = {"email": "user@example.com"}

    with pytest.raises(HTTPException) as exc_info:
        await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=MagicMock(),
            user_api_key_cache=DualCache(),
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
    assert exc_info.value.status_code == 403
    assert "missing from the JWT" in exc_info.value.detail


@pytest.mark.asyncio
async def test_fallback_team_mapping_returns_none_when_claim_field_missing_from_jwt():
    """
    Under FALLBACK_TEAM_MAPPING (the default, backward-compatible mode), a JWT
    without the configured claim field must still fall through to team-based
    JWT auth — not raise. This preserves the pre-existing contract.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING,
    )
    jwt_claims = {"email": "user@example.com"}

    result = await _resolve_jwt_to_virtual_key(
        jwt_claims=jwt_claims,
        jwt_handler=jwt_handler,
        prisma_client=MagicMock(),
        user_api_key_cache=DualCache(),
        parent_otel_span=None,
        proxy_logging_obj=None,
    )
    assert result is None


@pytest.mark.asyncio
async def test_fallback_team_mapping_returns_none_when_prisma_client_is_none():
    """
    When prisma_client is None and behavior is FALLBACK_TEAM_MAPPING, the
    function must return None (fall through to team auth) — not raise.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="email",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.FALLBACK_TEAM_MAPPING,
    )
    jwt_claims = {"email": "anyone@example.com"}

    result = await _resolve_jwt_to_virtual_key(
        jwt_claims=jwt_claims,
        jwt_handler=jwt_handler,
        prisma_client=None,
        user_api_key_cache=DualCache(),
        parent_otel_span=None,
        proxy_logging_obj=None,
    )
    assert result is None


@pytest.mark.asyncio
async def test_auto_register_raises_500_when_prisma_client_is_none():
    """
    AUTO_REGISTER without a DB connection must raise HTTP 500 with a clear
    message — it cannot create keys without a database.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
    )
    jwt_claims = {"sub": "new-user-42"}

    with pytest.raises(HTTPException) as exc_info:
        await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=None,
            user_api_key_cache=DualCache(),
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
    assert exc_info.value.status_code == 500
    assert "AUTO_REGISTER requires a database" in exc_info.value.detail


@pytest.mark.asyncio
async def test_auto_register_raises_500_when_sentinel_cached_and_no_db():
    """
    AUTO_REGISTER + cached __NO_MAPPING__ sentinel + prisma_client is None must
    raise HTTP 500, matching the fresh-path behavior. Previously this path
    silently returned None and let the request fall through to team auth,
    creating different access-control outcomes under identical configuration.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )
    jwt_claims = {"sub": "user-42"}

    user_api_key_cache = DualCache()
    # Stale sentinel written under a prior fallback_team_mapping config
    await user_api_key_cache.async_set_cache(
        "jwt_key_mapping:sub:user-42", "__NO_MAPPING__"
    )

    with pytest.raises(HTTPException) as exc_info:
        await _resolve_jwt_to_virtual_key(
            jwt_claims=jwt_claims,
            jwt_handler=jwt_handler,
            prisma_client=None,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
        )
    assert exc_info.value.status_code == 500
    assert "AUTO_REGISTER requires a database" in exc_info.value.detail


@pytest.mark.asyncio
async def test_auto_register_race_conflict_tolerates_delete_failure():
    """
    If deleting the orphaned virtual key after a race-condition conflict fails
    (e.g. transient DB error), the request must still succeed by returning the
    winner's mapping — the orphan is unmapped and inert.
    """
    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
        side_effect=Exception("Unique constraint failed (P2002)")
    )
    prisma_client.db.litellm_verificationtoken.delete = AsyncMock(
        side_effect=Exception("transient DB error")
    )
    winner_mapping = MagicMock()
    winner_mapping.token = "winner_token_hash"
    winner_mapping.is_active = True
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
        return_value=winner_mapping
    )

    user_api_key_cache = DualCache()
    mock_key_obj = UserAPIKeyAuth(token="winner_token_hash", team_id=None)

    with (
        patch(
            "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
            new_callable=AsyncMock,
        ) as mock_get_key,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            new_callable=AsyncMock,
            return_value={"token": "sk-loser", "key": "sk-loser"},
        ),
    ):
        mock_get_key.return_value = mock_key_obj

        result = await _auto_register_jwt_mapping(
            virtual_key_claim_field="sub",
            claim_value="user-42",
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
            cache_key="jwt_key_mapping:sub:user-42",
        )

    # Caller still receives the winner's mapping even when cleanup fails
    assert result == mock_key_obj
    prisma_client.db.litellm_verificationtoken.delete.assert_called_once()


@pytest.mark.asyncio
async def test_auto_register_raises_503_when_winner_mapping_vanishes():
    """
    Race edge case: this request loses the unique-constraint race, deletes its
    orphan, then refetches the winner's mapping — but the winner's row was
    concurrently deleted. Previously this returned None, silently falling
    through to less-restrictive team-based JWT auth (bypassing the configured
    AUTO_REGISTER policy). Must now raise HTTP 503 so the caller retries
    rather than getting unintended fallback access.
    """
    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
        side_effect=Exception("Unique constraint failed (P2002)")
    )
    prisma_client.db.litellm_verificationtoken.delete = AsyncMock()
    # Winner row no longer exists by the time we refetch
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(return_value=None)

    user_api_key_cache = DualCache()

    with (
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            new_callable=AsyncMock,
            return_value={"token": "sk-loser", "key": "sk-loser"},
        ),
        pytest.raises(HTTPException) as exc_info,
    ):
        await _auto_register_jwt_mapping(
            virtual_key_claim_field="sub",
            claim_value="user-42",
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            parent_otel_span=None,
            proxy_logging_obj=None,
            cache_key="jwt_key_mapping:sub:user-42",
        )

    assert exc_info.value.status_code == 503
    assert "concurrently removed" in exc_info.value.detail


@pytest.mark.asyncio
async def test_proxy_admin_sentinel_skips_db_lookup_on_cache_hit():
    """
    When the cache holds the proxy-admin sentinel (written after a prior
    request's is_proxy_admin early-return), _resolve_jwt_to_virtual_key must
    return None *without* hitting the DB. Caller proceeds to auth_builder.

    Without this, every subsequent proxy-admin request under AUTO_REGISTER
    would re-query get_jwt_key_mapping_object — a cache-miss regression
    introduced by the deferred-auto-register refactor.
    """
    from litellm.proxy._types import UnregisteredJWTClientBehavior

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        unregistered_jwt_client_behavior=UnregisteredJWTClientBehavior.AUTO_REGISTER,
        virtual_key_mapping_cache_ttl=300,
    )
    jwt_claims = {"sub": "admin-user"}

    prisma_client = MagicMock()
    # Will fail the test if accessed — proves the sentinel short-circuits DB
    prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
        side_effect=AssertionError("DB must not be hit when sentinel is cached")
    )

    user_api_key_cache = DualCache()
    await user_api_key_cache.async_set_cache(
        "jwt_key_mapping:sub:admin-user", "__JWT_PROXY_ADMIN__"
    )

    result = await _resolve_jwt_to_virtual_key(
        jwt_claims=jwt_claims,
        jwt_handler=jwt_handler,
        prisma_client=prisma_client,
        user_api_key_cache=user_api_key_cache,
        parent_otel_span=None,
        proxy_logging_obj=None,
    )

    assert result is None
    prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()


# ──────────────────────────────────────────────
# Tests: AUTO_REGISTER stamps validated identity from auth_builder
# ──────────────────────────────────────────────


@pytest.mark.asyncio
async def test_auto_register_helper_stamps_validated_identity_context():
    """
    The deferred-auto-register contract: _auto_register_jwt_mapping is called
    with identity fields from JWTAuthManager.auth_builder's *validated*
    result (after RBAC, scope mappings, custom_validate, email-domain policy).
    These must be passed to generate_key_helper_fn so the created key carries
    them — the cached future-request path then inherits the same team/user/org
    limits the auth_builder path would have applied.
    """
    from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping

    jwt_handler = JWTHandler()
    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
        virtual_key_claim_field="sub",
        virtual_key_mapping_cache_ttl=300,
    )

    prisma_client = MagicMock()
    prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
    mock_key_obj = UserAPIKeyAuth(
        token="hashed", team_id="validated-team", user_id="validated-user"
    )

    with (
        patch(
            "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
            new_callable=AsyncMock,
        ) as mock_get_key,
        patch(
            "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
            new_callable=AsyncMock,
        ) as mock_gen_key,
    ):
        mock_gen_key.return_value = {"token": "sk-newkey", "key": "sk-newkey"}
        mock_get_key.return_value = mock_key_obj

        result = await _auto_register_jwt_mapping(
            virtual_key_claim_field="sub",
            claim_value="new-user",
            jwt_handler=jwt_handler,
            prisma_client=prisma_client,
            user_api_key_cache=DualCache(),
            parent_otel_span=None,
            proxy_logging_obj=None,
            cache_key="jwt_key_mapping:sub:new-user",
            team_id="validated-team",
            user_id="validated-user",
            org_id="validated-org",
            end_user_id="validated-end-user",
        )

    assert result == mock_key_obj
    assert mock_gen_key.call_args.kwargs["team_id"] == "validated-team"
    assert mock_gen_key.call_args.kwargs["user_id"] == "validated-user"
    assert mock_gen_key.call_args.kwargs["organization_id"] == "validated-org"
    assert result.org_id == "validated-org"
    assert result.end_user_id == "validated-end-user"


# ──────────────────────────────────────────────
# Tests: backward-compat alias jwt_client_id_field
# ──────────────────────────────────────────────


def test_jwt_client_id_field_alias_maps_to_virtual_key_claim_field():
    """
    jwt_client_id_field (old doc name) must silently alias to virtual_key_claim_field.
    """
    auth = LiteLLM_JWTAuth(jwt_client_id_field="azp")
    assert auth.virtual_key_claim_field == "azp"


def test_jwt_client_id_field_does_not_raise_on_duplicate():
    """
    If both jwt_client_id_field and virtual_key_claim_field are supplied,
    virtual_key_claim_field takes precedence and no error is raised.
    """
    auth = LiteLLM_JWTAuth(
        jwt_client_id_field="old_field",
        virtual_key_claim_field="new_field",
    )
    assert auth.virtual_key_claim_field == "new_field"