From 45f188b04106a05fbdf2f441a1ddff9f9c17f0bf Mon Sep 17 00:00:00 2001 From: Jugal Bhatt Date: Thu, 14 Aug 2025 15:03:21 -0700 Subject: [PATCH] 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. --- .../test_tag_management_endpoints.py | 228 ++++++++++-------- 1 file changed, 134 insertions(+), 94 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 add08f5568..749ee4acd1 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,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