From 680b5731347b8cf0dfcd8f989cd8604f71a6b23a Mon Sep 17 00:00:00 2001 From: shin-bot-litellm Date: Sat, 31 Jan 2026 15:56:24 +0000 Subject: [PATCH] fix(lint): remove unused imports, fix None comparison, replace print with logging --- .../test_prometheus_logging_callbacks.py | 31 +++++++------------ 1 file changed, 11 insertions(+), 20 deletions(-) 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 ab57e520129..2c7d89f0eaa 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 @@ -1,4 +1,3 @@ -import io import os import sys @@ -10,13 +9,10 @@ from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, call, patch import pytest -from prometheus_client import REGISTRY, CollectorRegistry +from prometheus_client import REGISTRY import litellm -from litellm import completion from litellm._logging import verbose_logger -from litellm._uuid import uuid -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -37,7 +33,6 @@ from litellm.proxy._types import UserAPIKeyAuth verbose_logger.setLevel(logging.DEBUG) litellm.set_verbose = True -import time @pytest.fixture @@ -1147,7 +1142,7 @@ def test_prometheus_factory(monkeypatch, enable_end_user_cost_tracking_prometheu if enable_end_user_cost_tracking_prometheus_only is True: assert returned_dict["end_user"] == "test_end_user" else: - assert returned_dict["end_user"] == None + assert returned_dict["end_user"] is None def test_get_custom_labels_from_metadata(monkeypatch): @@ -1492,7 +1487,7 @@ async def test_initialize_remaining_budget_metrics(prometheus_logger): """ litellm.prometheus_initialize_budget_metrics = True # Mock the prisma client and get_paginated_teams function - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + with patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.management_endpoints.team_endpoints.get_paginated_teams" ) as mock_get_teams: @@ -1659,7 +1654,7 @@ async def test_initialize_api_key_budget_metrics(prometheus_logger): """ litellm.prometheus_initialize_budget_metrics = True # Mock the prisma client and _list_key_helper function - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + with patch("litellm.proxy.proxy_server.prisma_client"), patch( "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper" ) as mock_list_keys: @@ -1917,7 +1912,6 @@ def test_prometheus_label_factory_with_custom_tags(monkeypatch): Test that prometheus_label_factory correctly handles custom tags """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -1955,7 +1949,6 @@ def test_prometheus_label_factory_with_no_custom_tags(monkeypatch): Test that prometheus_label_factory works when no custom tags are configured """ from litellm.integrations.prometheus import ( - get_custom_labels_from_tags, prometheus_label_factory, ) from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @@ -2180,9 +2173,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): All three metrics should be properly incremented when making a successful completion request. """ - from prometheus_client import CollectorRegistry, Counter - import litellm from litellm.types.integrations.prometheus import PrometheusMetricsConfig # Clear registry before test @@ -2249,7 +2240,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): await asyncio.sleep(2) - print("final registry values", REGISTRY._collector_to_names) + verbose_logger.debug("final registry values %s", REGISTRY._collector_to_names) # Get metric collectors directly from registry metric_collectors = {} @@ -2257,7 +2248,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): metric_name = names[0] # First name is the base metric name metric_collectors[metric_name] = collector - print("=== Final Metric Values (Direct Access) ===") + verbose_logger.debug("=== Final Metric Values (Direct Access) ===") # Expected values expected_values = { @@ -2294,10 +2285,10 @@ async def test_prometheus_token_metrics_with_prometheus_config(): actual_value = total_sample.value actual_labels = total_sample.labels - print( - f"✓ {metric_name}: expected={expected_value}, actual={actual_value}" + verbose_logger.debug( + "✓ %s: expected=%s, actual=%s", metric_name, expected_value, actual_value ) - print(f" Labels: {actual_labels}") + verbose_logger.debug(" Labels: %s", actual_labels) # Validate the value assert ( @@ -2314,12 +2305,12 @@ async def test_prometheus_token_metrics_with_prometheus_config(): actual_label_value == expected_label_value ), f"Expected label {label_key}={expected_label_value}, got {actual_label_value}" - print(f" ✓ {metric_name} VALIDATED") + verbose_logger.debug(" ✓ %s VALIDATED", metric_name) else: raise AssertionError(f"No _total sample found for {metric_name}") else: raise AssertionError(f"Metric {metric_name} not found in registry") - print("✓ All token metrics validated successfully!") + verbose_logger.debug("✓ All token metrics validated successfully!") # check final value of metrics in registry