diff --git a/litellm/tests/test_jwt.py b/litellm/tests/test_jwt.py index a2edbcc2f..960b85af5 100644 --- a/litellm/tests/test_jwt.py +++ b/litellm/tests/test_jwt.py @@ -1,8 +1,14 @@ #### What this tests #### # Unit tests for JWT-Auth -import sys, os, asyncio, time, random, uuid +import asyncio +import os +import random +import sys +import time import traceback +import uuid + from dotenv import load_dotenv load_dotenv() @@ -11,13 +17,15 @@ import os sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +from datetime import datetime, timedelta + import pytest +from fastapi import Request + +from litellm.caching import DualCache from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLMRoutes from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.team_endpoints import new_team -from litellm.caching import DualCache -from datetime import datetime, timedelta -from fastapi import Request public_key = { "kty": "RSA", @@ -89,10 +97,12 @@ async def test_valid_invalid_token(audience): - valid token - invalid token """ - import jwt, json + import json + + import jwt + from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa - from cryptography.hazmat.backends import default_backend os.environ.pop("JWT_AUDIENCE", None) if audience: @@ -191,8 +201,8 @@ async def test_valid_invalid_token(audience): @pytest.fixture def prisma_client(): import litellm - from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.proxy.proxy_cli import append_query_params + from litellm.proxy.utils import PrismaClient, ProxyLogging proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) @@ -213,16 +223,19 @@ def prisma_client(): @pytest.mark.parametrize("audience", [None, "litellm-proxy"]) @pytest.mark.asyncio async def test_team_token_output(prisma_client, audience): - import jwt, json + import json + import uuid + + import jwt + from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa - from cryptography.hazmat.backends import default_backend from fastapi import Request from starlette.datastructures import URL - from litellm.proxy.proxy_server import user_api_key_auth - from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth + import litellm - import uuid + from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth + from litellm.proxy.proxy_server import user_api_key_auth setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) await litellm.proxy.proxy_server.prisma_client.connect() @@ -392,20 +405,23 @@ async def test_user_token_output( - create user - retry -> it should pass now """ - import jwt, json + import json + import uuid + + import jwt + from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa - from cryptography.hazmat.backends import default_backend from fastapi import Request from starlette.datastructures import URL - from litellm.proxy.proxy_server import ( - user_api_key_auth, + + import litellm + from litellm.proxy._types import NewTeamRequest, NewUserRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( new_user, user_info, ) - from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, NewUserRequest - import litellm - import uuid + from litellm.proxy.proxy_server import user_api_key_auth setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) await litellm.proxy.proxy_server.prisma_client.connect() @@ -617,16 +633,19 @@ async def test_allowed_routes_admin(prisma_client, audience): - iterate through allowed endpoints - check if admin passes user_api_key_auth for them """ - import jwt, json + import json + import uuid + + import jwt + from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa - from cryptography.hazmat.backends import default_backend from fastapi import Request from starlette.datastructures import URL - from litellm.proxy.proxy_server import user_api_key_auth - from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth + import litellm - import uuid + from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth + from litellm.proxy.proxy_server import user_api_key_auth setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) await litellm.proxy.proxy_server.prisma_client.connect()