diff --git a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py b/tests/test_litellm/responses/test_no_duplicate_spend_logs.py index 19b4d5f524..7c0d6c15d6 100644 --- a/tests/test_litellm/responses/test_no_duplicate_spend_logs.py +++ b/tests/test_litellm/responses/test_no_duplicate_spend_logs.py @@ -49,32 +49,41 @@ def test_logging_object_not_popped(): @pytest.mark.asyncio -async def test_no_duplicate_spend_logs(): +async def test_async_no_duplicate_spend_logs(): """ Test that spend logs are only created once, not duplicated. This integration test verifies the fix by using a custom logger - that counts log_success_event calls. Before the fix, it would be - called twice for non-OpenAI providers (Anthropic/Gemini). + that counts log_success_event calls for a specific request ID. + Before the fix, it would be called twice for non-OpenAI providers. """ - # Create a custom logger to count log_success_event calls + import uuid + + # Generate a unique ID to track only this test's request + test_request_id = f"test-no-dup-{uuid.uuid4()}" + + # Create a custom logger to count log_success_event calls for our specific request class SpendLogCounter(CustomLogger): - def __init__(self): + def __init__(self, tracking_id: str): super().__init__() + self.tracking_id = tracking_id self.log_count = 0 async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.log_count += 1 + # Only count logs for our specific test request + litellm_call_id = kwargs.get("litellm_call_id", "") + if litellm_call_id == self.tracking_id: + self.log_count += 1 - spend_logger = SpendLogCounter() + spend_logger = SpendLogCounter(tracking_id=test_request_id) - # Save original callbacks and set our custom logger - original_callbacks = litellm.callbacks - litellm.callbacks = [spend_logger] + # Save original callbacks and append our logger (don't replace to avoid affecting other tests) + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + litellm.callbacks = original_callbacks + [spend_logger] try: # Call responses API with Anthropic model using mock_response - # This prevents real API calls while still exercising the logging path + # Pass our unique ID as litellm_call_id to track this specific request response = await litellm.aresponses( model="anthropic/claude-3-7-sonnet-latest", input=[{ @@ -83,20 +92,19 @@ async def test_no_duplicate_spend_logs(): "type": "message" }], instructions="You are a helpful assistant.", - mock_response="Hello! I'm doing well." # Use mock to avoid real API call + mock_response="Hello! I'm doing well.", + litellm_call_id=test_request_id, ) - # Wait for async logging to complete using the logging worker's flush method - # Then add a small delay to ensure callbacks finish executing + # Wait for async logging to complete from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER await GLOBAL_LOGGING_WORKER.flush() - # flush() empties the queue but callbacks may still be running await asyncio.sleep(0.5) - # Verify that log_success_event was called exactly once + # Verify that log_success_event was called exactly once for our request assert spend_logger.log_count == 1, ( - f"FAIL: log_success_event called {spend_logger.log_count} times instead of 1. " - f"This indicates duplicate spend logs are being created." + f"FAIL: log_success_event called {spend_logger.log_count} times instead of 1 " + f"for request {test_request_id}. This indicates duplicate spend logs." ) finally: