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.

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}"
+    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)