From bfb0a3854ec1604bf6156b66489adccc32936d52 Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 14:21:22 -0700 Subject: [PATCH] Enhance logging in cost calculation tests to ensure DEBUG level captures are accurate. Updated tests to set logger level before assertions and restored original logger level after execution. This improves reliability of log level checks in both cost and batch cost calculation tests. --- .../test_cost_calculation_log_level.py | 144 ++++++++++-------- 1 file changed, 81 insertions(+), 63 deletions(-) diff --git a/tests/test_litellm/test_cost_calculation_log_level.py b/tests/test_litellm/test_cost_calculation_log_level.py index 4380ae8bf62..3925ea751af 100644 --- a/tests/test_litellm/test_cost_calculation_log_level.py +++ b/tests/test_litellm/test_cost_calculation_log_level.py @@ -17,46 +17,55 @@ def test_cost_calculation_uses_debug_level(caplog): This ensures cost calculation details don't appear in production logs. Part of fix for issue #9815. """ - # Create a mock completion response - mock_response = { - "id": "test", - "object": "chat.completion", - "created": 1234567890, - "model": "gpt-3.5-turbo", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop" - }], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "total_tokens": 30 + # Ensure verbose_logger is set to DEBUG level to capture the debug logs + from litellm._logging import verbose_logger + original_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) + + try: + # Create a mock completion response + mock_response = { + "id": "test", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-3.5-turbo", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Test response"}, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + } } - } - - # Test that cost calculation logs are at DEBUG level - with caplog.at_level(logging.DEBUG): - try: - cost = completion_cost( - completion_response=mock_response, - model="gpt-3.5-turbo" - ) - except Exception: - pass # Cost calculation may fail, but we're checking log levels - - # Find the cost calculation log records - cost_calc_records = [ - record for record in caplog.records - if "selected model name for cost calculation" in record.message - ] - - # Verify that cost calculation logs are at DEBUG level - assert len(cost_calc_records) > 0, "No cost calculation logs found" - - for record in cost_calc_records: - assert record.levelno == logging.DEBUG, \ - f"Cost calculation log should be DEBUG level, but was {record.levelname}" + + # Test that cost calculation logs are at DEBUG level + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + try: + cost = completion_cost( + completion_response=mock_response, + model="gpt-3.5-turbo" + ) + except Exception: + pass # Cost calculation may fail, but we're checking log levels + + # Find the cost calculation log records + cost_calc_records = [ + record for record in caplog.records + if "selected model name for cost calculation" in record.message + ] + + # Verify that cost calculation logs are at DEBUG level + assert len(cost_calc_records) > 0, "No cost calculation logs found" + + for record in cost_calc_records: + assert record.levelno == logging.DEBUG, \ + f"Cost calculation log should be DEBUG level, but was {record.levelname}" + finally: + # Restore original logger level + verbose_logger.setLevel(original_level) def test_batch_cost_calculation_uses_debug_level(caplog): @@ -65,29 +74,38 @@ def test_batch_cost_calculation_uses_debug_level(caplog): """ from litellm.cost_calculator import batch_cost_calculator from litellm.types.utils import Usage + from litellm._logging import verbose_logger - # Create a mock usage object - usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + # Ensure verbose_logger is set to DEBUG level to capture the debug logs + original_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) - # Test that batch cost calculation logs are at DEBUG level - with caplog.at_level(logging.DEBUG): - try: - batch_cost_calculator( - usage=usage, - model="gpt-3.5-turbo", - custom_llm_provider="openai" - ) - except Exception: - pass # May fail, but we're checking log levels - - # Find batch cost calculation log records - batch_cost_records = [ - record for record in caplog.records - if "Calculating batch cost per token" in record.message - ] - - # Verify logs exist and are at DEBUG level - if batch_cost_records: # May not always log depending on the code path - for record in batch_cost_records: - assert record.levelno == logging.DEBUG, \ - f"Batch cost calculation log should be DEBUG level, but was {record.levelname}" \ No newline at end of file + try: + # Create a mock usage object + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + + # Test that batch cost calculation logs are at DEBUG level + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + try: + batch_cost_calculator( + usage=usage, + model="gpt-3.5-turbo", + custom_llm_provider="openai" + ) + except Exception: + pass # May fail, but we're checking log levels + + # Find batch cost calculation log records + batch_cost_records = [ + record for record in caplog.records + if "Calculating batch cost per token" in record.message + ] + + # Verify logs exist and are at DEBUG level + if batch_cost_records: # May not always log depending on the code path + for record in batch_cost_records: + assert record.levelno == logging.DEBUG, \ + f"Batch cost calculation log should be DEBUG level, but was {record.levelname}" + finally: + # Restore original logger level + verbose_logger.setLevel(original_level) \ No newline at end of file