Fix test_async_post_call_success_hook_includes_client_ip_user_agent

Author: Sameerlite

tests/test_litellm/integrations/test_prometheus_client_ip_user_agent.py

@@ -1,10 +1,12 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
 import pytest
-from unittest.mock import MagicMock, patch
+
 from litellm.integrations.prometheus import PrometheusLogger
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.integrations.prometheus import (
     UserAPIKeyLabelValues,
 )
-from litellm.proxy._types import UserAPIKeyAuth
 
 
 @pytest.mark.asyncio
@@ -72,10 +74,12 @@ async def test_async_post_call_failure_hook_includes_client_ip_user_agent():
 @pytest.mark.asyncio
 async def test_async_post_call_success_hook_includes_client_ip_user_agent():
     """
-    Test that async_post_call_success_hook includes client_ip and user_agent in UserAPIKeyLabelValues
+    Test that async_log_success_event includes client_ip and user_agent in UserAPIKeyLabelValues.
+    
+    Note: After PR #21159, the metric increment was moved from async_post_call_success_hook 
+    to async_log_success_event to prevent double-counting.
     """
     # Mocking
-    # Mocking
     with patch(
         "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None
     ):
@@ -84,27 +88,55 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent():
         logger.get_labels_for_metric = MagicMock(
             return_value=["client_ip", "user_agent"]
         )
-
-    data = {
+        logger._should_skip_metrics_for_invalid_key = MagicMock(return_value=False)
+        logger._increment_top_level_request_and_spend_metrics = MagicMock()
+        logger._increment_token_metrics = MagicMock()
+        logger._increment_remaining_budget_metrics = AsyncMock()
+        logger._set_virtual_key_rate_limit_metrics = MagicMock()
+        logger._set_latency_metrics = MagicMock()
+        logger.set_llm_deployment_success_metrics = MagicMock()
+        logger._increment_cache_metrics = MagicMock()
+
+    kwargs = {
         "model": "gpt-4",
-        "metadata": {
-            "requester_ip_address": "192.168.1.1",
-            "user_agent": "success-agent",
+        "litellm_params": {
+            "metadata": {}
+        },
+        "start_time": None,
+        "standard_logging_object": {
+            "model_group": "gpt-4",
+            "model_id": "model_1",
+            "api_base": "http://api.base",
+            "custom_llm_provider": "openai",
+            "completion_tokens": 10,
+            "total_tokens": 20,
+            "response_cost": 0.01,
+            "request_tags": [],
+            "metadata": {
+                "user_api_key_user_id": "user_1",
+                "user_api_key_hash": "hash_1",
+                "user_api_key_alias": "alias_1",
+                "user_api_key_team_id": "team_1",
+                "user_api_key_team_alias": "team_alias_1",
+                "user_api_key_user_email": "test@example.com",
+                "user_api_key_request_route": "/chat/completions",
+                "requester_ip_address": "192.168.1.1",
+                "user_agent": "success-agent",
+            },
         },
     }
-    user_api_key_dict = UserAPIKeyAuth(token="test_token")
-    response = MagicMock()
 
     # Mock prometheus_label_factory to inspect arguments
     with patch(
         "litellm.integrations.prometheus.prometheus_label_factory"
     ) as mock_label_factory:
         mock_label_factory.return_value = {}
 
-        await logger.async_post_call_success_hook(
-            data=data,
-            user_api_key_dict=user_api_key_dict,
-            response=response,
+        await logger.async_log_success_event(
+            kwargs=kwargs,
+            response_obj=None,
+            start_time=None,
+            end_time=None,
         )
 
         # Verification
@@ -114,8 +146,8 @@ async def test_async_post_call_success_hook_includes_client_ip_user_agent():
         calls = mock_label_factory.call_args_list
         found = False
         for call in calls:
-            kwargs = call.kwargs
-            enum_values = kwargs.get("enum_values")
+            kwargs_args = call.kwargs
+            enum_values = kwargs_args.get("enum_values")
             if isinstance(enum_values, UserAPIKeyLabelValues):
                 if (
                     enum_values.client_ip == "192.168.1.1"