From 952c2b5215cbcd10ae54f189e43e854ac4decfdb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 1 Aug 2025 09:07:53 -0700 Subject: [PATCH] test: update test --- .../test_key_generate_prisma.py | 79 ++++++++++++++----- 1 file changed, 60 insertions(+), 19 deletions(-) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 35c969886b..983f68b89e 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -94,7 +94,9 @@ verbose_proxy_logger.setLevel(level=logging.DEBUG) from starlette.datastructures import URL from litellm.caching.caching import DualCache -from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams +from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, +) from litellm.proxy._types import ( DynamoDBArgs, GenerateKeyRequest, @@ -259,7 +261,9 @@ def test_generate_and_call_with_valid_key(prisma_client, api_route): # check /user/info to verify user_role was set correctly request_mock = MagicMock() new_user_info = await user_info( - request=request_mock, user_id=user_id, user_api_key_dict=user_api_key_dict + request=request_mock, + user_id=user_id, + user_api_key_dict=user_api_key_dict, ) new_user_info = new_user_info.user_info print("new_user_info=", new_user_info) @@ -1358,7 +1362,9 @@ def test_generate_and_update_key(prisma_client): # budget_reset_at should exist for "1mo" duration assert result["info"]["budget_reset_at"] is not None - budget_reset_at = result["info"]["budget_reset_at"].replace(tzinfo=timezone.utc) + budget_reset_at = result["info"]["budget_reset_at"].replace( + tzinfo=timezone.utc + ) current_time = datetime.now(timezone.utc) print(f"Budget reset time: {budget_reset_at}") @@ -1372,14 +1378,20 @@ def test_generate_and_update_key(prisma_client): month_diff = budget_reset_at.month - current_time.month if budget_reset_at.year > current_time.year: month_diff += 12 - + # Should be scheduled for next month (at least 0.5 month away) - assert month_diff >= 1, f"Expected reset to be at least 1 month ahead, got {month_diff} months" - assert month_diff <= 2, f"Expected reset to be at most 2 months ahead, got {month_diff} months" + assert ( + month_diff >= 1 + ), f"Expected reset to be at least 1 month ahead, got {month_diff} months" + assert ( + month_diff <= 2 + ), f"Expected reset to be at most 2 months ahead, got {month_diff} months" else: # Just ensure the date is reasonable (not more than 40 days away) days_diff = (budget_reset_at - current_time).days - assert 0 <= days_diff <= 40, f"Expected reset date to be reasonable, got {days_diff} days from now" + assert ( + 0 <= days_diff <= 40 + ), f"Expected reset date to be reasonable, got {days_diff} days from now" # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) @@ -2298,16 +2310,20 @@ def test_get_bearer_token(): result = _get_bearer_token(api_key) assert result == "sk-1234", f"Expected 'valid_token', got '{result}'" + @pytest.mark.asyncio async def test_update_logs_with_spend_logs_url(prisma_client): """ Unit test for making sure spend logs list is still updated when url passed in """ from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + db_spend_update_writer = DBSpendUpdateWriter() payload = {"startTime": datetime.now(), "endTime": datetime.now()} - await db_spend_update_writer._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + await db_spend_update_writer._insert_spend_log_to_db( + payload=payload, prisma_client=prisma_client + ) assert len(prisma_client.spend_log_transactions) > 0 @@ -2349,7 +2365,7 @@ async def test_user_api_key_auth(prisma_client): print(exc.message) assert ( exc.message - == "Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix. Passed in: my_token" + == "Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix." ) # Test case: User passes empty string API Key @@ -2358,7 +2374,10 @@ async def test_user_api_key_auth(prisma_client): pytest.fail(f"This should have failed!. IT's an invalid key") except ProxyException as exc: print(exc.message) - assert "Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix. Passed in:" in exc.message + assert ( + "Authentication Error, Malformed API Key passed in. Ensure Key has `Bearer ` prefix." + in exc.message + ) @pytest.mark.asyncio @@ -2881,12 +2900,12 @@ async def test_update_user_unit_test(prisma_client): # budget_reset_at should be at midnight 10 days from now budget_reset_at = _user_info["budget_reset_at"].replace(tzinfo=timezone.utc) current_time = datetime.now(timezone.utc) - + # Verify that budget_reset_at is at midnight (hour, minute, second are all 0) assert budget_reset_at.hour == 0 assert budget_reset_at.minute == 0 assert budget_reset_at.second == 0 - + # Calculate days difference - should be close to 10 days (within 1 day to account for time of test execution) days_diff = (budget_reset_at.date() - current_time.date()).days assert 9 <= days_diff <= 10 @@ -3605,7 +3624,10 @@ async def test_key_generate_with_secret_manager_call(prisma_client): assert it is deleted from the secret manager """ from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 - from litellm.types.secret_managers.main import KeyManagementSystem, KeyManagementSettings + from litellm.types.secret_managers.main import ( + KeyManagementSystem, + KeyManagementSettings, + ) from litellm.proxy.hooks.key_management_event_hooks import ( LITELLM_PREFIX_STORED_VIRTUAL_KEYS, @@ -4057,11 +4079,19 @@ async def test_reset_budget_job(prisma_client, entity_type): assert entity_after is not None assert entity_after.spend == 0.0 + def test_delete_nonexistent_key_returns_404(prisma_client): # Try to delete a key that does not exist, expect a 404 error import random, string - from litellm.proxy._types import KeyRequest, UserAPIKeyAuth, LitellmUserRoles, ProxyException - from litellm.proxy.management_endpoints.key_management_endpoints import delete_key_fn + from litellm.proxy._types import ( + KeyRequest, + UserAPIKeyAuth, + LitellmUserRoles, + ProxyException, + ) + from litellm.proxy.management_endpoints.key_management_endpoints import ( + delete_key_fn, + ) from starlette.datastructures import URL from fastapi import Request @@ -4069,25 +4099,36 @@ def test_delete_nonexistent_key_returns_404(prisma_client): setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") try: + async def test(): await litellm.proxy.proxy_server.prisma_client.connect() # Generate a random key that does not exist - random_key = "sk-" + ''.join(random.choices(string.ascii_letters + string.digits, k=24)) + random_key = "sk-" + "".join( + random.choices(string.ascii_letters + string.digits, k=24) + ) delete_key_request = KeyRequest(keys=[random_key]) bearer_token = "Bearer sk-1234" request = Request(scope={"type": "http"}) request._url = URL(url="/key/delete") # use admin to auth in - result = await litellm.proxy.proxy_server.user_api_key_auth(request=request, api_key=bearer_token) + result = await litellm.proxy.proxy_server.user_api_key_auth( + request=request, api_key=bearer_token + ) result.user_role = LitellmUserRoles.PROXY_ADMIN try: await delete_key_fn(data=delete_key_request, user_api_key_dict=result) - pytest.fail("Expected ProxyException 404 for non-existent key, but delete_key_fn did not raise.") + pytest.fail( + "Expected ProxyException 404 for non-existent key, but delete_key_fn did not raise." + ) except ProxyException as e: print("Caught ProxyException:", e) assert str(e.code) == "404" - assert "No keys found" in str(e.message) or "No matching keys or aliases found to delete" in str(e.message) + assert "No keys found" in str( + e.message + ) or "No matching keys or aliases found to delete" in str(e.message) + import asyncio + asyncio.run(test()) except Exception as e: pytest.fail(f"An exception occurred - {str(e)}")