mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
style: ruff-format the internal-user tests
Generated-by: termag
This commit is contained in:
JackDanger
Jack Danger
parent
606ead50bb
commit
6d9d6d9eb5
1 changed files with 69 additions and 54 deletions
|
|
@ -1672,9 +1672,7 @@ async def test_new_user_clears_recent_missing_user_lookup(mocker: MockerFixture)
|
|||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn",
|
||||
new=mocker.AsyncMock(
|
||||
return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None}
|
||||
),
|
||||
new=mocker.AsyncMock(return_value={"user_id": user_id, "token": "sk-sso-created-user", "expires": None}),
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -2301,7 +2299,11 @@ def test_update_internal_user_params_preserves_model_budget_presence_and_neighbo
|
|||
"user_alias": "Spruce",
|
||||
}
|
||||
|
||||
replacement: Final = {"model-spruce": {"budget_limit": 0, "time_period": "1d"}, "model-birch": 5.0, "model-cedar": 0}
|
||||
replacement: Final = {
|
||||
"model-spruce": {"budget_limit": 0, "time_period": "1d"},
|
||||
"model-birch": 5.0,
|
||||
"model-cedar": 0,
|
||||
}
|
||||
request: Final = UpdateUserRequest(
|
||||
user_id="user-spruce",
|
||||
model_max_budget=replacement,
|
||||
|
|
@ -2343,10 +2345,14 @@ async def test_user_model_budget_update_by_email_refreshes_cached_user(mocker: M
|
|||
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
|
||||
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
|
||||
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", prisma_client
|
||||
) # test-quality-ok: substitute the database dependency
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: exercise a real isolated cache
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
|
|
@ -2382,10 +2388,14 @@ async def test_user_status_update_refreshes_cached_user(
|
|||
prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user)
|
||||
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
|
||||
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", prisma_client
|
||||
) # test-quality-ok: substitute the database dependency
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: exercise a real isolated cache
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
|
|
@ -2419,10 +2429,14 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
|
|||
prisma_client: Final = mocker.MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user])
|
||||
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", prisma_client
|
||||
) # test-quality-ok: substitute the database dependency
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: exercise a real isolated cache
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
|
|
@ -2454,7 +2468,9 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
|
||||
async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: MockerFixture, all_users: bool) -> None:
|
||||
async def test_user_max_budget_update_evicts_cached_user_on_every_worker(
|
||||
mocker: MockerFixture, all_users: bool
|
||||
) -> None:
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper, bulk_user_update
|
||||
|
|
@ -2467,10 +2483,14 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker:
|
|||
prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1)
|
||||
prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user])
|
||||
prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user})
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", prisma_client
|
||||
) # test-quality-ok: substitute the database dependency
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: exercise a real isolated cache
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
|
|
@ -2875,24 +2895,12 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker):
|
|||
mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[])
|
||||
|
||||
# Mock all delete_many calls
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0)
|
||||
mock_prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=1)
|
||||
mock_prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0)
|
||||
mock_prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
|
||||
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
|
||||
_serve_locked_user_delete(mock_prisma_client, mock_user_row)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
|
@ -2985,9 +2993,15 @@ async def test_delete_user_evicts_jwt_key_mapping_cache_of_its_keys(mocker):
|
|||
mock_prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
|
||||
_serve_locked_user_delete(mock_prisma_client, user_row)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
||||
) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", None
|
||||
) # test-quality-ok: delete_user reads it off proxy_server at call time
|
||||
|
||||
await delete_user(
|
||||
data=DeleteUserRequest(user_ids=["jwt-user"]),
|
||||
|
|
@ -3160,9 +3174,7 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker: MockerFixture
|
|||
return None
|
||||
|
||||
mock_prisma_client: Final = _PrismaClient(
|
||||
db=_Database(
|
||||
litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique))
|
||||
)
|
||||
db=_Database(litellm_usertable=_UserTable(find_unique=mocker.AsyncMock(side_effect=mock_find_unique)))
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
|
@ -4398,7 +4410,9 @@ async def test_user_update_broadcasts_the_entitlement_invalidation_to_other_work
|
|||
_object_permission_mocks(mocker)
|
||||
cache: Final = mocker.MagicMock()
|
||||
cache.async_delete_cache = mocker.AsyncMock()
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: substitute the cache dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: substitute the cache dependency
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
|
|
@ -4945,12 +4959,18 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non
|
|||
prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0)
|
||||
prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1)
|
||||
_serve_locked_user_delete(prisma_client, deleted)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", prisma_client
|
||||
) # test-quality-ok: substitute the database dependency
|
||||
cache: Final = UserApiKeyCache()
|
||||
for row in (deleted, survivor):
|
||||
await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", cache
|
||||
) # test-quality-ok: exercise a real isolated cache
|
||||
mocker.patch(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", None
|
||||
) # test-quality-ok: delete_user reads it off proxy_server at call time
|
||||
broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
|
|
@ -5056,7 +5076,9 @@ async def test_ui_view_users_answers_a_db_outage_as_503_no_db_connection_not_as_
|
|||
assert raised.value.code == "503"
|
||||
assert raised.value.type == ProxyErrorTypes.no_db_connection
|
||||
assert isinstance(raised.value.__cause__, httpx.ConnectError)
|
||||
outage_logs: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()]
|
||||
outage_logs: Final = [
|
||||
r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and "ConnectError" in r.getMessage()
|
||||
]
|
||||
assert outage_logs == ["Database unavailable during user search: ConnectError"]
|
||||
|
||||
|
||||
|
|
@ -5242,9 +5264,7 @@ def test_user_update_self_write_of_ceiling_fields_rejected(field, value):
|
|||
_reject_ceiling_self_update(
|
||||
sent_fields=set(data.fields_set()),
|
||||
non_default_values={field: value},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="self", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="self", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
assert int(getattr(exc.value, "status_code", 0)) == 403
|
||||
assert f"'{field}'" in str(exc.value.detail)
|
||||
|
|
@ -5261,9 +5281,7 @@ def test_user_update_self_write_of_innocuous_field_allowed():
|
|||
_reject_ceiling_self_update(
|
||||
sent_fields=set(data.fields_set()),
|
||||
non_default_values={"user_alias": "me"},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="self", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="self", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -5278,9 +5296,7 @@ def test_user_update_admin_self_write_unrestricted():
|
|||
_reject_ceiling_self_update(
|
||||
sent_fields=set(data.fields_set()),
|
||||
non_default_values={"rpm_limit": 1000000, "models": []},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -5299,6 +5315,7 @@ async def test_user_update_self_write_of_membership_rejected():
|
|||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
_reject_ceiling_self_update,
|
||||
)
|
||||
|
||||
for field, value in (
|
||||
("teams", ["victim-team"]),
|
||||
("organizations", ["victim-org"]),
|
||||
|
|
@ -5308,9 +5325,7 @@ async def test_user_update_self_write_of_membership_rejected():
|
|||
_reject_ceiling_self_update(
|
||||
sent_fields=set(data.fields_set()),
|
||||
non_default_values={field: value},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_id="self", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="self", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
assert int(getattr(exc.value, "status_code", 0)) == 403
|
||||
assert f"'{field}'" in str(exc.value.detail)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue