From 2a60af7164b5770a2db1f249db44cf11ec7cf0e5 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 26 Jan 2024 14:43:16 -0800 Subject: [PATCH] (test) key crossing budget --- tests/test_keys.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/test_keys.py b/tests/test_keys.py index a296ef13e..283776b96 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -115,7 +115,9 @@ async def chat_completion(session, key, model="gpt-4"): print() if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") + raise Exception( + f"Request did not return a 200 status code: {status}. Response: {response_text}" + ) return await response.json() @@ -386,3 +388,31 @@ async def test_key_with_budgets(): key_info = await get_key_info(session=session, get_key=key, call_key=key) reset_at_new_value = key_info["info"]["budget_reset_at"] assert reset_at_init_value != reset_at_new_value + + +@pytest.mark.asyncio +async def test_key_crossing_budget(): + """ + - Create key with budget with budget=0.00000001 + - make a /chat/completions call + - wait 5s + - make a /chat/completions call - should fail with key crossed it's budget + + - Check if value updated + """ + from litellm.proxy.utils import hash_token + + async with aiohttp.ClientSession() as session: + key_gen = await generate_key(session=session, i=0, budget=0.0000001) + key = key_gen["key"] + hashed_token = hash_token(token=key) + print(f"hashed_token: {hashed_token}") + + response = await chat_completion(session=session, key=key) + print("response 1: ", response) + await asyncio.sleep(2) + try: + response = await chat_completion(session=session, key=key) + pytest.fail("Should have failed - Key crossed it's budget") + except Exception as e: + assert "ExceededTokenBudget: Current spend for token:" in str(e)