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