scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

Add mock user API key authentication in tag management tests

Browse source 
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.
This commit is contained in:
Jugal Bhatt 2025-08-14 14:40:49 -07:00
parent a6e55c0447
commit 5fc0803b94
1 changed files with 30 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py
View file

@ -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",