mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
test(identity): cover endpoint identity-invalidation wiring
Add mutation-quality regression tests asserting that key deletion (_delete_cache_key_object) invalidates the token identity entry, and that update_organization and delete_organization bump the org identity generation counter. These endpoint call sites previously had no test asserting they invoke invalidate_identity_*, and they were the only uncovered invalidation lines in the diff (team and user paths are already covered by existing endpoint tests).
This commit is contained in:
parent
01359d8eae
commit
e08b6c18d6
2 changed files with 132 additions and 0 deletions
|
|
@ -3714,3 +3714,31 @@ async def test_inference_route_still_enforces_team_budget():
|
|||
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cache_key_object_invalidates_identity_entry():
|
||||
"""Deleting a key's cache must also drop its identity-cache entry.
|
||||
|
||||
`_delete_cache_key_object` removes the legacy hashed-token key and then
|
||||
calls `invalidate_identity_for_token`. The identity entry lives under a
|
||||
different cache key (`identity:v1:<hash>`), so only the invalidation call
|
||||
can remove it; this pins that wiring.
|
||||
"""
|
||||
from litellm.identity.cache import IdentityCache
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_key_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
backend = UserApiKeyCache()
|
||||
identity_cache = IdentityCache(dual_cache=backend)
|
||||
uak = UserAPIKeyAuth(api_key="sk-delete-me", user_id="u1")
|
||||
await identity_cache.set(uak.token, uak)
|
||||
assert await identity_cache.get(uak.token) is not None
|
||||
|
||||
await _delete_cache_key_object(
|
||||
hashed_token=uak.token,
|
||||
user_api_key_cache=backend,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
assert await identity_cache.get(uak.token) is None
|
||||
|
|
|
|||
|
|
@ -695,3 +695,107 @@ async def test_organization_member_delete_rejects_unauthorized_caller(
|
|||
user_api_key_dict=unauthorized_caller,
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_organization_invalidates_identity_for_each_org():
|
||||
"""Deleting organizations must bump the identity generation for each id."""
|
||||
from litellm.proxy._types import (
|
||||
DeleteOrganizationRequest,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints import organization_endpoints as org_mod
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin", api_key="sk-admin"
|
||||
)
|
||||
|
||||
def _repo_with_async_table(**table_methods):
|
||||
repo = MagicMock()
|
||||
for name, ret in table_methods.items():
|
||||
setattr(repo.table, name, AsyncMock(return_value=ret))
|
||||
return MagicMock(return_value=repo)
|
||||
|
||||
invalidate = AsyncMock()
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
|
||||
patch.object(
|
||||
org_mod, "TeamRepository", _repo_with_async_table(delete_many=None)
|
||||
),
|
||||
patch.object(
|
||||
org_mod,
|
||||
"OrganizationMembershipRepository",
|
||||
_repo_with_async_table(delete_many=None),
|
||||
),
|
||||
patch.object(
|
||||
org_mod,
|
||||
"VerificationTokenRepository",
|
||||
_repo_with_async_table(delete_many=None),
|
||||
),
|
||||
patch.object(
|
||||
org_mod,
|
||||
"OrganizationRepository",
|
||||
_repo_with_async_table(delete=MagicMock(organization_id="org-del-1")),
|
||||
),
|
||||
patch(
|
||||
"litellm.identity.invalidation.invalidate_identity_for_org", new=invalidate
|
||||
),
|
||||
):
|
||||
await org_mod.delete_organization(
|
||||
data=DeleteOrganizationRequest(organization_ids=["org-del-1", "org-del-2"]),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
invalidated = sorted(c.kwargs["org_id"] for c in invalidate.await_args_list)
|
||||
assert invalidated == ["org-del-1", "org-del-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_organization_invalidates_identity_for_org():
|
||||
"""Updating an organization must bump its identity generation counter."""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints import organization_endpoints as org_mod
|
||||
|
||||
admin = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin", api_key="sk-admin"
|
||||
)
|
||||
|
||||
request = MagicMock(spec=Request)
|
||||
request.json = AsyncMock(
|
||||
return_value={"organization_id": "org-upd-9", "metadata": {"tier": "gold"}}
|
||||
)
|
||||
|
||||
existing = MagicMock()
|
||||
existing.metadata = {}
|
||||
existing.budget_id = None
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.jsonify_object = lambda d: d
|
||||
|
||||
org_repo = MagicMock()
|
||||
org_repo.table.find_unique = AsyncMock(return_value=existing)
|
||||
org_repo.table.update = AsyncMock(
|
||||
return_value=MagicMock(organization_id="org-upd-9")
|
||||
)
|
||||
|
||||
invalidate = AsyncMock()
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", fake_prisma),
|
||||
patch.object(org_mod, "_verify_org_access", new=AsyncMock()),
|
||||
patch.object(
|
||||
org_mod, "OrganizationRepository", MagicMock(return_value=org_repo)
|
||||
),
|
||||
patch(
|
||||
"litellm.identity.invalidation.invalidate_identity_for_org", new=invalidate
|
||||
),
|
||||
):
|
||||
await org_mod.update_organization(
|
||||
request=request,
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
assert invalidate.await_count == 1
|
||||
assert invalidate.await_args.kwargs["org_id"] == "org-upd-9"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue