mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 21:04:10 +00:00
user
509 lines
18 KiB
Python
509 lines
18 KiB
Python
"""
|
||
Tests for the invite-link onboarding endpoints.
|
||
|
||
Covers the security behavior of:
|
||
GET /onboarding/get_token – rejects already-used links and returns only a
|
||
short-lived onboarding token, not a UI session key
|
||
POST /onboarding/claim_token – requires that onboarding token; mints the UI
|
||
session key only after the password is written
|
||
"""
|
||
|
||
from datetime import timedelta
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import jwt
|
||
import pytest
|
||
from fastapi import HTTPException
|
||
|
||
import litellm
|
||
from litellm.proxy._types import InvitationClaim
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class _AsyncTx:
|
||
def __init__(self, db: MagicMock):
|
||
self.db = db
|
||
|
||
async def __aenter__(self) -> MagicMock:
|
||
return self.db
|
||
|
||
async def __aexit__(self, exc_type, exc, tb):
|
||
return False
|
||
|
||
|
||
def _make_invite(
|
||
*, is_accepted: bool, expired: bool = False, claimed: bool = False
|
||
) -> MagicMock:
|
||
now = litellm.utils.get_utc_datetime()
|
||
invite = MagicMock()
|
||
invite.id = "invite-abc"
|
||
invite.user_id = "user-123"
|
||
invite.is_accepted = is_accepted
|
||
invite.expires_at = now - timedelta(days=1) if expired else now + timedelta(days=6)
|
||
invite.accepted_at = now if claimed else None
|
||
return invite
|
||
|
||
|
||
def _make_user() -> MagicMock:
|
||
user = MagicMock()
|
||
user.user_id = "user-123"
|
||
user.user_email = "alice@example.com"
|
||
user.user_role = "internal_user"
|
||
return user
|
||
|
||
|
||
def _make_prisma(invite: MagicMock, user: MagicMock | None = None) -> MagicMock:
|
||
prisma = MagicMock()
|
||
prisma.db.litellm_invitationlink.find_unique = AsyncMock(return_value=invite)
|
||
prisma.db.litellm_invitationlink.update = AsyncMock()
|
||
prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=1)
|
||
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user)
|
||
prisma.db.litellm_usertable.update = AsyncMock(return_value=user)
|
||
prisma.db.tx = MagicMock(return_value=_AsyncTx(prisma.db))
|
||
return prisma
|
||
|
||
|
||
def _make_onboarding_token(
|
||
*,
|
||
invitation_link: str = "invite-abc",
|
||
user_id: str = "user-123",
|
||
token_type: str = "litellm_onboarding",
|
||
master_key: str = "sk-test",
|
||
) -> str:
|
||
return jwt.encode(
|
||
{
|
||
"token_type": token_type,
|
||
"invitation_link": invitation_link,
|
||
"user_id": user_id,
|
||
"exp": litellm.utils.get_utc_datetime() + timedelta(minutes=15),
|
||
},
|
||
master_key,
|
||
algorithm="HS256",
|
||
)
|
||
|
||
|
||
def _make_claim_request(token: str | None = None) -> MagicMock:
|
||
request = MagicMock()
|
||
request.headers = {"Authorization": f"Bearer {token}"} if token is not None else {}
|
||
request.base_url = "http://localhost:4000/"
|
||
return request
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GET /onboarding/get_token
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_rejects_already_used_link():
|
||
"""
|
||
If is_accepted is True the link was already claimed.
|
||
The endpoint must raise 401 *before* returning any user data.
|
||
"""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
invite = _make_invite(is_accepted=True)
|
||
prisma = _make_prisma(invite)
|
||
request = MagicMock()
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await onboarding(invite_link="invite-abc", request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "already been used" in exc_info.value.detail["error"]
|
||
# The user table must never have been queried
|
||
prisma.db.litellm_usertable.find_unique.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_rejects_expired_link():
|
||
"""An expired link must raise 401 regardless of is_accepted."""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
invite = _make_invite(is_accepted=False, expired=True)
|
||
prisma = _make_prisma(invite)
|
||
request = MagicMock()
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await onboarding(invite_link="invite-abc", request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "expired" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_rejects_missing_link():
|
||
"""A link that does not exist in the DB must raise 401."""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
prisma = _make_prisma(invite=None) # type: ignore[arg-type]
|
||
request = MagicMock()
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await onboarding(invite_link="nonexistent", request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "does not exist" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_get_token_returns_onboarding_token_without_minting_ui_key():
|
||
"""
|
||
A valid, unused link should return a short-lived onboarding token, but
|
||
must not reserve the invite or mint a usable UI/API key on GET.
|
||
"""
|
||
from litellm.proxy.proxy_server import onboarding
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = MagicMock()
|
||
request.base_url = "http://localhost:4000/"
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
) as mock_generate_key,
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_custom_url",
|
||
return_value="http://localhost:4000/",
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
|
||
return_value=False,
|
||
),
|
||
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""),
|
||
):
|
||
result = await onboarding(invite_link="invite-abc", request=request)
|
||
|
||
# Endpoint succeeded
|
||
assert "token" in result
|
||
assert "login_url" in result
|
||
|
||
outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"])
|
||
onboarding_token = outer_claims["key"]
|
||
onboarding_claims = jwt.decode(onboarding_token, "sk-test", algorithms=["HS256"])
|
||
assert onboarding_claims["token_type"] == "litellm_onboarding"
|
||
assert onboarding_claims["invitation_link"] == "invite-abc"
|
||
assert onboarding_claims["user_id"] == "user-123"
|
||
assert not onboarding_token.startswith("sk-")
|
||
|
||
mock_generate_key.assert_not_called()
|
||
prisma.db.litellm_invitationlink.update_many.assert_not_called()
|
||
prisma.db.litellm_invitationlink.update.assert_not_called()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /onboarding/claim_token
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_already_used_link():
|
||
"""
|
||
If is_accepted is True, the password has already been set.
|
||
A second claim attempt must be rejected with 401.
|
||
"""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=True, claimed=True)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "already been used" in exc_info.value.detail["error"]
|
||
# Password must never have been written
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_expired_link():
|
||
"""An expired link must be rejected even if is_accepted is False."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False, expired=True)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "expired" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_mismatched_user_id():
|
||
"""The user_id in the request must match the one on the invite."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="wrong-user",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
with patch("litellm.proxy.proxy_server.prisma_client", prisma):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "does not match" in exc_info.value.detail["error"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_missing_onboarding_token():
|
||
"""The password endpoint must require the onboarding token returned by get_token."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=_make_claim_request())
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "Missing onboarding session" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_wrong_onboarding_session():
|
||
"""The onboarding token must be bound to the invite and user being claimed."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
request = _make_claim_request(
|
||
_make_onboarding_token(invitation_link="other-invite")
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "Invalid onboarding session" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_invalid_bearer_token():
|
||
"""A regular API key must not be accepted as an onboarding token."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
request = _make_claim_request("sk-regular-key")
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "Invalid onboarding session" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rejects_concurrent_reuse_before_password_write():
|
||
"""Only the first valid claim may reserve the invitation."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
prisma = _make_prisma(invite)
|
||
prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=0)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
) as mock_generate_key,
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 401
|
||
assert "already been used" in exc_info.value.detail["error"]
|
||
prisma.db.litellm_usertable.update.assert_not_called()
|
||
mock_generate_key.assert_not_called()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_sets_accepted_at_after_password_written():
|
||
"""
|
||
A valid first-time claim must:
|
||
1. Write the hashed password to the user table.
|
||
2. Set accepted_at on the invitation link after the password write succeeds.
|
||
"""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"}
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
return_value=mock_token_response,
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_custom_url",
|
||
return_value="http://localhost:4000/",
|
||
),
|
||
patch(
|
||
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
|
||
return_value=False,
|
||
),
|
||
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""),
|
||
):
|
||
result = await claim_onboarding_link(data=data, request=request)
|
||
|
||
# Password was written
|
||
prisma.db.litellm_invitationlink.update_many.assert_called_once()
|
||
reserve_kwargs = prisma.db.litellm_invitationlink.update_many.call_args.kwargs
|
||
assert reserve_kwargs["where"] == {"id": "invite-abc", "is_accepted": False}
|
||
assert reserve_kwargs["data"]["is_accepted"] is True
|
||
prisma.db.litellm_usertable.update.assert_called_once()
|
||
call_kwargs = prisma.db.litellm_usertable.update.call_args
|
||
assert call_kwargs.kwargs["where"] == {"user_id": "user-123"}
|
||
assert "password" in call_kwargs.kwargs["data"]
|
||
|
||
# is_accepted was flipped to True on the invitation link
|
||
prisma.db.litellm_invitationlink.update.assert_called_once()
|
||
link_update_data = prisma.db.litellm_invitationlink.update.call_args.kwargs["data"]
|
||
assert "is_accepted" not in link_update_data
|
||
assert link_update_data["accepted_at"] is not None
|
||
outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"])
|
||
assert outer_claims["key"] == "sk-generated-key"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_claim_token_rolls_back_invite_when_session_key_mint_fails():
|
||
"""A session key failure must not leave the invite permanently consumed."""
|
||
from litellm.proxy.proxy_server import claim_onboarding_link
|
||
|
||
invite = _make_invite(is_accepted=False)
|
||
user = _make_user()
|
||
prisma = _make_prisma(invite, user)
|
||
request = _make_claim_request(_make_onboarding_token())
|
||
|
||
data = InvitationClaim(
|
||
invitation_link="invite-abc",
|
||
user_id="user-123",
|
||
password="NewP@ssw0rd",
|
||
)
|
||
|
||
with (
|
||
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
||
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
|
||
patch("litellm.proxy.proxy_server.general_settings", {}),
|
||
patch(
|
||
"litellm.proxy.proxy_server.generate_key_helper_fn",
|
||
new_callable=AsyncMock,
|
||
side_effect=Exception("key mint failed"),
|
||
),
|
||
):
|
||
with pytest.raises(HTTPException) as exc_info:
|
||
await claim_onboarding_link(data=data, request=request)
|
||
|
||
assert exc_info.value.status_code == 500
|
||
assert "Failed to create onboarding session" in exc_info.value.detail["error"]
|
||
assert prisma.db.litellm_invitationlink.update_many.call_count == 2
|
||
rollback_kwargs = prisma.db.litellm_invitationlink.update_many.call_args_list[
|
||
1
|
||
].kwargs
|
||
assert rollback_kwargs["where"] == {
|
||
"id": "invite-abc",
|
||
"is_accepted": True,
|
||
}
|
||
assert rollback_kwargs["data"]["accepted_at"] is None
|
||
assert rollback_kwargs["data"]["is_accepted"] is False
|