mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-14 02:22:54 +00:00
Add mock user API key authentication in tag management tests
This update integrates mock user API key authentication into the tag management endpoint tests, ensuring accurate simulation of user roles for creating, updating, and deleting tags. The changes enhance the reliability of the tests by properly setting up user authentication before executing test cases.
This commit is contained in:
Jugal Bhatt
@@ -14,6 +14,7 @@ from unittest.mock import patch
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
|
||||
from litellm.types.tag_management import TagDeleteRequest, TagInfoRequest, TagNewRequest
|
||||
|
||||
client = TestClient(app)
|
||||
@@ -24,58 +25,71 @@ async def test_create_and_get_tag():
|
||||
"""
|
||||
Test creation of a new tag and retrieving its information
|
||||
"""
|
||||
# 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"
|
||||
) as mock_router, 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, patch(
|
||||
"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:
|
||||
# Setup mocks
|
||||
mock_get_tags.return_value = {}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
|
||||
# Mock the user authentication
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
# 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"
|
||||
) as mock_router, 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, patch(
|
||||
"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:
|
||||
# Setup mocks
|
||||
mock_get_tags.return_value = {}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo"}
|
||||
|
||||
# Create a new tag
|
||||
tag_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for unit testing",
|
||||
"models": ["model-1"],
|
||||
}
|
||||
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag creation
|
||||
response = client.post("/tag/new", json=tag_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag created successfully"
|
||||
assert result["tag"]["name"] == "test-tag"
|
||||
assert result["tag"]["description"] == "Test tag for unit testing"
|
||||
|
||||
# Mock updated tag config for the get request
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
# Create a new tag
|
||||
tag_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for unit testing",
|
||||
"models": ["model-1"],
|
||||
"model_info": {"model-1": "gpt-3.5-turbo"},
|
||||
}
|
||||
}
|
||||
|
||||
# Test retrieving tag info
|
||||
info_data = {"names": ["test-tag"]}
|
||||
response = client.post("/tag/info", json=info_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "test-tag" in result
|
||||
assert result["test-tag"]["description"] == "Test tag for unit testing"
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag creation
|
||||
response = client.post("/tag/new", json=tag_data, headers=headers)
|
||||
print(f"response: {response.text}")
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag created successfully"
|
||||
assert result["tag"]["name"] == "test-tag"
|
||||
assert result["tag"]["description"] == "Test tag for unit testing"
|
||||
|
||||
# Mock updated tag config for the get request
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for unit testing",
|
||||
"models": ["model-1"],
|
||||
"model_info": {"model-1": "gpt-3.5-turbo"},
|
||||
}
|
||||
}
|
||||
|
||||
# Test retrieving tag info
|
||||
info_data = {"names": ["test-tag"]}
|
||||
response = client.post("/tag/info", json=info_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "test-tag" in result
|
||||
assert result["test-tag"]["description"] == "Test tag for unit testing"
|
||||
finally:
|
||||
# Clean up dependency overrides
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -83,16 +97,26 @@ async def test_update_tag():
|
||||
"""
|
||||
Test updating an existing tag
|
||||
"""
|
||||
# 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, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
|
||||
) as mock_get_models:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
# Mock the user authentication
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
# 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, patch(
|
||||
"litellm.proxy.management_endpoints.tag_management_endpoints._get_model_names"
|
||||
) as mock_get_models:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
"name": "test-tag",
|
||||
"description": "Original description",
|
||||
@@ -101,27 +125,30 @@ async def test_update_tag():
|
||||
"updated_at": "2023-01-01T00:00:00",
|
||||
"created_by": "user-123",
|
||||
}
|
||||
}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"}
|
||||
}
|
||||
mock_get_models.return_value = {"model-1": "gpt-3.5-turbo", "model-2": "gpt-4"}
|
||||
|
||||
# Update tag data
|
||||
update_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Updated description",
|
||||
"models": ["model-1", "model-2"],
|
||||
}
|
||||
# Update tag data
|
||||
update_data = {
|
||||
"name": "test-tag",
|
||||
"description": "Updated description",
|
||||
"models": ["model-1", "model-2"],
|
||||
}
|
||||
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag update
|
||||
response = client.post("/tag/update", json=update_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag updated successfully"
|
||||
assert result["tag"]["description"] == "Updated description"
|
||||
assert len(result["tag"]["models"]) == 2
|
||||
assert "model-2" in result["tag"]["models"]
|
||||
# Test tag update
|
||||
response = client.post("/tag/update", json=update_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag updated successfully"
|
||||
assert result["tag"]["description"] == "Updated description"
|
||||
assert len(result["tag"]["models"]) == 2
|
||||
assert "model-2" in result["tag"]["models"]
|
||||
finally:
|
||||
# Clean up dependency overrides
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -129,14 +156,24 @@ async def test_delete_tag():
|
||||
"""
|
||||
Test deleting a tag
|
||||
"""
|
||||
# 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:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
# Mock the user authentication
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
try:
|
||||
# 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:
|
||||
# Setup mocks for existing tag
|
||||
mock_get_tags.return_value = {
|
||||
"test-tag": {
|
||||
"name": "test-tag",
|
||||
"description": "Test tag for deletion",
|
||||
@@ -145,22 +182,25 @@ async def test_delete_tag():
|
||||
"updated_at": "2023-01-01T00:00:00",
|
||||
"created_by": "user-123",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Delete tag data
|
||||
delete_data = {"name": "test-tag"}
|
||||
# Delete tag data
|
||||
delete_data = {"name": "test-tag"}
|
||||
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
# Set admin access for the test
|
||||
headers = {"Authorization": f"Bearer sk-1234"}
|
||||
|
||||
# Test tag deletion
|
||||
response = client.post("/tag/delete", json=delete_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag deleted successfully"
|
||||
# Test tag deletion
|
||||
response = client.post("/tag/delete", json=delete_data, headers=headers)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["message"] == "Tag test-tag deleted successfully"
|
||||
|
||||
# Verify _save_tags_config was called without the deleted tag
|
||||
mock_save_tags.assert_called_once()
|
||||
# Verify _save_tags_config was called without the deleted tag
|
||||
mock_save_tags.assert_called_once()
|
||||
finally:
|
||||
# Clean up dependency overrides
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
Reference in New Issue
Block a user