From 3d7c1db14ac1752fd00d00462a1289930e5b8f6a Mon Sep 17 00:00:00 2001 From: Praising Harris Ratnam Date: Mon, 19 Oct 2026 01:46:40 +0530 Subject: [PATCH] style: fix ruff formatting and imports in test_auth_checks --- tests/unit/proxy/auth/test_auth_checks.py | 267 +++++++++++----------- 1 file changed, 138 insertions(+), 129 deletions(-) diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index 14ad7596fb4..14f7d1559dd 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -1,10 +1,11 @@ # What is this? ## Tests if 'get_end_user_object' works as expected -import sys, os, asyncio, time, random, uuid +import asyncio import hashlib import json -import traceback +import uuid + from dotenv import load_dotenv load_dotenv() @@ -15,39 +16,40 @@ from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -import pytest, litellm -import httpx +import pytest from fastapi import Request, status from prisma import Prisma + +import litellm +import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module from litellm._service_logger import ServiceTypes -from litellm.proxy._types import LiteLLM_OrganizationTable, UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import get_org_object, get_user_object -from litellm.proxy.db.prisma_client import PrismaWrapper -from litellm.proxy.auth.auth_checks import get_end_user_object from litellm.caching.caching import DualCache -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy._types import ( - LiteLLM_EndUserTable, LiteLLM_BudgetTable, - LiteLLM_UserTable, - LiteLLM_TeamTable, + LiteLLM_EndUserTable, Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, ProxyException, + UserAPIKeyAuth, ) -from litellm.proxy.utils import PrismaClient -import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module -from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.auth.auth_checks import ( + _team_soft_budget_check, can_team_access_model, + common_checks, delete_cache_key_objects, + get_end_user_object, + get_org_object, + get_user_object, is_model_cost_zero, virtual_key_soft_budget_check, - _team_soft_budget_check, - common_checks, ) -from litellm.proxy.utils import ProxyLogging -from litellm.proxy.utils import CallInfo +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.utils import CallInfo, ProxyLogging from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER @@ -160,12 +162,15 @@ async def test_team_require_trace_id_accepts_supported_trace_carriers( @pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"]) @pytest.mark.asyncio async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None: - assert await _run_team_trace_id_check( - route=route, - method="POST", - request_body={}, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route=route, + method="POST", + request_body={}, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.asyncio @@ -197,11 +202,14 @@ async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat() async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route( request_body: dict[str, object], ) -> None: - assert await _run_team_trace_id_check( - route="/v1/responses", - request_body=request_body, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route="/v1/responses", + request_body=request_body, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.parametrize( @@ -240,12 +248,15 @@ async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trac @pytest.mark.asyncio async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None: - assert await _run_team_trace_id_check( - route="/mcp/", - request_body={}, - headers={"x-litellm-trace-id": "trace-mcp-123"}, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route="/mcp/", + request_body={}, + headers={"x-litellm-trace-id": "trace-mcp-123"}, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.parametrize( @@ -278,12 +289,15 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_ @pytest.mark.asyncio async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None: - assert await _run_team_trace_id_check( - route="/anthropic/v1/messages", - request_body={}, - headers={"x-litellm-trace-id": "trace-provider-123"}, - team_metadata={"require_trace_id": True}, - ) is True + assert ( + await _run_team_trace_id_check( + route="/anthropic/v1/messages", + request_body={}, + headers={"x-litellm-trace-id": "trace-provider-123"}, + team_metadata={"require_trace_id": True}, + ) + is True + ) @pytest.mark.parametrize( @@ -329,7 +343,9 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pa "jsonrpc": "2.0", "id": "1", "method": "message/send", - "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}}, + "params": { + "message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]} + }, }, ), ], @@ -370,7 +386,9 @@ async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes "jsonrpc": "2.0", "id": "1", "method": "message/send", - "params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}}, + "params": { + "message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]} + }, "metadata": {"trace_id": "trace-body-123"}, }, ), @@ -463,7 +481,7 @@ def test_lens_management_does_not_require_an_inference_user_parameter(route: str async def test_get_end_user_object(customer_spend, customer_budget): """ Scenario 1: normal - get_end_user_object returns the cached user - Scenario 2: user over budget - NOTE: budget enforcement now happens in + Scenario 2: user over budget - NOTE: budget enforcement now happens in common_checks() via _check_end_user_budget(), not in get_end_user_object() This test verifies that get_end_user_object correctly retrieves the end user @@ -480,7 +498,7 @@ async def test_get_end_user_object(customer_spend, customer_budget): # UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts # and breaks get_end_user_object's typed async_get_cache path. _cache = UserApiKeyCache() - _key = "end_user_id:{}".format(end_user_id) + _key = f"end_user_id:{end_user_id}" await _cache.async_set_cache( key=_key, value=end_user_obj, @@ -550,8 +568,8 @@ async def test_can_key_call_model(model, expect_to_work): """ If wildcard model + specific model is used, choose the specific model settings """ + from litellm.proxy.auth.auth_checks import can_key_call_model - from fastapi import HTTPException llm_model_list = [ { @@ -591,7 +609,7 @@ async def test_can_key_call_model(model, expect_to_work): if expect_to_work: await can_key_call_model(**args) else: - with pytest.raises(Exception, match='is not available for this API key') as e: + with pytest.raises(Exception, match="is not available for this API key") as e: await can_key_call_model(**args) print(e) @@ -603,8 +621,8 @@ async def test_can_key_call_model(model, expect_to_work): ) @pytest.mark.asyncio async def test_can_team_call_model(model, expect_to_work): + from litellm.proxy.auth.auth_checks import model_in_access_group - from fastapi import HTTPException llm_model_list = [ { @@ -850,8 +868,8 @@ async def test_add_known_models_explicit_map_updates_provider_sets(): @pytest.mark.asyncio async def test_is_valid_fallback_model(): - from litellm.proxy.auth.auth_checks import is_valid_fallback_model from litellm import Router + from litellm.proxy.auth.auth_checks import is_valid_fallback_model router = Router( model_list=[ @@ -863,16 +881,12 @@ async def test_is_valid_fallback_model(): ) try: - await is_valid_fallback_model( - model="gpt-3.5-turbo", llm_router=router, user_model=None - ) + await is_valid_fallback_model(model="gpt-3.5-turbo", llm_router=router, user_model=None) except Exception as e: pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}") with pytest.raises(Exception, match="Invalid") as exc_info: - await is_valid_fallback_model( - model="gpt-4o", llm_router=router, user_model=None - ) + await is_valid_fallback_model(model="gpt-4o", llm_router=router, user_model=None) e = exc_info.value assert "Invalid" in str(e) @@ -886,9 +900,7 @@ async def test_is_valid_fallback_model(): ], ) @pytest.mark.asyncio -async def test_virtual_key_max_budget_check( - token_spend, max_budget, expect_budget_error -): +async def test_virtual_key_max_budget_check(token_spend, max_budget, expect_budget_error): """ Test if virtual key budget checks work as expected: 1. Triggers budget alert for all cases @@ -1003,13 +1015,11 @@ async def test_can_team_access_model(model, team_models, expect_to_work): team_model_aliases=None, ) if not expect_to_work: - pytest.fail( - f"Expected model access check to fail for model={model}, team_models={team_models}" - ) + pytest.fail(f"Expected model access check to fail for model={model}, team_models={team_models}") except Exception as e: if expect_to_work: pytest.fail( - f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {str(e)}" + f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {e!s}" ) @@ -1058,9 +1068,9 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert): await asyncio.sleep(0.1) # Allow time for the alert task to complete - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + ) @pytest.mark.parametrize( @@ -1113,9 +1123,7 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert): ], ) @pytest.mark.asyncio -async def test_team_soft_budget_check( - spend, soft_budget, expect_alert, metadata, expected_alert_emails -): +async def test_team_soft_budget_check(spend, soft_budget, expect_alert, metadata, expected_alert_emails): """ Test cases for _team_soft_budget_check: 1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured) @@ -1163,9 +1171,9 @@ async def test_team_soft_budget_check( await asyncio.sleep(0.1) # Allow time for the alert task to complete - assert ( - alert_triggered == expect_alert - ), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + assert alert_triggered == expect_alert, ( + f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}" + ) if expect_alert: assert captured_call_info is not None @@ -1177,17 +1185,14 @@ async def test_team_soft_budget_check( if expected_alert_emails is not None: assert captured_call_info.alert_emails == expected_alert_emails else: - assert ( - captured_call_info.alert_emails is None - or captured_call_info.alert_emails == [] - ) + assert captured_call_info.alert_emails is None or captured_call_info.alert_emails == [] @pytest.mark.asyncio async def test_can_user_call_model(): - from litellm.proxy.auth.auth_checks import can_user_call_model - from litellm.proxy._types import ProxyException from litellm import Router + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_user_call_model router = Router( model_list=[ @@ -1225,10 +1230,11 @@ async def test_can_user_call_model(): @pytest.mark.asyncio async def test_can_user_call_model_with_no_default_models(): - from litellm.proxy.auth.auth_checks import can_user_call_model - from litellm.proxy._types import ProxyException, SpecialModelNames from unittest.mock import MagicMock + from litellm.proxy._types import ProxyException, SpecialModelNames + from litellm.proxy.auth.auth_checks import can_user_call_model + args = { "model": "anthropic-claude", "llm_router": MagicMock(), @@ -1249,9 +1255,10 @@ async def test_can_user_call_model_with_no_default_models(): @pytest.mark.asyncio async def test_get_fuzzy_user_object(): - from litellm.proxy.auth.auth_checks import _get_fuzzy_user_object from unittest.mock import AsyncMock, MagicMock + from litellm.proxy.auth.auth_checks import _get_fuzzy_user_object + # Setup mock Prisma client mock_prisma = MagicMock() mock_prisma.db = MagicMock() @@ -1311,9 +1318,7 @@ async def test_get_fuzzy_user_object(): # Test 5: Only email provided (no SSO ID) mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user) - result = await _get_fuzzy_user_object( - prisma_client=mock_prisma, user_email="test@example.com" - ) + result = await _get_fuzzy_user_object(prisma_client=mock_prisma, user_email="test@example.com") assert result == test_user mock_prisma.db.litellm_usertable.find_first.assert_called_with( where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}}, @@ -1322,9 +1327,7 @@ async def test_get_fuzzy_user_object(): # Test 6: Only SSO ID provided (no email) mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user) - result = await _get_fuzzy_user_object( - prisma_client=mock_prisma, sso_user_id="sso_123" - ) + result = await _get_fuzzy_user_object(prisma_client=mock_prisma, sso_user_id="sso_123") assert result == test_user mock_prisma.db.litellm_usertable.find_unique.assert_called_with( where={"sso_user_id": "sso_123"}, include={"organization_memberships": True} @@ -1371,7 +1374,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work) llm_router=router, ) else: - with pytest.raises(Exception, match='is not available for this API key') as e: + with pytest.raises(Exception, match="is not available for this API key") as e: await can_key_call_model( model=model, llm_model_list=llm_model_list, @@ -1388,8 +1391,8 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work) @pytest.mark.asyncio async def test_cache_access_object(): """Test _cache_access_object stores access group in cache with correct key.""" - from litellm.proxy.auth.auth_checks import cache_access_object from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import cache_access_object cache = DualCache() ag_id = "ag-test-123" @@ -1416,8 +1419,8 @@ async def test_cache_access_object(): @pytest.mark.asyncio async def test_delete_cache_access_object(): """Test _delete_cache_access_object removes access group from in-memory cache.""" - from litellm.proxy.auth.auth_checks import delete_cache_access_object from litellm.proxy._types import LiteLLM_AccessGroupTable + from litellm.proxy.auth.auth_checks import delete_cache_access_object cache = DualCache() ag_id = "ag-delete-test" @@ -1457,9 +1460,7 @@ async def test_delete_cache_access_object(): ], ) @pytest.mark.asyncio -async def test_get_resources_from_access_groups( - resource_field, access_group_data, expected -): +async def test_get_resources_from_access_groups(resource_field, access_group_data, expected): """Test _get_resources_from_access_groups returns correct resource list from access groups.""" from unittest.mock import AsyncMock, MagicMock, patch @@ -1542,8 +1543,8 @@ async def test_can_team_access_model_access_group_ids_denied(): """Test can_team_access_model denies when neither team models nor access_group_ids grant access.""" from unittest.mock import AsyncMock, patch - from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_team_access_model team_object = LiteLLM_TeamTable( team_id="test-team", @@ -1609,8 +1610,8 @@ async def test_can_key_call_model_empty_models_empty_access_groups_denied(): """Test can_key_call_model fails closed (denies) when key has empty models and access_group_ids resolve to empty.""" from unittest.mock import AsyncMock, patch - from litellm.proxy.auth.auth_checks import can_key_call_model from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_key_call_model user_api_key_object = UserAPIKeyAuth( token="test-token", @@ -1671,20 +1672,26 @@ async def test_can_key_call_model_union_preserves_both_key_and_group_models(): return_value=["claude-3"], ): # Access gpt-3.5-turbo (from key_models) - assert await can_key_call_model( - model="gpt-3.5-turbo", - llm_model_list=[], - valid_token=user_api_key_object, - llm_router=router, - ) is True + assert ( + await can_key_call_model( + model="gpt-3.5-turbo", + llm_model_list=[], + valid_token=user_api_key_object, + llm_router=router, + ) + is True + ) # Access claude-3 (from access_group_ids) - assert await can_key_call_model( - model="claude-3", - llm_model_list=[], - valid_token=user_api_key_object, - llm_router=router, - ) is True + assert ( + await can_key_call_model( + model="claude-3", + llm_model_list=[], + valid_token=user_api_key_object, + llm_router=router, + ) + is True + ) @pytest.mark.asyncio @@ -1692,8 +1699,8 @@ async def test_can_team_access_model_empty_models_empty_access_groups_denied(): """Test can_team_access_model fails closed (denies) when team has empty models and access_group_ids resolve to empty.""" from unittest.mock import AsyncMock, patch - from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_checks import can_team_access_model team_object = LiteLLM_TeamTable( team_id="test-team", @@ -1734,20 +1741,26 @@ async def test_can_team_access_model_union_preserves_both_team_and_group_models( return_value=["claude-3"], ): # Access gpt-4 (from team_models) - assert await can_team_access_model( - model="gpt-4", - team_object=team_object, - llm_router=None, - team_model_aliases=None, - ) is True + assert ( + await can_team_access_model( + model="gpt-4", + team_object=team_object, + llm_router=None, + team_model_aliases=None, + ) + is True + ) # Access claude-3 (from access_group_ids) - assert await can_team_access_model( - model="claude-3", - team_object=team_object, - llm_router=None, - team_model_aliases=None, - ) is True + assert ( + await can_team_access_model( + model="claude-3", + team_object=team_object, + llm_router=None, + team_model_aliases=None, + ) + is True + ) # --------------------------------------------------------------------------- @@ -2195,7 +2208,9 @@ def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_ro model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"}, ) expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False} - order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry") + order: Final = ( + ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry") + ) verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order} @@ -2438,10 +2453,7 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) client = _WedgedPublishRedisClient() cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) - hashed_tokens: Final = tuple( - hashlib.sha256(f"token-{index}".encode()).hexdigest() - for index in range(600) - ) + hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600)) for hashed_token in hashed_tokens: cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token}) seeded: Final = tuple( @@ -2501,10 +2513,7 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) client = _WedgedPublishRedisClient() - hashed_tokens: Final = tuple( - hashlib.sha256(f"token-{index}".encode()).hexdigest() - for index in range(600) - ) + hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600)) redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens}) cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000)) cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer