mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-13 18:23:20 +00:00
feat(proxy): include user_email in jwt upsert user creation (#22915)
* Include user_email in new user creation within get_user_object Enhance the get_user_object function to include user_email in the parameters when creating a new user. This change is accompanied by a new test to verify that user_email is correctly included during the upsert process. * Improve error handling in test_get_user_object by logging exceptions Updated the test_get_user_object_upsert_includes_user_email function to log exceptions when they occur, enhancing the visibility of potential issues during testing. This change helps in diagnosing failures related to the mock LiteLLM_UserTable.
This commit is contained in:
@@ -1316,6 +1316,8 @@ async def get_user_object(
|
||||
new_user_params: Dict[str, Any] = {
|
||||
"user_id": user_id,
|
||||
}
|
||||
if user_email is not None:
|
||||
new_user_params["user_email"] = user_email
|
||||
if litellm.default_internal_user_params is not None:
|
||||
new_user_params.update(litellm.default_internal_user_params)
|
||||
|
||||
|
||||
@@ -379,6 +379,65 @@ async def test_default_internal_user_params_with_get_user_object(monkeypatch):
|
||||
assert creation_args["user_role"] == "internal_user"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_object_upsert_includes_user_email():
|
||||
"""Test that user_email is included when creating a new user via get_user_object upsert"""
|
||||
# Mock the necessary dependencies
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_db = AsyncMock()
|
||||
mock_prisma_client.db = mock_db
|
||||
|
||||
# Set up the user creation mock
|
||||
mock_user = MagicMock()
|
||||
mock_user.user_id = "new_test_user"
|
||||
mock_user.user_email = "test@example.com"
|
||||
mock_user.models = []
|
||||
mock_user.max_budget = None
|
||||
mock_user.user_role = None
|
||||
mock_user.organization_memberships = []
|
||||
|
||||
mock_user.dict = lambda: {
|
||||
"user_id": "new_test_user",
|
||||
"user_email": "test@example.com",
|
||||
"models": [],
|
||||
"max_budget": None,
|
||||
"user_role": None,
|
||||
"organization_memberships": [],
|
||||
}
|
||||
|
||||
# Setup the mock returns - user does not exist
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
|
||||
mock_prisma_client.db.litellm_usertable.create = AsyncMock(return_value=mock_user)
|
||||
|
||||
# Create a mock cache
|
||||
mock_cache = MagicMock()
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
mock_cache.async_set_cache = AsyncMock()
|
||||
|
||||
# Call get_user_object with user_id_upsert=True and user_email
|
||||
try:
|
||||
await get_user_object(
|
||||
user_id="new_test_user",
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
user_id_upsert=True,
|
||||
proxy_logging_obj=None,
|
||||
user_email="test@example.com",
|
||||
)
|
||||
except Exception as e:
|
||||
# May fail since mock object is not a real LiteLLM_UserTable
|
||||
print(e)
|
||||
|
||||
# Verify the user was created with user_email included
|
||||
mock_prisma_client.db.litellm_usertable.create.assert_called_once()
|
||||
creation_args = mock_prisma_client.db.litellm_usertable.create.call_args[1]["data"]
|
||||
|
||||
assert "user_email" in creation_args, "user_email should be included when upserting a new user"
|
||||
assert creation_args["user_email"] == "test@example.com"
|
||||
assert creation_args["user_id"] == "new_test_user"
|
||||
|
||||
|
||||
def test_log_budget_lookup_failure_dry_run():
|
||||
"""Dry run: verify _log_budget_lookup_failure logs for schema/DB errors."""
|
||||
with patch("litellm.proxy.auth.auth_checks.verbose_proxy_logger") as mock_logger:
|
||||
|
||||
Reference in New Issue
Block a user