mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
refactor(redis): migrate user MCP permission links to typed store
This commit is contained in:
parent
06affb432f
commit
e03e42fedd
7 changed files with 484 additions and 105 deletions
|
|
@ -3,16 +3,28 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from contextlib import suppress
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_batch import active_request_redis_batch
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.redis.stores.mcp_user_permission_links import (
|
||||
USER_NO_MCP_PERMISSION_SENTINEL as _USER_NO_PERMISSION,
|
||||
)
|
||||
from litellm.redis.stores.mcp_user_permission_links import (
|
||||
McpUserPermissionLinkStore,
|
||||
TransportMcpUserPermissionLinkStore,
|
||||
user_permission_link_key,
|
||||
)
|
||||
from litellm.redis.transport.contract import RedisFailure
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
|
@ -291,21 +303,64 @@ class UserApiKeyCache(DualCache):
|
|||
await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs)
|
||||
|
||||
|
||||
#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row,
|
||||
#: so a human without an entitlement costs no DB read per request. Lives beside the key builder
|
||||
#: because it is part of the same cache protocol: a reader that knows the key must know this value.
|
||||
USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__"
|
||||
_PERMISSION_LINK_VALUE: Final = TypeAdapter(object)
|
||||
USER_NO_MCP_PERMISSION_SENTINEL: Final = _USER_NO_PERMISSION
|
||||
|
||||
|
||||
def user_object_permission_id_cache_key(user_id: str) -> str:
|
||||
"""Cache key for the ``user_id -> object_permission_id`` link.
|
||||
return user_permission_link_key(user_id)
|
||||
|
||||
Lives here rather than next to either user because two modules own the two halves: the MCP auth
|
||||
resolver writes it on read, and ``/user/update`` deletes it after changing the link. A key format
|
||||
duplicated across those two drifts silently, and the failure is an entitlement change that never
|
||||
takes effect.
|
||||
"""
|
||||
return f"user_object_permission_id:{user_id}"
|
||||
|
||||
async def get_user_permission_link(cache: UserApiKeyCache, user_id: str) -> str | None:
|
||||
key: Final = user_permission_link_key(user_id)
|
||||
local: Final = _PERMISSION_LINK_VALUE.validate_python(await cache.async_get_cache(key=key, local_only=True))
|
||||
if local is not None:
|
||||
return local if isinstance(local, str) else None
|
||||
if cache.redis_cache is None:
|
||||
return None
|
||||
batch: Final = active_request_redis_batch(cache.redis_cache)
|
||||
if batch is not None and batch.read_as_missing(key):
|
||||
return None
|
||||
with suppress(Exception):
|
||||
store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore(cache.redis_cache.transport())
|
||||
result: Final = await store.get(user_id)
|
||||
if isinstance(result, RedisFailure):
|
||||
return None
|
||||
if result is not None:
|
||||
await cache.async_set_cache(key=key, value=result, local_only=True)
|
||||
return result
|
||||
return None
|
||||
|
||||
|
||||
async def fill_user_permission_link(
|
||||
cache: UserApiKeyCache, user_id: str, permission_id: str | None, *, ttl: float
|
||||
) -> None:
|
||||
await cache.async_set_cache(
|
||||
key=user_permission_link_key(user_id),
|
||||
value=permission_id or USER_NO_MCP_PERMISSION_SENTINEL,
|
||||
ttl=ttl,
|
||||
local_only=True,
|
||||
)
|
||||
if cache.redis_cache is not None:
|
||||
with suppress(Exception):
|
||||
store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore(
|
||||
cache.redis_cache.transport()
|
||||
)
|
||||
_ = await store.fill(user_id, permission_id, ttl=timedelta(seconds=int(ttl)))
|
||||
|
||||
|
||||
async def invalidate_user_permission_link(cache: UserApiKeyCache, user_id: str) -> None:
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
|
||||
key: Final = user_permission_link_key(user_id)
|
||||
cache.in_memory_cache.delete_cache(key)
|
||||
if cache.redis_cache is not None:
|
||||
with suppress(Exception):
|
||||
store: Final[McpUserPermissionLinkStore] = TransportMcpUserPermissionLinkStore(
|
||||
cache.redis_cache.transport()
|
||||
)
|
||||
_ = await store.invalidate(user_id)
|
||||
await publish_auth_cache_invalidation(cache_key=key)
|
||||
|
||||
|
||||
def object_permission_cache_key(object_permission_id: str) -> str:
|
||||
|
|
|
|||
|
|
@ -50,8 +50,8 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_
|
|||
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
invalidate_user_permission_link,
|
||||
object_permission_cache_key,
|
||||
user_object_permission_id_cache_key,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
|
|
@ -1481,9 +1481,11 @@ async def _invalidate_cached_user_entitlement(user_id: str | None, object_permis
|
|||
|
||||
keys: Final = (
|
||||
*(object_permission_cache_key(permission_id) for permission_id in dict.fromkeys(object_permission_ids)),
|
||||
*((user_object_permission_id_cache_key(user_id), user_id) if user_id is not None else ()),
|
||||
)
|
||||
await evict_and_broadcast(cache_keys=keys, user_api_key_cache=user_api_key_cache)
|
||||
if user_id is not None:
|
||||
await invalidate_user_permission_link(user_api_key_cache, user_id)
|
||||
await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache)
|
||||
|
||||
|
||||
async def _update_single_user_helper(
|
||||
|
|
|
|||
|
|
@ -61,8 +61,9 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
|
|||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
USER_NO_MCP_PERMISSION_SENTINEL,
|
||||
fill_user_permission_link,
|
||||
get_management_object_ttl,
|
||||
user_object_permission_id_cache_key,
|
||||
get_user_permission_link,
|
||||
)
|
||||
from litellm.proxy.mcp_server.catalog import global_manager
|
||||
from litellm.proxy.mcp_server.oauth_utils import (
|
||||
|
|
@ -4149,9 +4150,8 @@ class MCPRequestHandler:
|
|||
from litellm.proxy.auth.auth_checks import get_user_object
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
cache_key: Final = user_object_permission_id_cache_key(user_id)
|
||||
try:
|
||||
cached: Final[object] = None if check_db_only else await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached: Final = None if check_db_only else await get_user_permission_link(user_api_key_cache, user_id)
|
||||
if cached == USER_NO_MCP_PERMISSION_SENTINEL:
|
||||
return None
|
||||
if isinstance(cached, str) and cached:
|
||||
|
|
@ -4165,9 +4165,10 @@ class MCPRequestHandler:
|
|||
)
|
||||
linked: Final[object] = getattr(user_row, "object_permission_id", None) if user_row is not None else None
|
||||
object_permission_id: Final = linked if isinstance(linked, str) and linked else None
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=object_permission_id or USER_NO_MCP_PERMISSION_SENTINEL,
|
||||
await fill_user_permission_link(
|
||||
user_api_key_cache,
|
||||
user_id,
|
||||
object_permission_id,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
)
|
||||
return object_permission_id
|
||||
|
|
|
|||
64
litellm/redis/stores/mcp_user_permission_links.py
Normal file
64
litellm/redis/stores/mcp_user_permission_links.py
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
import ast
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.redis.transport.contract import RedisFailure, RedisTransport
|
||||
|
||||
USER_NO_MCP_PERMISSION_SENTINEL: Final = "__user_no_mcp_permission__"
|
||||
_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
def user_permission_link_key(user_id: str) -> str:
|
||||
return f"user_object_permission_id:{user_id}"
|
||||
|
||||
|
||||
class McpUserPermissionLinkStore(Protocol):
|
||||
async def get(self, user_id: str) -> str | None | RedisFailure: ...
|
||||
|
||||
async def fill(self, user_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure: ...
|
||||
|
||||
async def invalidate(self, user_id: str) -> int | RedisFailure: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TransportMcpUserPermissionLinkStore:
|
||||
transport: RedisTransport
|
||||
|
||||
async def get(self, user_id: str) -> str | None | RedisFailure:
|
||||
reply: Final = await self.transport.get(user_permission_link_key(user_id))
|
||||
if reply is None or isinstance(reply, RedisFailure):
|
||||
return reply
|
||||
try:
|
||||
decoded: Final = reply.decode("utf-8")
|
||||
try:
|
||||
value: Final = _VALUE.validate_json(decoded)
|
||||
except ValidationError:
|
||||
literal: Final = _VALUE.validate_python(ast.literal_eval(decoded))
|
||||
return _permission_id(literal)
|
||||
return _permission_id(value)
|
||||
except (UnicodeDecodeError, ValueError, SyntaxError):
|
||||
return RedisFailure(kind="decode", message="Invalid user permission link")
|
||||
|
||||
async def fill(self, user_id: str, permission_id: str | None, *, ttl: timedelta) -> bool | None | RedisFailure:
|
||||
seconds: Final = int(ttl.total_seconds())
|
||||
key: Final = user_permission_link_key(user_id)
|
||||
value: Final = json.dumps(permission_id or USER_NO_MCP_PERMISSION_SENTINEL)
|
||||
expiry: Final = timedelta(seconds=seconds)
|
||||
batch: Final = self.transport.request_batch()
|
||||
if batch is not None:
|
||||
_ = batch.set_seconds(key, value, ttl=expiry)
|
||||
return None
|
||||
return await self.transport.set_seconds(key, value, ttl=expiry)
|
||||
|
||||
async def invalidate(self, user_id: str) -> int | RedisFailure:
|
||||
return await self.transport.delete(user_permission_link_key(user_id))
|
||||
|
||||
|
||||
def _permission_id(value: object) -> str | RedisFailure:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return RedisFailure(kind="decode", message="User permission link is not a string")
|
||||
|
|
@ -9,6 +9,7 @@ from types import SimpleNamespace
|
|||
from typing import TYPE_CHECKING, Final, cast
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import fakeredis
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -41,6 +42,7 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
|||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import (
|
||||
CascadingJWTMappingTable,
|
||||
JWTMappingRow,
|
||||
|
|
@ -870,7 +872,6 @@ async def test_new_user_license_over_limit(mocker):
|
|||
"""
|
||||
Test that /user/new endpoint raises an error when license is over the user limit
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import NewUserRequest, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
|
||||
|
|
@ -1389,9 +1390,7 @@ async def test_user_info_url_encoding_plus_character(mocker):
|
|||
|
||||
# Mock get_data to return user when called with user_id, empty list for keys
|
||||
async def mock_get_data(*args, **kwargs):
|
||||
if kwargs.get("table_name") == "key":
|
||||
return []
|
||||
elif kwargs.get("table_name") == "team":
|
||||
if kwargs.get("table_name") == "key" or kwargs.get("table_name") == "team":
|
||||
return []
|
||||
elif kwargs.get("user_id") is not None:
|
||||
return mock_user
|
||||
|
|
@ -1524,7 +1523,7 @@ async def test_new_user_default_teams_flow(mocker):
|
|||
- Teams ARE sent to _add_user_to_team
|
||||
"""
|
||||
import litellm
|
||||
from litellm.proxy._types import NewUserRequest, NewUserRequestTeam, UserAPIKeyAuth
|
||||
from litellm.proxy._types import NewUserRequest, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
|
||||
|
||||
# Mock the prisma client
|
||||
|
|
@ -1672,9 +1671,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:
|
||||
|
|
@ -1884,8 +1881,6 @@ async def test_check_duplicate_user_email_case_insensitive(mocker):
|
|||
assert user_email_clause.get("mode") == "insensitive"
|
||||
assert user_email_clause.get("equals") == "newuser@example.com"
|
||||
|
||||
return None # No existing user found
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate
|
||||
|
||||
# Should not raise any exception when no duplicate is found
|
||||
|
|
@ -1938,7 +1933,6 @@ async def test_check_duplicate_user_id(mocker):
|
|||
user_id_clause = where_clause.get("user_id", {})
|
||||
assert user_id_clause.get("equals") == "new-user-id"
|
||||
assert "mode" not in user_id_clause
|
||||
return None
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.find_first = mock_find_first_no_duplicate
|
||||
|
||||
|
|
@ -2301,7 +2295,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 +2341,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 +2384,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 +2425,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 +2464,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 +2479,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 +2891,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 +2989,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 +3170,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)
|
||||
|
|
@ -3556,7 +3564,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker):
|
|||
"""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, UserInfoV2Response
|
||||
from litellm.proxy._types import UserInfoV2Response
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
|
@ -4419,8 +4427,7 @@ async def test_user_update_persists_mcp_entitlement_and_links_it(mocker):
|
|||
)
|
||||
|
||||
mock_prisma_client = _object_permission_mocks(mocker)
|
||||
cache = mocker.MagicMock()
|
||||
cache.async_delete_cache = mocker.AsyncMock()
|
||||
cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements"))
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
|
||||
await _update_single_user_helper(
|
||||
|
|
@ -4456,8 +4463,10 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker):
|
|||
)
|
||||
|
||||
_object_permission_mocks(mocker)
|
||||
cache = mocker.MagicMock()
|
||||
cache.async_delete_cache = mocker.AsyncMock()
|
||||
cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements"))
|
||||
keys: Final = ("object_permission_id:perm-new", "user_object_permission_id:target-user", "target-user")
|
||||
for key in keys:
|
||||
await cache.async_set_cache(key, "stale", ttl=60)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
|
||||
await _update_single_user_helper(
|
||||
|
|
@ -4468,12 +4477,8 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker):
|
|||
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list}
|
||||
assert deleted == {
|
||||
"object_permission_id:perm-new",
|
||||
"user_object_permission_id:target-user",
|
||||
"target-user",
|
||||
}
|
||||
for key in keys:
|
||||
assert await cache.async_get_cache(key) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4483,9 +4488,10 @@ 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
|
||||
cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements"))
|
||||
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,
|
||||
|
|
@ -4520,8 +4526,10 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker):
|
|||
)
|
||||
|
||||
mock_prisma_client = _object_permission_mocks(mocker, "perm-existing")
|
||||
cache = mocker.MagicMock()
|
||||
cache.async_delete_cache = mocker.AsyncMock()
|
||||
cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements"))
|
||||
keys: Final = ("object_permission_id:perm-existing", "user_object_permission_id:target-user", "target-user")
|
||||
for key in keys:
|
||||
await cache.async_set_cache(key, "stale", ttl=60)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
|
||||
await _update_single_user_helper(
|
||||
|
|
@ -4534,12 +4542,8 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker):
|
|||
assert "object_permission" not in written
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called()
|
||||
|
||||
deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list}
|
||||
assert deleted == {
|
||||
"object_permission_id:perm-existing",
|
||||
"user_object_permission_id:target-user",
|
||||
"target-user",
|
||||
}
|
||||
for key in keys:
|
||||
assert await cache.async_get_cache(key) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4554,8 +4558,15 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock
|
|||
)
|
||||
|
||||
_object_permission_mocks(mocker, "perm-existing")
|
||||
cache = mocker.MagicMock()
|
||||
cache.async_delete_cache = mocker.AsyncMock()
|
||||
cache: Final = UserApiKeyCache(redis_cache=fake_redis_cache(fakeredis.FakeServer(), namespace="entitlements"))
|
||||
keys: Final = (
|
||||
"object_permission_id:perm-existing",
|
||||
"object_permission_id:perm-new",
|
||||
"user_object_permission_id:target-user",
|
||||
"target-user",
|
||||
)
|
||||
for key in keys:
|
||||
await cache.async_set_cache(key, "stale", ttl=60)
|
||||
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
|
||||
await _update_single_user_helper(
|
||||
|
|
@ -4566,13 +4577,8 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock
|
|||
user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list}
|
||||
assert deleted == {
|
||||
"object_permission_id:perm-existing",
|
||||
"object_permission_id:perm-new",
|
||||
"user_object_permission_id:target-user",
|
||||
"target-user",
|
||||
}
|
||||
for key in keys:
|
||||
assert await cache.async_get_cache(key) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4755,7 +4761,7 @@ async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget,
|
|||
async def get_data(self, *args, **kwargs):
|
||||
return None
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", _FakePrisma(), raising=False)
|
||||
# model_max_budget is an enterprise feature; without this the call is rejected
|
||||
|
|
@ -5032,12 +5038,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,
|
||||
|
|
@ -5143,7 +5155,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"]
|
||||
|
||||
|
||||
|
|
@ -5301,7 +5315,9 @@ async def test_user_update_without_password_revokes_nothing(_admin_prisma, mocke
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(_admin_prisma: MagicMock) -> None:
|
||||
async def test_user_update_reports_a_bad_budget_duration_as_a_request_error_not_an_auth_error(
|
||||
_admin_prisma: MagicMock,
|
||||
) -> None:
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from types import SimpleNamespace
|
|||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
|
@ -26,6 +27,7 @@ from litellm.proxy._types import (
|
|||
SpecialMCPServerNames,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
UnloadableEntitlementError,
|
||||
|
|
@ -33,6 +35,48 @@ from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import (
|
|||
is_mcp_admitted_user_subject,
|
||||
)
|
||||
from litellm.types.agents import AgentCaller
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("permission_id", ["policy", None])
|
||||
async def test_user_permission_link_fills_from_cached_user_row_and_survives_redis_outage(
|
||||
monkeypatch: pytest.MonkeyPatch, permission_id: str | None
|
||||
) -> None:
|
||||
server: Final = fakeredis.FakeServer()
|
||||
redis: Final = fake_redis_cache(server, namespace="auth")
|
||||
cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=41)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
await cache.async_set_cache(
|
||||
"link-user", LiteLLM_UserTable(user_id="link-user", object_permission_id=permission_id), local_only=True
|
||||
)
|
||||
prisma: Final = MagicMock()
|
||||
assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) == permission_id
|
||||
assert await redis.async_get_cache("user_object_permission_id:link-user") == (
|
||||
permission_id or "__user_no_mcp_permission__"
|
||||
)
|
||||
cache.in_memory_cache.delete_cache("user_object_permission_id:link-user")
|
||||
server.connected = False
|
||||
assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) == permission_id
|
||||
assert await cache.async_get_cache("user_object_permission_id:link-user", local_only=True) == (
|
||||
permission_id or "__user_no_mcp_permission__"
|
||||
)
|
||||
prisma.db.litellm_usertable.find_unique.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fresh_user_policy_does_not_use_redis_sentinel(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
redis: Final = fake_redis_cache(fakeredis.FakeServer(), namespace="auth")
|
||||
cache: Final = UserApiKeyCache(redis_cache=redis)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
|
||||
await redis.async_set_cache("user_object_permission_id:link-user", "__user_no_mcp_permission__", ttl=60)
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable"))
|
||||
assert await MCPRequestHandler._user_object_permission_id("link-user", prisma) is None # pyright: ignore[reportPrivateUsage] # resolver internals under test
|
||||
prisma.db.litellm_usertable.find_unique.assert_not_called()
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await MCPRequestHandler._user_object_permission_id("link-user", prisma, check_db_only=True) # pyright: ignore[reportPrivateUsage] # resolver internals under test
|
||||
assert error.value.status_code == 503
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -11309,11 +11353,11 @@ async def test_fresh_mcp_user_permission_link_ignores_cached_and_replica_grants(
|
|||
database.writer_db.litellm_usertable.find_unique = AsyncMock(return_value=current)
|
||||
database.db.litellm_usertable.find_unique = AsyncMock(return_value=cached)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
|
||||
assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current"
|
||||
assert await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) == "current" # pyright: ignore[reportPrivateUsage] # resolver internals under test
|
||||
database.db.litellm_usertable.find_unique.assert_not_awaited()
|
||||
database.writer_db.litellm_usertable.find_unique.side_effect = failure
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True)
|
||||
await MCPRequestHandler._user_object_permission_id("fresh-human", database, check_db_only=True) # pyright: ignore[reportPrivateUsage] # resolver internals under test
|
||||
assert denied.value.status_code == 503
|
||||
|
||||
|
||||
|
|
@ -12451,7 +12495,6 @@ async def test_personal_key_display_name_billing_uses_original_tool_grant(
|
|||
async def test_personal_key_display_name_billing_uses_allowed_servers(
|
||||
monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.auth import auth_checks
|
||||
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
from litellm.types.mcp import MCPTransport
|
||||
|
|
|
|||
198
tests/unit/redis/stores/test_mcp_user_permission_links.py
Normal file
198
tests/unit/redis/stores/test_mcp_user_permission_links.py
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
from datetime import timedelta
|
||||
from typing import Final
|
||||
|
||||
import fakeredis
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from pytest import MonkeyPatch
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
|
||||
from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
fill_user_permission_link,
|
||||
get_user_permission_link,
|
||||
invalidate_user_permission_link,
|
||||
)
|
||||
from litellm.redis.stores.mcp_user_permission_links import TransportMcpUserPermissionLinkStore
|
||||
from litellm.redis.transport.contract import RedisFailure
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(loop_scope="function")
|
||||
async def link_cache(monkeypatch: MonkeyPatch) -> tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache]:
|
||||
monkeypatch.setattr("time.time", lambda: 1000.0)
|
||||
server: Final = fakeredis.FakeServer()
|
||||
redis: Final = fake_redis_cache(server, namespace="links")
|
||||
cache: Final = UserApiKeyCache(redis_cache=redis, default_in_memory_ttl=37.9)
|
||||
return server, redis, cache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", ['policy é "quoted"', "__user_no_mcp_permission__", ""])
|
||||
async def test_old_writer_new_reader(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str
|
||||
) -> None:
|
||||
_, redis, cache = link_cache
|
||||
await redis.async_set_cache("user_object_permission_id:u", value, ttl=37.9)
|
||||
assert await get_user_permission_link(cache, "u") == value
|
||||
assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) == value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"value,wire",
|
||||
[(None, b'"__user_no_mcp_permission__"'), ('policy é "quoted"', b'"policy \\u00e9 \\"quoted\\""')],
|
||||
)
|
||||
@pytest.mark.parametrize("ttl", [37.9, 37.999999999])
|
||||
async def test_new_writer_old_reader_and_whole_second_expiry(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], value: str | None, wire: bytes, ttl: float
|
||||
) -> None:
|
||||
server, redis, cache = link_cache
|
||||
raw: Final = fakeredis.FakeAsyncRedis(server=server)
|
||||
await fill_user_permission_link(cache, "u", value, ttl=ttl)
|
||||
assert await raw.get("links:user_object_permission_id:u") == wire
|
||||
assert await raw.pttl("links:user_object_permission_id:u") == 37000
|
||||
assert await redis.async_get_cache("user_object_permission_id:u") == (value or "__user_no_mcp_permission__")
|
||||
assert await raw.keys("*") == [b"links:user_object_permission_id:u"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("wire", [b"broken", b"42", b"null", b"{}", b"[]", b"\xff"])
|
||||
async def test_bad_payload_is_decode_failure_and_caller_cache_miss(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], wire: bytes
|
||||
) -> None:
|
||||
server, redis, cache = link_cache
|
||||
await fakeredis.FakeAsyncRedis(server=server).set("links:user_object_permission_id:u", wire)
|
||||
result: Final = await TransportMcpUserPermissionLinkStore(redis.transport()).get("u")
|
||||
assert isinstance(result, RedisFailure) and result.kind == "decode"
|
||||
assert await get_user_permission_link(cache, "u") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_literal_legacy_string_and_local_non_string(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache],
|
||||
) -> None:
|
||||
server, _, cache = link_cache
|
||||
await fakeredis.FakeAsyncRedis(server=server).set("links:user_object_permission_id:u", b"'legacy'")
|
||||
assert await get_user_permission_link(cache, "u") == "legacy"
|
||||
await cache.async_set_cache("user_object_permission_id:u", {"invalid": True}, local_only=True)
|
||||
assert await get_user_permission_link(cache, "u") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_outage_preserves_local_fill_and_invalidation(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache],
|
||||
) -> None:
|
||||
server, redis, cache = link_cache
|
||||
server.connected = False
|
||||
assert await get_user_permission_link(cache, "u") is None
|
||||
await fill_user_permission_link(cache, "u", "policy", ttl=30)
|
||||
assert await get_user_permission_link(cache, "u") == "policy"
|
||||
await invalidate_user_permission_link(cache, "u")
|
||||
assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) is None
|
||||
store: Final = TransportMcpUserPermissionLinkStore(redis.transport())
|
||||
for _ in range(10):
|
||||
assert isinstance(await store.get("u"), RedisFailure)
|
||||
result: Final = await store.get("u")
|
||||
assert isinstance(result, RedisFailure) and result.kind == "breaker_open"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("ttl", [37.9, 37.999999999])
|
||||
async def test_request_fill_stays_queued_and_clears_prior_miss(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float
|
||||
) -> None:
|
||||
server, redis, cache = link_cache
|
||||
raw: Final = fakeredis.FakeAsyncRedis(server=server)
|
||||
with request_redis_batch_scope():
|
||||
batch: Final = active_request_redis_batch(redis)
|
||||
assert batch is not None
|
||||
assert await batch.mget(("user_object_permission_id:u",)) == {"user_object_permission_id:u": None}
|
||||
await redis.async_set_cache("user_object_permission_id:u", "external", ttl=60)
|
||||
assert await get_user_permission_link(cache, "u") is None
|
||||
await fill_user_permission_link(cache, "u", "policy", ttl=ttl)
|
||||
assert await get_user_permission_link(cache, "u") == "policy"
|
||||
assert await raw.get("links:user_object_permission_id:u") == b'"external"'
|
||||
await batch.flush()
|
||||
assert await raw.get("links:user_object_permission_id:u") == b'"policy"'
|
||||
assert await raw.pttl("links:user_object_permission_id:u") == 37000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("ttl", [0.9, 0.999999999, 0, -1])
|
||||
async def test_rejected_expiry_keeps_existing_redis_value_and_new_local_value(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float
|
||||
) -> None:
|
||||
_, redis, cache = link_cache
|
||||
await redis.async_set_cache("user_object_permission_id:u", "old", ttl=60)
|
||||
await fill_user_permission_link(cache, "u", "new", ttl=ttl)
|
||||
assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) == ("new" if ttl >= 0 else None)
|
||||
assert await redis.async_get_cache("user_object_permission_id:u") == "old"
|
||||
result: Final = await TransportMcpUserPermissionLinkStore(redis.transport()).fill(
|
||||
"u", "new", ttl=timedelta(seconds=int(ttl))
|
||||
)
|
||||
assert isinstance(result, RedisFailure) and result.kind == "server"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("ttl", [0.999999999, 0, -1])
|
||||
async def test_rejected_batched_expiry_keeps_existing_value(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], ttl: float
|
||||
) -> None:
|
||||
_, redis, cache = link_cache
|
||||
await redis.async_set_cache("user_object_permission_id:u", "old", ttl=60)
|
||||
with request_redis_batch_scope() as scope:
|
||||
await fill_user_permission_link(cache, "u", "new", ttl=ttl)
|
||||
await scope.flush_all()
|
||||
assert await redis.async_get_cache("user_object_permission_id:u") == "old"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unclassified_client_failure_keeps_cache_optional(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch
|
||||
) -> None:
|
||||
_, redis, cache = link_cache
|
||||
client: Final = redis.init_async_client()
|
||||
assert isinstance(client, fakeredis.FakeAsyncRedis)
|
||||
await client.ping()
|
||||
connection: Final = await client.connection_pool.get_connection()
|
||||
read_response: Final = connection.read_response
|
||||
|
||||
async def unavailable_reply(**kwargs: object) -> object:
|
||||
await read_response(**kwargs)
|
||||
raise RuntimeError("client callback failed")
|
||||
|
||||
monkeypatch.setattr(connection, "read_response", unavailable_reply)
|
||||
await client.connection_pool.release(connection)
|
||||
assert await cache.async_get_cache("user_object_permission_id:u") is None
|
||||
assert await get_user_permission_link(cache, "u") is None
|
||||
await fill_user_permission_link(cache, "u", "policy", ttl=30)
|
||||
assert await get_user_permission_link(cache, "u") == "policy"
|
||||
await invalidate_user_permission_link(cache, "u")
|
||||
assert await cache.async_get_cache("user_object_permission_id:u", local_only=True) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lost_set_reply_does_not_roll_back_applied_write(
|
||||
link_cache: tuple[fakeredis.FakeServer, RedisCache, UserApiKeyCache], monkeypatch: MonkeyPatch
|
||||
) -> None:
|
||||
server, redis, cache = link_cache
|
||||
client: Final = redis.init_async_client()
|
||||
assert isinstance(client, fakeredis.FakeAsyncRedis)
|
||||
await client.ping()
|
||||
connection: Final = await client.connection_pool.get_connection()
|
||||
read_response: Final = connection.read_response
|
||||
|
||||
async def lost_reply(**kwargs: object) -> object:
|
||||
await read_response(**kwargs)
|
||||
raise RedisConnectionError("reply lost after write")
|
||||
|
||||
monkeypatch.setattr(connection, "read_response", lost_reply)
|
||||
await client.connection_pool.release(connection)
|
||||
await fill_user_permission_link(cache, "u", "applied", ttl=37.9)
|
||||
raw: Final = fakeredis.FakeAsyncRedis(server=server)
|
||||
assert await raw.get("links:user_object_permission_id:u") == b'"applied"'
|
||||
assert await raw.pttl("links:user_object_permission_id:u") == 37000
|
||||
assert await get_user_permission_link(cache, "u") == "applied"
|
||||
Loading…
Add table
Reference in a new issue