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 9d0c9dbb5f..9993b25dfd 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 @@ -231,8 +231,8 @@ async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue): # 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): + # Add 30 updates (200 for user1, 10 for key1) + for i in range(200): await spend_queue.add_update( { "entity_type": Litellm_EntityType.USER, @@ -243,23 +243,23 @@ async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue): # 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 + assert spend_queue.update_queue.qsize() <= 10 - for i in range(10): + for i in range(300): await spend_queue.add_update( { - "entity_type": Litellm_EntityType.USER, - "entity_id": "user2", + "entity_type": Litellm_EntityType.KEY, + "entity_id": "key1", "response_cost": 1.0, } ) # Queue should have at most 2 items after all this activity - assert spend_queue.update_queue.qsize() <= 2 + assert spend_queue.update_queue.qsize() <= 10 # 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 + assert aggregated["user_list_transactions"]["user1"] == 200 * 0.5 + assert aggregated["key_list_transactions"]["key1"] == 300 * 1.0