From 5bed0b7557a918b7113ff94982dd804dfa219ff7 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 4 Apr 2025 15:55:14 -0700
Subject: [PATCH] aggregated values

---
 .../test_spend_update_queue.py                | 113 ++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py
index 98d3b4e4c73..9d0c9dbb5f0 100644
--- a/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py
+++ b/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py
@@ -6,6 +6,7 @@ import sys
 import pytest
 from fastapi.testclient import TestClient
 
+from litellm.constants import MAX_SIZE_IN_MEMORY_QUEUE
 from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
 from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
 
@@ -150,3 +151,115 @@ async def test_missing_entity_type(spend_queue):
 
     # Should ignore updates without entity type
     assert all(len(transactions) == 0 for transactions in aggregated.values())
+
+
+@pytest.mark.asyncio
+async def test_queue_max_size_triggers_aggregation(monkeypatch, spend_queue):
+    """Test that reaching MAX_SIZE_IN_MEMORY_QUEUE triggers aggregation"""
+    # Override MAX_SIZE_IN_MEMORY_QUEUE for testing
+    monkeypatch.setattr(spend_queue, "MAX_SIZE_IN_MEMORY_QUEUE", 6)
+
+    # Add 6 updates for the same user (exceeding the max size)
+    for i in range(6):
+        update: SpendUpdateQueueItem = {
+            "entity_type": Litellm_EntityType.USER,
+            "entity_id": "user123",
+            "response_cost": 1.0,
+        }
+        await spend_queue.add_update(update)
+
+    # Queue should have been aggregated, resulting in a single entry
+    assert spend_queue.update_queue.qsize() == 1
+
+    # Verify the aggregated cost is correct
+    aggregated = (
+        await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
+    )
+    assert aggregated["user_list_transactions"]["user123"] == 6.0
+
+
+@pytest.mark.asyncio
+async def test_aggregate_queue_updates_accuracy(spend_queue):
+    """Test that queue aggregation correctly combines costs by entity type and ID"""
+    # Add multiple updates for different entities
+    updates = [
+        {
+            "entity_type": Litellm_EntityType.USER,
+            "entity_id": "user1",
+            "response_cost": 1.5,
+        },
+        {
+            "entity_type": Litellm_EntityType.USER,
+            "entity_id": "user1",
+            "response_cost": 2.5,
+        },
+        {
+            "entity_type": Litellm_EntityType.USER,
+            "entity_id": "user2",
+            "response_cost": 3.0,
+        },
+        {
+            "entity_type": Litellm_EntityType.TEAM,
+            "entity_id": "team1",
+            "response_cost": 5.0,
+        },
+    ]
+
+    for update in updates:
+        await spend_queue.update_queue.put(update)
+
+    # Force aggregation
+    await spend_queue.aggregate_queue_updates()
+
+    # Queue size should now be 3 (user1, user2, team1)
+    assert spend_queue.update_queue.qsize() == 3
+
+    # Flush and verify aggregated values
+    aggregated = (
+        await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
+    )
+    print("aggregated values", aggregated)
+
+    assert aggregated["user_list_transactions"]["user1"] == 4.0  # 1.5 + 2.5
+    assert aggregated["user_list_transactions"]["user2"] == 3.0
+    assert aggregated["team_list_transactions"]["team1"] == 5.0
+
+
+@pytest.mark.asyncio
+async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue):
+    """Test that queue size is actually reduced when dealing with many items"""
+    # Set a smaller MAX_SIZE for testing
+    monkeypatch.setattr(spend_queue, "MAX_SIZE_IN_MEMORY_QUEUE", 10)
+
+    # Add 30 updates (20 for user1, 10 for user2)
+    for i in range(20):
+        await spend_queue.add_update(
+            {
+                "entity_type": Litellm_EntityType.USER,
+                "entity_id": "user1",
+                "response_cost": 0.5,
+            }
+        )
+
+    # At this point, aggregation should have happened at least once
+    # Queue size should be much less than 20
+    assert spend_queue.update_queue.qsize() <= 2
+
+    for i in range(10):
+        await spend_queue.add_update(
+            {
+                "entity_type": Litellm_EntityType.USER,
+                "entity_id": "user2",
+                "response_cost": 1.0,
+            }
+        )
+
+    # Queue should have at most 2 items after all this activity
+    assert spend_queue.update_queue.qsize() <= 2
+
+    # Verify total costs are correct
+    aggregated = (
+        await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
+    )
+    assert aggregated["user_list_transactions"]["user1"] == 10.0  # 20 * 0.5
+    assert aggregated["user_list_transactions"]["user2"] == 10.0  # 10 * 1.0