mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
test(auth): freeze prefetch cache clock so the org getter cannot miss on a slow runner
Backport of f4e21430c0 (originally on main via PR #41102). The org cache
entries written by prefetch_auth_objects carry the 5s DEFAULT_IN_MEMORY_TTL,
so on a slow CI runner the entry expires before get_org_object reads it and
the getter falls through to the MagicMock database, producing a
TypeError: object MagicMock can't be used in 'await' expression at
tests/proxy_behavior/auth/test_auth_object_prefetch.py::test_join_binds_the_membership_to_the_requested_team.
Inject a frozen clock into InMemoryCache so the test asserts the SQL join,
not import or scheduler latency.
Co-authored-by: Krrish Dholakia <krrish-berri-2@users.noreply.github.com>
This commit is contained in:
parent
af6dc1db08
commit
2ae1cf608f
1 changed files with 9 additions and 3 deletions
|
|
@ -22,6 +22,11 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|||
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
||||
|
||||
|
||||
def _frozen_cache() -> UserApiKeyCache:
|
||||
"""The org entries carry a 5s TTL; a frozen clock keeps a slow first call from expiring them mid-test."""
|
||||
return UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1_000_000.0), redis_cache=None)
|
||||
|
||||
|
||||
def _dead_db() -> MagicMock:
|
||||
prisma = MagicMock(name="prisma_client")
|
||||
prisma.db.query_first = AsyncMock(return_value=None)
|
||||
|
|
@ -58,9 +63,10 @@ async def test_join_binds_the_membership_to_the_requested_team(prisma):
|
|||
data={"user_id": user_id, "team_id": team_b, "litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}}}
|
||||
)
|
||||
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None)
|
||||
cache = _frozen_cache()
|
||||
refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id)
|
||||
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
|
||||
assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None
|
||||
|
||||
dead_db = _dead_db()
|
||||
membership = await get_team_membership(
|
||||
|
|
@ -100,7 +106,7 @@ async def test_join_reads_team_model_aliases_from_the_mapped_column(prisma):
|
|||
where={"team_id": team_id}, include={"litellm_model_table": True}
|
||||
)
|
||||
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None)
|
||||
cache = _frozen_cache()
|
||||
refs = AuthObjectRefs(user_id=None, team_id=team_id, membership_user_id=None, organization_id=None)
|
||||
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
|
||||
|
||||
|
|
@ -144,7 +150,7 @@ async def test_join_reads_null_nested_lists_the_way_prisma_does(prisma):
|
|||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}
|
||||
)
|
||||
|
||||
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=None)
|
||||
cache = _frozen_cache()
|
||||
refs = AuthObjectRefs(user_id=user_id, team_id=team_id, membership_user_id=user_id, organization_id=None)
|
||||
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue