From feba274a89eb5b5924352a9a994268cc1313eddc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 18:39:23 -0700 Subject: [PATCH] test DailySpendUpdateQueue --- .../test_daily_spend_update_queue.py | 264 ++++++++++++++++++ 1 file changed, 264 insertions(+) create mode 100644 tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py new file mode 100644 index 00000000000..228dfe64b1e --- /dev/null +++ b/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py @@ -0,0 +1,264 @@ +import asyncio +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +from litellm.proxy._types import ( + DailyUserSpendTransaction, + Litellm_EntityType, + SpendUpdateQueueItem, +) +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( + DailySpendUpdateQueue, +) +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +@pytest.fixture +def daily_spend_update_queue(): + return DailySpendUpdateQueue() + + +@pytest.mark.asyncio +async def test_empty_queue_flush(daily_spend_update_queue): + """Test flushing an empty queue returns an empty list""" + result = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue() + assert result == [] + + +@pytest.mark.asyncio +async def test_add_single_update(daily_spend_update_queue): + """Test adding a single update to the queue""" + test_key = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + # Add update to queue + await daily_spend_update_queue.add_update({test_key: test_transaction}) + + # Flush and check + updates = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue() + assert len(updates) == 1 + assert test_key in updates[0] + assert updates[0][test_key] == test_transaction + + +@pytest.mark.asyncio +async def test_add_multiple_updates(daily_spend_update_queue): + """Test adding multiple updates to the queue""" + test_key1 = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction1 = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + test_key2 = "user2_2023-01-01_key456_gpt-3.5-turbo_openai" + test_transaction2 = { + "spend": 5.0, + "prompt_tokens": 200, + "completion_tokens": 30, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + # Add updates to queue + await daily_spend_update_queue.add_update({test_key1: test_transaction1}) + await daily_spend_update_queue.add_update({test_key2: test_transaction2}) + + # Flush and check + updates = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue() + assert len(updates) == 2 + + # Find each transaction in the list of updates + found_transaction1 = False + found_transaction2 = False + + for update in updates: + if test_key1 in update: + assert update[test_key1] == test_transaction1 + found_transaction1 = True + if test_key2 in update: + assert update[test_key2] == test_transaction2 + found_transaction2 = True + + assert found_transaction1 + assert found_transaction2 + + +@pytest.mark.asyncio +async def test_aggregated_daily_spend_update_empty(daily_spend_update_queue): + """Test aggregating updates from an empty queue""" + result = ( + await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert result == {} + + +@pytest.mark.asyncio +async def test_get_aggregated_daily_spend_update_transactions_single_key(): + """Test static method for aggregating a single key""" + test_key = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + updates = [{test_key: test_transaction}] + + # Test aggregation + result = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + updates + ) + + assert len(result) == 1 + assert test_key in result + assert result[test_key] == test_transaction + + +@pytest.mark.asyncio +async def test_get_aggregated_daily_spend_update_transactions_multiple_keys(): + """Test static method for aggregating multiple different keys""" + test_key1 = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction1 = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + test_key2 = "user2_2023-01-01_key456_gpt-3.5-turbo_openai" + test_transaction2 = { + "spend": 5.0, + "prompt_tokens": 200, + "completion_tokens": 30, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + updates = [{test_key1: test_transaction1}, {test_key2: test_transaction2}] + + # Test aggregation + result = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + updates + ) + + assert len(result) == 2 + assert test_key1 in result + assert test_key2 in result + assert result[test_key1] == test_transaction1 + assert result[test_key2] == test_transaction2 + + +@pytest.mark.asyncio +async def test_get_aggregated_daily_spend_update_transactions_same_key(): + """Test static method for aggregating updates with the same key""" + test_key = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction1 = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + test_transaction2 = { + "spend": 5.0, + "prompt_tokens": 200, + "completion_tokens": 30, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + expected_transaction = { + "spend": 15.0, # 10 + 5 + "prompt_tokens": 300, # 100 + 200 + "completion_tokens": 80, # 50 + 30 + "api_requests": 2, # 1 + 1 + "successful_requests": 2, # 1 + 1 + "failed_requests": 0, # 0 + 0 + } + + updates = [{test_key: test_transaction1}, {test_key: test_transaction2}] + + # Test aggregation + result = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + updates + ) + + assert len(result) == 1 + assert test_key in result + assert result[test_key] == expected_transaction + + +@pytest.mark.asyncio +async def test_flush_and_get_aggregated_daily_spend_update_transactions( + daily_spend_update_queue, +): + """Test the full workflow of adding, flushing, and aggregating updates""" + test_key = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction1 = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + test_transaction2 = { + "spend": 5.0, + "prompt_tokens": 200, + "completion_tokens": 30, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + expected_transaction = { + "spend": 15.0, # 10 + 5 + "prompt_tokens": 300, # 100 + 200 + "completion_tokens": 80, # 50 + 30 + "api_requests": 2, # 1 + 1 + "successful_requests": 2, # 1 + 1 + "failed_requests": 0, # 0 + 0 + } + + # Add updates to queue + await daily_spend_update_queue.add_update({test_key: test_transaction1}) + await daily_spend_update_queue.add_update({test_key: test_transaction2}) + + # Test full workflow + result = ( + await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + + assert len(result) == 1 + assert test_key in result + assert result[test_key] == expected_transaction