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