From 2ae1cf608f3d9cb6837a293aeb6e9941870b5fb6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 15 Sep 2026 19:16:31 +0000 Subject: [PATCH] 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 --- .../proxy_behavior/auth/test_auth_object_prefetch.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/proxy_behavior/auth/test_auth_object_prefetch.py b/tests/proxy_behavior/auth/test_auth_object_prefetch.py index e2d947f4284..cfa958500af 100644 --- a/tests/proxy_behavior/auth/test_auth_object_prefetch.py +++ b/tests/proxy_behavior/auth/test_auth_object_prefetch.py @@ -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)