Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions tests/proxy_behavior/auth/test_auth_object_prefetch.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,11 @@
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)
Expand Down Expand Up @@ -58,7 +63,7 @@ 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)

Expand Down Expand Up @@ -100,7 +105,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)

Expand Down Expand Up @@ -144,7 +149,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)

Expand Down
Loading