From 86e0dd68c3c6f74f7955ba09dfde99691c3daab6 Mon Sep 17 00:00:00 2001 From: Sara Ghaemi Date: Tue, 7 May 2024 13:28:57 -0400 Subject: [PATCH] updated tests --- litellm/tests/test_jwt.py | 41 +++++++++++++++++++++++---------------- 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/litellm/tests/test_jwt.py b/litellm/tests/test_jwt.py index 3dacebe2702..b3af9913f66 100644 --- a/litellm/tests/test_jwt.py +++ b/litellm/tests/test_jwt.py @@ -77,9 +77,9 @@ async def test_token_single_public_key(): == "qIgOQfEVrrErJC0E7gsHXi6rs_V0nyFY5qPFui2-tv0o4CwpwDzgfBtLO7o_wLiguq0lnu54sMT2eLNoRiiPuLvv6bg7Iy1H9yc5_4Jf5oYEOrqN5o9ZBOoYp1q68Pv0oNJYyZdGu5ZJfd7V4y953vB2XfEKgXCsAkhVhlvIUMiDNKWoMDWsyb2xela5tRURZ2mJAXcHfSC_sYdZxIA2YYrIHfoevq_vTlaz0qVSe_uOKjEpgOAS08UUrgda4CQL11nzICiIQzc6qmjIQt2cjzB2D_9zb4BYndzEtfl0kwAT0z_I85S3mkwTqHU-1BvKe_4MG4VG3dAAeffLPXJyXQ" ) - +@pytest.mark.parametrize('audience', [None, "litellm-proxy"]) @pytest.mark.asyncio -async def test_valid_invalid_token(): +async def test_valid_invalid_token(audience): """ Tests - valid token @@ -90,6 +90,10 @@ async def test_valid_invalid_token(): from cryptography.hazmat.primitives.asymmetric import rsa from cryptography.hazmat.backends import default_backend + os.environ.pop('JWT_AUDIENCE', None) + if audience: + os.environ["JWT_AUDIENCE"] = audience + # Generate a private / public key pair using RSA algorithm key = rsa.generate_private_key( public_exponent=65537, key_size=2048, backend=default_backend() @@ -134,9 +138,8 @@ async def test_valid_invalid_token(): "sub": "user123", "exp": expiration_time, # set the token to expire in 10 minutes "scope": "litellm-proxy-admin", + "aud": audience } - if os.getenv("JWT_AUDIENCE"): - payload["aud"] = os.getenv("JWT_AUDIENCE") # Generate the JWT token # But before, you should convert bytes to string @@ -163,9 +166,8 @@ async def test_valid_invalid_token(): "sub": "user123", "exp": expiration_time, # set the token to expire in 10 minutes "scope": "litellm-NO-SCOPE", + "aud": audience } - if os.getenv("JWT_AUDIENCE"): - payload["aud"] = os.getenv("JWT_AUDIENCE") # Generate the JWT token # But before, you should convert bytes to string @@ -181,7 +183,6 @@ async def test_valid_invalid_token(): except Exception as e: pytest.fail(f"An exception occurred - {str(e)}") - @pytest.fixture def prisma_client(): import litellm @@ -204,8 +205,9 @@ def prisma_client(): return prisma_client +@pytest.mark.parametrize('audience', [None, "litellm-proxy"]) @pytest.mark.asyncio -async def test_team_token_output(prisma_client): +async def test_team_token_output(prisma_client, audience): import jwt, json from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa @@ -220,6 +222,10 @@ async def test_team_token_output(prisma_client): setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) await litellm.proxy.proxy_server.prisma_client.connect() + os.environ.pop('JWT_AUDIENCE', None) + if audience: + os.environ["JWT_AUDIENCE"] = audience + # Generate a private / public key pair using RSA algorithm key = rsa.generate_private_key( public_exponent=65537, key_size=2048, backend=default_backend() @@ -268,9 +274,8 @@ async def test_team_token_output(prisma_client): "exp": expiration_time, # set the token to expire in 10 minutes "scope": "litellm_team", "client_id": team_id, + "aud": audience } - if os.getenv("JWT_AUDIENCE"): - payload["aud"] = os.getenv("JWT_AUDIENCE") # Generate the JWT token # But before, you should convert bytes to string @@ -284,9 +289,8 @@ async def test_team_token_output(prisma_client): "sub": "user123", "exp": expiration_time, # set the token to expire in 10 minutes "scope": "litellm_proxy_admin", + "aud": audience } - if os.getenv("JWT_AUDIENCE"): - payload["aud"] = os.getenv("JWT_AUDIENCE") admin_token = jwt.encode(payload, private_key_str, algorithm="RS256") @@ -354,8 +358,9 @@ async def test_team_token_output(prisma_client): assert team_result.team_models == ["gpt-3.5-turbo", "gpt-4"] +@pytest.mark.parametrize('audience', [None, "litellm-proxy"]) @pytest.mark.asyncio -async def test_user_token_output(prisma_client): +async def test_user_token_output(prisma_client, audience): """ - If user required, check if it exists - fail initial request (when user doesn't exist) @@ -376,6 +381,10 @@ async def test_user_token_output(prisma_client): setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) await litellm.proxy.proxy_server.prisma_client.connect() + os.environ.pop('JWT_AUDIENCE', None) + if audience: + os.environ["JWT_AUDIENCE"] = audience + # Generate a private / public key pair using RSA algorithm key = rsa.generate_private_key( public_exponent=65537, key_size=2048, backend=default_backend() @@ -427,9 +436,8 @@ async def test_user_token_output(prisma_client): "exp": expiration_time, # set the token to expire in 10 minutes "scope": "litellm_team", "client_id": team_id, + "aud": audience } - if os.getenv("JWT_AUDIENCE"): - payload["aud"] = os.getenv("JWT_AUDIENCE") # Generate the JWT token # But before, you should convert bytes to string @@ -443,9 +451,8 @@ async def test_user_token_output(prisma_client): "sub": user_id, "exp": expiration_time, # set the token to expire in 10 minutes "scope": "litellm_proxy_admin", + "aud": audience } - if os.getenv("JWT_AUDIENCE"): - payload["aud"] = os.getenv("JWT_AUDIENCE") admin_token = jwt.encode(payload, private_key_str, algorithm="RS256")