fix(test): update prometheus tests to use new increment_deployment_cooled_down and set_litellm_deployment_state API

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

litellm/integrations/prometheus.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
 import os
 import sys
 from datetime import datetime, timedelta
@@ -2821,8 +2822,8 @@ class PrometheusLogger(CustomLogger):
             supported_enum_labels=self.get_labels_for_metric(
                 metric_name="litellm_deployment_cooled_down"
             ),
-            enum_values=UserAPIKeyLabelValues(
-                **{**enum_values.__dict__, "exception_status": exception_status}
+            enum_values=dataclasses.replace(
+                enum_values, exception_status=exception_status
             ),
         )
         self.litellm_deployment_cooled_down.labels(**_labels).inc()

tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py

@@ -1179,39 +1179,28 @@ def test_deployment_state_management(prometheus_logger, monkeypatch):
 
 
 def test_increment_deployment_cooled_down(prometheus_logger, monkeypatch):
-    import inspect
+    from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
 
     monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True)
     prometheus_logger.get_labels_for_metric = (
         lambda metric_name: PrometheusMetricLabels.get_labels(metric_name)
     )
 
-    method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down)
-    expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"])
-
     mock_chain = MagicMock()
-
-    def validating_labels(*label_values, **label_kwargs):
-        """Validate label count matches metric definition"""
-        total = len(label_values) + len(label_kwargs)
-        if total != expected_label_count:
-            raise ValueError(
-                f"Incorrect label count: expected {expected_label_count}, got {total}"
-            )
-        return mock_chain
-
     prometheus_logger.litellm_deployment_cooled_down = MagicMock()
     prometheus_logger.litellm_deployment_cooled_down.labels = MagicMock(
-        side_effect=validating_labels
+        return_value=mock_chain
     )
 
     prometheus_logger.increment_deployment_cooled_down(
-        litellm_model_name="gpt-5-mini",
-        model_id="model-123",
-        api_base="https://api.openai.com",
-        api_provider="openai",
+        enum_values=UserAPIKeyLabelValues(
+            litellm_model_name="gpt-5-mini",
+            model_id="model-123",
+            api_base="https://api.openai.com",
+            api_provider="openai",
+            model_group="openai-gpt",
+        ),
         exception_status="429",
-        model_group="openai-gpt",
     )
 
     prometheus_logger.litellm_deployment_cooled_down.labels.assert_called_once_with(

tests/test_litellm/integrations/test_prometheus_model_group_labels.py

@@ -100,13 +100,17 @@ def test_set_litellm_deployment_state_emits_model_group(monkeypatch):
         lambda metric_name: PrometheusMetricLabels.get_labels(metric_name)
     )
 
+    from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
+
     logger.set_litellm_deployment_state(
         state=2,
-        litellm_model_name="gpt-4o-mini",
-        model_id="model-123",
-        api_base="https://api.openai.com",
-        api_provider="openai",
-        model_group="gpt-group",
+        enum_values=UserAPIKeyLabelValues(
+            litellm_model_name="gpt-4o-mini",
+            model_id="model-123",
+            api_base="https://api.openai.com",
+            api_provider="openai",
+            model_group="gpt-group",
+        ),
     )
 
     labels = logger.litellm_deployment_state.labels.call_args.kwargs
@@ -154,13 +158,17 @@ def test_deployment_metrics_omit_model_group_when_flag_disabled(monkeypatch):
         lambda metric_name: PrometheusMetricLabels.get_labels(metric_name)
     )
 
+    from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
+
     logger.set_litellm_deployment_state(
         state=0,
-        litellm_model_name="gpt-4o-mini",
-        model_id="model-123",
-        api_base="https://api.openai.com",
-        api_provider="openai",
-        model_group="gpt-group",
+        enum_values=UserAPIKeyLabelValues(
+            litellm_model_name="gpt-4o-mini",
+            model_id="model-123",
+            api_base="https://api.openai.com",
+            api_provider="openai",
+            model_group="gpt-group",
+        ),
     )
 
     assert "model_group" not in logger.litellm_deployment_state.labels.call_args.kwargs