mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(users)!: validate custom IDs and reject concurrent duplicate creation (#49748)
* fix(users): validate custom IDs and reject concurrent duplicate creation * fix(users): return precise conflicts without retrying them * fix(users): preserve retries for existing-user upserts * fix(scim): preserve opaque IDs during user provisioning * fix(sso): preserve identity provider user IDs
This commit is contained in:
parent
9984ed40a1
commit
fb2a1450dd
14 changed files with 267 additions and 24 deletions
|
|
@ -5482,7 +5482,9 @@ async def generate_key_helper_fn(
|
|||
if query_type == "insert_data":
|
||||
user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row
|
||||
"prisma_models.LiteLLM_UserTable | None",
|
||||
await prisma_client.insert_data(data=user_data, table_name="user"),
|
||||
await prisma_client.insert_data(
|
||||
data=user_data, table_name="user", allow_existing_user=request_type != "user"
|
||||
),
|
||||
)
|
||||
|
||||
if user_row is None:
|
||||
|
|
|
|||
|
|
@ -1294,7 +1294,7 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou
|
|||
if litellm.default_internal_user_params:
|
||||
default_role = litellm.default_internal_user_params.get("user_role")
|
||||
|
||||
new_user_request: Final = NewUserRequest(
|
||||
new_user_request: Final = SCIMNewUserRequest(
|
||||
user_id=user_id,
|
||||
user_email=user_id, # We don't have email from group membership
|
||||
user_alias=None,
|
||||
|
|
@ -1873,7 +1873,7 @@ async def create_user(
|
|||
admin_group: Final = await _get_scim_admin_group()
|
||||
resolved_role: Final = _resolve_scim_user_role(user.groups or [], admin_group, default_role)
|
||||
|
||||
new_user_request: Final = NewUserRequest(
|
||||
new_user_request: Final = SCIMNewUserRequest(
|
||||
user_id=user_id,
|
||||
user_email=user_data["user_email"],
|
||||
user_alias=user_data["user_alias"],
|
||||
|
|
|
|||
|
|
@ -78,7 +78,6 @@ from litellm.proxy._types import (
|
|||
LitellmUserRoles,
|
||||
Member,
|
||||
NewTeamRequest,
|
||||
NewUserRequest,
|
||||
NewUserRequestTeam,
|
||||
NewUserResponse,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -164,6 +163,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
|
|||
MicrosoftGraphAPIUserGroupResponse,
|
||||
MicrosoftServicePrincipalTeam,
|
||||
RoleMappings,
|
||||
SSONewUserRequest,
|
||||
TeamMappings,
|
||||
)
|
||||
from litellm.types.proxy.ui_sso import ParsedOpenIDResult
|
||||
|
|
@ -2915,7 +2915,7 @@ async def insert_sso_user(
|
|||
if user_defined_values["user_role"] is None:
|
||||
user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
|
||||
|
||||
new_user_request: Final = NewUserRequest(
|
||||
new_user_request: Final = SSONewUserRequest(
|
||||
user_id=user_defined_values["user_id"],
|
||||
user_email=normalize_email(user_defined_values["user_email"]),
|
||||
user_alias=_get_sso_user_alias(result_openid),
|
||||
|
|
@ -2924,11 +2924,9 @@ async def insert_sso_user(
|
|||
budget_duration=user_defined_values["budget_duration"],
|
||||
sso_user_id=user_defined_values["user_id"],
|
||||
auto_create_key=False,
|
||||
metadata={"auth_provider": getattr(result_openid, "provider")} if hasattr(result_openid, "provider") else {},
|
||||
)
|
||||
|
||||
if result_openid and hasattr(result_openid, "provider"):
|
||||
new_user_request.metadata = {"auth_provider": getattr(result_openid, "provider")}
|
||||
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
with free_sso_user_creation(new_user_request.user_id, enabled=premium_user is not True):
|
||||
|
|
|
|||
|
|
@ -313,6 +313,7 @@ if TYPE_CHECKING:
|
|||
from prisma import models as prisma_models
|
||||
from prisma.actions import LiteLLM_DeprecatedVerificationTokenActions
|
||||
from prisma.client import TransactionManager
|
||||
from prisma.errors import UniqueViolationError
|
||||
from prisma.models import LiteLLM_DeprecatedVerificationToken
|
||||
from prisma.types import HttpConfig, LiteLLM_VerificationTokenInclude
|
||||
|
||||
|
|
@ -5559,6 +5560,25 @@ _LOCK_TEAMS_BY_ID_SQL: Final = (
|
|||
)
|
||||
|
||||
|
||||
def _user_unique_conflict_message(metadata: object, user_id: object) -> str:
|
||||
match metadata:
|
||||
case {"target": ["user_id"]}:
|
||||
return f"User with id {user_id} already exists"
|
||||
case {"target": ["sso_user_id"]}:
|
||||
return "User with this sso_user_id already exists"
|
||||
case _:
|
||||
return "User conflicts with an existing unique value"
|
||||
|
||||
|
||||
def _raise_user_unique_conflict(error: "UniqueViolationError", user_id: object, *, allow_existing: bool) -> NoReturn:
|
||||
if allow_existing:
|
||||
raise error
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail={"error": _user_unique_conflict_message(error.meta, user_id)},
|
||||
) from error
|
||||
|
||||
|
||||
class PrismaClient:
|
||||
spend_log_transactions: list = []
|
||||
_spend_log_transactions_lock = asyncio.Lock()
|
||||
|
|
@ -6446,12 +6466,15 @@ class PrismaClient:
|
|||
max_tries=3, # maximum number of retries
|
||||
max_time=10, # maximum total time to retry for
|
||||
on_backoff=on_backoff, # specifying the function to call on backoff
|
||||
giveup=lambda exc: isinstance(exc, HTTPException) and exc.status_code == 409,
|
||||
)
|
||||
@log_db_metrics
|
||||
async def insert_data(
|
||||
self,
|
||||
data: Mapping[str, object],
|
||||
table_name: Literal["user", "key", "config", "spend", "team", "user_notification"],
|
||||
*,
|
||||
allow_existing_user: bool = True,
|
||||
):
|
||||
"""
|
||||
Add a key to the database. If it already exists, do nothing.
|
||||
|
|
@ -6487,14 +6510,20 @@ class PrismaClient:
|
|||
verbose_proxy_logger.info("Data Inserted into Keys Table")
|
||||
return new_verification_token
|
||||
elif table_name == "user":
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
db_data = self.jsonify_object(data=data)
|
||||
try:
|
||||
new_user_row: Final = await UserRepository(self).upsert_user(db_data)
|
||||
new_user_row: Final = await UserRepository(self).upsert_user(
|
||||
db_data, allow_existing=allow_existing_user
|
||||
)
|
||||
if new_user_row is None:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="SSO for more than 5 users requires a LiteLLM Enterprise license. Existing users can still sign in.",
|
||||
)
|
||||
except UniqueViolationError as e:
|
||||
_raise_user_unique_conflict(e, db_data["user_id"], allow_existing=allow_existing_user)
|
||||
except Exception as e:
|
||||
if (
|
||||
"Foreign key constraint failed on the field: `LiteLLM_UserTable_organization_id_fkey (index)`"
|
||||
|
|
|
|||
|
|
@ -168,7 +168,9 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
|
|||
"""
|
||||
return await _count_billable_users(self.table)
|
||||
|
||||
async def upsert_user(self, data: Mapping[str, object]) -> "prisma_models.LiteLLM_UserTable | None":
|
||||
async def upsert_user(
|
||||
self, data: Mapping[str, object], *, allow_existing: bool = True
|
||||
) -> "prisma_models.LiteLLM_UserTable | None":
|
||||
"""Return the user row, or None when a new free SSO user would exceed capacity."""
|
||||
scope: Final = _FREE_SSO_CREATION.get()
|
||||
if (
|
||||
|
|
@ -176,6 +178,8 @@ class UserRepository(BaseRepository[LiteLLM_UserTable]):
|
|||
or (scope[0] is not None and scope[0] != data["user_id"])
|
||||
or scope[1] is not asyncio.current_task()
|
||||
):
|
||||
if not allow_existing:
|
||||
return await self.table.create(data=dict(data))
|
||||
return await self.table.upsert(
|
||||
where={"user_id": data["user_id"]}, data={"create": dict(data), "update": {}}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -120,6 +120,7 @@ class BulkNewUserItem(NewUserRequest):
|
|||
|
||||
model_config = ConfigDict(extra="forbid", protected_namespaces=())
|
||||
|
||||
user_id: str | None = None
|
||||
auto_create_key: bool = False
|
||||
|
||||
@field_validator("send_invite_email")
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from pydantic import (
|
|||
from pydantic_core.core_schema import SerializerFunctionWrapHandler
|
||||
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest
|
||||
|
||||
SCIM_ENTERPRISE_USER_SCHEMA: Final = "urn:ietf:params:scim:schemas:extension:enterprise:2.0:User"
|
||||
SCIM_ENTERPRISE_METADATA_KEY: Final = "scim_enterprise"
|
||||
|
|
@ -26,6 +27,14 @@ _SCIM_SERIALIZED_FIELDS_ADAPTER: Final = TypeAdapter(
|
|||
)
|
||||
|
||||
|
||||
class SCIMNewUserRequest(NewUserRequest):
|
||||
"""Preserve opaque directory IDs when provisioning SCIM users and group members."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
class LiteLLM_UserScimMetadata(LiteLLMBaseModel):
|
||||
"""
|
||||
Scim metadata stored in LiteLLM_UserTable.metadata
|
||||
|
|
|
|||
|
|
@ -1,12 +1,19 @@
|
|||
from typing import Literal
|
||||
|
||||
from pydantic import Field
|
||||
from pydantic import ConfigDict, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm.types.proxy.auth.user_roles import KeyManagementRoutes, LitellmUserRoles
|
||||
from litellm.types.proxy.management_endpoints.user_requests import NewUserRequest
|
||||
from litellm.types.utils import LiteLLMPydanticObjectBase
|
||||
|
||||
|
||||
class SSONewUserRequest(NewUserRequest):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase):
|
||||
"""
|
||||
Set default upperbound to max budget a key called via `/key/generate` can be.
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Literal
|
||||
from typing import Annotated, Final, Literal
|
||||
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
from pydantic import AfterValidator, Field, field_validator, model_validator
|
||||
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
|
||||
|
|
@ -15,7 +15,17 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase):
|
|||
user_role: Literal["user", "admin"] = "user"
|
||||
|
||||
|
||||
def _validate_custom_user_id(value: str) -> str:
|
||||
normalized: Final = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("user_id must not be empty or contain only whitespace")
|
||||
if "\x00" in normalized:
|
||||
raise ValueError("user_id must not contain NUL characters")
|
||||
return normalized
|
||||
|
||||
|
||||
class NewUserRequest(GenerateRequestBase):
|
||||
user_id: Annotated[str, AfterValidator(_validate_custom_user_id)] | None = None
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
max_budget: float | None = None
|
||||
|
|
|
|||
|
|
@ -411,10 +411,11 @@ async def test_create_user_without_groups_defers_to_default_team(mocker: MockerF
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch):
|
||||
@pytest.mark.parametrize("user_id", ["new-user", " employee-42 "])
|
||||
async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch, user_id: str):
|
||||
scim_user = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
userName="new-user",
|
||||
userName=user_id,
|
||||
emails=[SCIMUserEmail(value="new@example.com")],
|
||||
groups=[SCIMUserGroup(value="idp-team")],
|
||||
)
|
||||
|
|
@ -428,10 +429,12 @@ async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, mo
|
|||
await create_user(user=scim_user)
|
||||
|
||||
assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"]
|
||||
assert new_user_mock.call_args.kwargs["data"].user_id == user_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch):
|
||||
@pytest.mark.parametrize("user_id", ["group-user", " employee-42 "])
|
||||
async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch, user_id: str):
|
||||
monkeypatch.setattr(
|
||||
"litellm.default_internal_user_params",
|
||||
{"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]},
|
||||
|
|
@ -439,14 +442,15 @@ async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFi
|
|||
)
|
||||
new_user_mock = mocker.patch( # test-quality-ok: new_user is imported inside the helper, not injectable
|
||||
"litellm.proxy.management_endpoints.internal_user_endpoints.new_user",
|
||||
AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")),
|
||||
AsyncMock(return_value=NewUserResponse(user_id=user_id, key="k")),
|
||||
)
|
||||
|
||||
created = await _create_user_if_not_exists(user_id="group-user")
|
||||
created = await _create_user_if_not_exists(user_id=user_id)
|
||||
|
||||
assert created is not None
|
||||
assert new_user_mock.call_args.kwargs["data"].teams is None
|
||||
assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
assert new_user_mock.call_args.kwargs["data"].user_id == user_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import httpx
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import ValidationError
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
|
@ -48,6 +49,51 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import (
|
|||
|
||||
client = TestClient(app)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("user_id", ["", " \t\n", "user\x00id", 123, True, [], {}])
|
||||
def test_new_user_rejects_invalid_custom_id(user_id: object) -> None:
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
NewUserRequest.model_validate({"user_id": user_id})
|
||||
assert exc.value.errors()[0]["loc"] == ("user_id",)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("user_id", [None, "employee-42", "社員@example.test", " employee-42 "])
|
||||
def test_new_user_normalizes_custom_id_and_preserves_auto_generation(user_id: str | None) -> None:
|
||||
request: Final = NewUserRequest(user_id=user_id)
|
||||
assert request.user_id == (user_id.strip() if user_id is not None else None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_user_returns_conflict_when_id_is_taken_after_duplicate_check(mocker: MockerFixture) -> None:
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
table: Final = SimpleNamespace(
|
||||
find_first=AsyncMock(return_value=None),
|
||||
count=AsyncMock(return_value=0),
|
||||
create=AsyncMock(side_effect=UniqueViolationError(
|
||||
{"user_facing_error": {"meta": {"target": ["user_id"]}}},
|
||||
message="Unique constraint failed on user_id",
|
||||
)),
|
||||
upsert=AsyncMock(return_value=LiteLLM_UserTable(user_id="employee-42")),
|
||||
)
|
||||
mocker.patch("prisma.Prisma", return_value=SimpleNamespace(litellm_usertable=table))
|
||||
database: Final = PrismaClient(
|
||||
database_url="postgresql://unused:unused@localhost/unused",
|
||||
proxy_logging_obj=proxy_server.proxy_logging_obj,
|
||||
)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", database)
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await new_user(
|
||||
data=NewUserRequest(user_id="employee-42", auto_create_key=False, send_invite_email=False),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
assert str(exc.value.code) == "409"
|
||||
assert "employee-42" in exc.value.message
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
|
@ -4862,7 +4908,9 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget,
|
|||
models = []
|
||||
|
||||
class _FakePrisma:
|
||||
async def insert_data(self, data, table_name):
|
||||
async def insert_data(
|
||||
self, data: Mapping[str, object], table_name: str, *, allow_existing_user: bool = True
|
||||
) -> _FakeUserRow | SimpleNamespace:
|
||||
if table_name == "user":
|
||||
captured["user_data"] = dict(data)
|
||||
return _FakeUserRow()
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.proxy._types import (
|
|||
LiteLLM_JWTAuth,
|
||||
LiteLLM_UserTable,
|
||||
NewUserResponse,
|
||||
SSOUserDefinedValues,
|
||||
ProxyException,
|
||||
)
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
|
@ -52,6 +53,62 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
|
|||
from tests._master_key import MASTER_KEY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("billable_users", "premium", "admitted"),
|
||||
[(4, False, True), (5, False, False), (26, False, False), (26, True, True)],
|
||||
)
|
||||
async def test_sso_auto_creation_enforces_free_user_cap_through_new_user(
|
||||
monkeypatch: pytest.MonkeyPatch, billable_users: int, premium: bool, admitted: bool
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
row: Final = LiteLLM_UserTable(user_id="new-sso-user", user_email="new-sso-user@example.test")
|
||||
table: Final = SimpleNamespace(
|
||||
find_first=AsyncMock(return_value=None),
|
||||
find_unique=AsyncMock(return_value=None),
|
||||
count=AsyncMock(side_effect=lambda where=None: billable_users if where is None else 0),
|
||||
create=AsyncMock(return_value=row),
|
||||
upsert=AsyncMock(return_value=row),
|
||||
)
|
||||
transaction: Final = MagicMock()
|
||||
transaction.__aenter__ = AsyncMock(
|
||||
return_value=SimpleNamespace(litellm_usertable=table, execute_raw=AsyncMock(), query_raw=AsyncMock())
|
||||
)
|
||||
transaction.__aexit__ = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"prisma.Prisma",
|
||||
lambda **kwargs: SimpleNamespace(litellm_usertable=table, tx=MagicMock(return_value=transaction)),
|
||||
)
|
||||
database: Final = PrismaClient(
|
||||
database_url="postgresql://unused:unused@localhost/unused",
|
||||
proxy_logging_obj=proxy_server.proxy_logging_obj,
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", database)
|
||||
monkeypatch.setattr(proxy_server, "premium_user", premium)
|
||||
monkeypatch.setattr("asyncio.sleep", AsyncMock())
|
||||
values: Final[SSOUserDefinedValues] = {
|
||||
"models": [],
|
||||
"user_id": row.user_id,
|
||||
"user_email": row.user_email,
|
||||
"user_role": "internal_user_viewer",
|
||||
"max_budget": None,
|
||||
"budget_duration": None,
|
||||
}
|
||||
if admitted:
|
||||
response: Final = await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values)
|
||||
assert response.user_id == row.user_id
|
||||
assert table.create.await_count + table.upsert.await_count == 1
|
||||
else:
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await insert_sso_user(CustomOpenID(id=row.user_id, team_ids=[]), values)
|
||||
assert str(error.value.code) == "403"
|
||||
assert "SSO for more than 5 users" in error.value.message
|
||||
table.create.assert_not_awaited()
|
||||
table.upsert.assert_not_awaited()
|
||||
|
||||
|
||||
def _sso_http_handler(route: Callable[[httpx.Request], httpx.Response] | None = None) -> AsyncHTTPHandler:
|
||||
"""litellm's SSO httpx client answered in-process, by ``route`` or by the active respx mock."""
|
||||
handler = AsyncHTTPHandler()
|
||||
|
|
@ -1281,7 +1338,8 @@ async def test_upsert_sso_user_fills_user_alias_for_existing_user():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_sso_user_sets_user_alias_from_display_name():
|
||||
@pytest.mark.parametrize("user_id", ["S-1-5-21-adfs-user", " employee-42 "])
|
||||
async def test_insert_sso_user_sets_user_alias_from_display_name(user_id: str):
|
||||
"""
|
||||
A newly created SSO user is inserted with the IdP display name as user_alias.
|
||||
"""
|
||||
|
|
@ -1290,7 +1348,7 @@ async def test_insert_sso_user_sets_user_alias_from_display_name():
|
|||
from litellm.proxy.management_endpoints.ui_sso import insert_sso_user
|
||||
|
||||
sso_result = CustomOpenID(
|
||||
id="S-1-5-21-adfs-user",
|
||||
id=user_id,
|
||||
email="jane.doe@example.com",
|
||||
first_name="Jane",
|
||||
last_name="Doe",
|
||||
|
|
@ -1300,7 +1358,7 @@ async def test_insert_sso_user_sets_user_alias_from_display_name():
|
|||
)
|
||||
user_defined_values: SSOUserDefinedValues = {
|
||||
"models": [],
|
||||
"user_id": "S-1-5-21-adfs-user",
|
||||
"user_id": user_id,
|
||||
"user_email": "jane.doe@example.com",
|
||||
"max_budget": None,
|
||||
"user_role": "internal_user",
|
||||
|
|
@ -1309,12 +1367,14 @@ async def test_insert_sso_user_sets_user_alias_from_display_name():
|
|||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.ui_sso.new_user",
|
||||
return_value=NewUserResponse(user_id="S-1-5-21-adfs-user", key="sk-xxxxx", teams=None),
|
||||
return_value=NewUserResponse(user_id=user_id, key="sk-xxxxx", teams=None),
|
||||
) as mock_new_user:
|
||||
await insert_sso_user(result_openid=sso_result, user_defined_values=user_defined_values)
|
||||
|
||||
new_user_request = mock_new_user.call_args.kwargs["data"]
|
||||
assert new_user_request.user_id == "S-1-5-21-adfs-user"
|
||||
assert new_user_request.user_id == user_id
|
||||
assert new_user_request.sso_user_id == user_id
|
||||
assert new_user_request.metadata == {"auth_provider": "generic"}
|
||||
assert new_user_request.user_email == "jane.doe@example.com"
|
||||
assert new_user_request.user_alias == "Doe, Jane"
|
||||
|
||||
|
|
|
|||
|
|
@ -191,6 +191,19 @@ async def _run(prisma, users, caller=ADMIN, license=None, generate_key=_no_keys)
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("user_id", [None, "", " employee-42 ", " \t"])
|
||||
async def test_bulk_keeps_existing_user_id_generation_and_whitespace_behavior(user_id: str | None) -> None:
|
||||
prisma: Final = _FakePrisma()
|
||||
response: Final = await _run(prisma, [{"user_id": user_id}])
|
||||
result: Final = response.data[0]
|
||||
assert result.success
|
||||
assert result.user_id
|
||||
assert result.user_id in prisma.db.litellm_usertable.rows
|
||||
if user_id:
|
||||
assert result.user_id == user_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_users_and_team_membership_in_every_store():
|
||||
prisma = _FakePrisma(teams=[_team("t1", [Member(user_id="existing", role="admin")]), _team("t2")])
|
||||
|
|
|
|||
|
|
@ -18,8 +18,10 @@ from types import SimpleNamespace
|
|||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
from litellm._service_logger import ServiceTypes
|
||||
from litellm.models.user import LiteLLM_UserTable
|
||||
|
|
@ -192,6 +194,62 @@ async def test_insert_data_team_serializes_members(prisma_client: PrismaClient)
|
|||
assert create_payload["team_id"] == "t1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "message"),
|
||||
[
|
||||
({"target": ["user_id"]}, "User with id requested-user already exists"),
|
||||
({"target": ["sso_user_id"]}, "User with this sso_user_id already exists"),
|
||||
(None, "User conflicts with an existing unique value"),
|
||||
({"target": ["unknown_field"]}, "User conflicts with an existing unique value"),
|
||||
],
|
||||
)
|
||||
async def test_user_conflict_names_constraint_and_does_not_retry(
|
||||
prisma_client: PrismaClient,
|
||||
metadata: Mapping[str, object] | None,
|
||||
message: str,
|
||||
) -> None:
|
||||
create: Final = AsyncMock(side_effect=UniqueViolationError({"user_facing_error": {"meta": metadata}}))
|
||||
prisma_client.db.litellm_usertable.create = create
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await prisma_client.insert_data(
|
||||
data={"user_id": "requested-user", "sso_user_id": "sso-subject"},
|
||||
table_name="user", allow_existing_user=False,
|
||||
)
|
||||
assert error.value.status_code == 409
|
||||
assert error.value.detail == {"error": message}
|
||||
assert create.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_upsert_recovers_from_concurrent_creation(
|
||||
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
row: Final = LiteLLM_UserTable(user_id="requested-user")
|
||||
conflict: Final = UniqueViolationError({"user_facing_error": {"meta": {"target": ["user_id"]}}})
|
||||
prisma_client.db.litellm_usertable.upsert = AsyncMock(side_effect=[conflict, row])
|
||||
monkeypatch.setattr("asyncio.sleep", AsyncMock())
|
||||
|
||||
result: Final = await prisma_client.insert_data(data={"user_id": row.user_id}, table_name="user")
|
||||
|
||||
assert result is row
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_insert_still_retries_transient_connection_failure(
|
||||
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
row: Final = LiteLLM_UserTable(user_id="requested-user")
|
||||
create: Final = AsyncMock(side_effect=[httpx.ConnectError("connection unavailable"), row])
|
||||
prisma_client.db.litellm_usertable.create = create
|
||||
monkeypatch.setattr("asyncio.sleep", AsyncMock())
|
||||
result: Final = await prisma_client.insert_data(
|
||||
data={"user_id": row.user_id}, table_name="user", allow_existing_user=False
|
||||
)
|
||||
assert result is row
|
||||
assert create.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_insert_data_user_organization_fk_raises_400(
|
||||
prisma_client: PrismaClient,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue