From 5fc0803b945b48484926d8d18f67d7f370ce6f97 Mon Sep 17 00:00:00 2001
From: Jugal Bhatt <jugalbhatt3@gmail.com>
Date: Thu, 14 Aug 2025 14:40:49 -0700
Subject: [PATCH] Add mock user API key authentication in tag management tests

This update introduces a helper function to create a mock user API key authentication object, which is utilized in the tag management endpoint tests. The mock authentication is integrated into the test cases for creating, updating, and deleting tags, enhancing the reliability of the tests by simulating user roles accurately.
---
 .../test_tag_management_endpoints.py          | 33 +++++++++++++++++--
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
index add08f55683..bd68618e56c 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
@@ -14,16 +14,28 @@ from unittest.mock import patch
 
 import litellm
 from litellm.proxy.proxy_server import app
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest
 
 client = TestClient(app)
 
 
+def create_mock_user_api_key_auth():
+    """Helper function to create a mock auth object"""
+    return UserAPIKeyAuth(
+        user_id="test-user",
+        user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+
+
 @pytest.mark.asyncio
 async def test_create_and_get_tag():
     """
     Test creation of a new tag and retrieving its information
     """
+    # Create a mock auth object
+    mock_user_api_key_auth = create_mock_user_api_key_auth()
+    
     # Mock the prisma client and _get_tags_config and _save_tags_config
     with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
         "litellm.proxy.proxy_server.llm_router"
@@ -35,8 +47,11 @@ async def test_create_and_get_tag():
         "litellm.proxy.management_endpoints.tag_management_endpoints._add_tag_to_deployment"
     ) as mock_add_tag, patch(
         "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
-    ) as mock_get_models:
+    ) as mock_get_models, patch(
+        "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
+    ) as mock_auth:
         # Setup mocks
+        mock_auth.return_value = mock_user_api_key_auth
         mock_get_tags.return_value = {}
         mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
 
@@ -83,6 +98,9 @@ async def test_update_tag():
     """
     Test updating an existing tag
     """
+    # Create a mock auth object
+    mock_user_api_key_auth = create_mock_user_api_key_auth()
+    
     # Mock the prisma client and _get_tags_config and _save_tags_config
     with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
         "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
@@ -90,8 +108,11 @@ async def test_update_tag():
         "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
     ) as mock_save_tags, patch(
         "litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
-    ) as mock_get_models:
+    ) as mock_get_models, patch(
+        "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
+    ) as mock_auth:
         # Setup mocks for existing tag
+        mock_auth.return_value = mock_user_api_key_auth
         mock_get_tags.return_value = {
             "test-tag": {
                 "name": "test-tag",
@@ -129,13 +150,19 @@ async def test_delete_tag():
     """
     Test deleting a tag
     """
+    # Create a mock auth object
+    mock_user_api_key_auth = create_mock_user_api_key_auth()
+    
     # Mock the prisma client and _get_tags_config and _save_tags_config
     with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
         "litellm.proxy.management_endpoints.tag_management_endpoints._get_tags_config"
     ) as mock_get_tags, patch(
         "litellm.proxy.management_endpoints.tag_management_endpoints._save_tags_config"
-    ) as mock_save_tags:
+    ) as mock_save_tags, patch(
+        "litellm.proxy.auth.user_api_key_auth.user_api_key_auth"
+    ) as mock_auth:
         # Setup mocks for existing tag
+        mock_auth.return_value = mock_user_api_key_auth
         mock_get_tags.return_value = {
             "test-tag": {
                 "name": "test-tag",