From 53b2724e99c85e3daa26519d4e2c4dd77526b839 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 24 Jun 2026 17:28:11 +0530 Subject: [PATCH] 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 --- litellm/integrations/prometheus.py | 5 ++-- .../test_prometheus_logging_callbacks.py | 29 ++++++------------- .../test_prometheus_model_group_labels.py | 28 +++++++++++------- 3 files changed, 30 insertions(+), 32 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 83a40684f53..19c64aaa535 100644 --- a/litellm/integrations/prometheus.py +++ b/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() diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 3c4b5c0b21c..ef4ac8bb785 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/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( diff --git a/tests/test_litellm/integrations/test_prometheus_model_group_labels.py b/tests/test_litellm/integrations/test_prometheus_model_group_labels.py index f039fe5b81f..b914bbe6520 100644 --- a/tests/test_litellm/integrations/test_prometheus_model_group_labels.py +++ b/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