diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 4ec3dcbb4cf..7e67aee8d73 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -7,7 +7,7 @@ on: jobs: test: runs-on: ubuntu-latest - timeout-minutes: 20 + timeout-minutes: 25 steps: - uses: actions/checkout@v4 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..749ee4acd16 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 diff --git a/tests/test_litellm/test_cost_calculation_log_level.py b/tests/test_litellm/test_cost_calculation_log_level.py index 4380ae8bf62..3925ea751af 100644 --- a/tests/test_litellm/test_cost_calculation_log_level.py +++ b/tests/test_litellm/test_cost_calculation_log_level.py @@ -17,46 +17,55 @@ def test_cost_calculation_uses_debug_level(caplog): This ensures cost calculation details don't appear in production logs. Part of fix for issue #9815. """ - # Create a mock completion response - mock_response = { - "id": "test", - "object": "chat.completion", - "created": 1234567890, - "model": "gpt-3.5-turbo", - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "Test response"}, - "finish_reason": "stop" - }], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "total_tokens": 30 + # Ensure verbose_logger is set to DEBUG level to capture the debug logs + from litellm._logging import verbose_logger + original_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) + + try: + # Create a mock completion response + mock_response = { + "id": "test", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-3.5-turbo", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Test response"}, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + } } - } - - # Test that cost calculation logs are at DEBUG level - with caplog.at_level(logging.DEBUG): - try: - cost = completion_cost( - completion_response=mock_response, - model="gpt-3.5-turbo" - ) - except Exception: - pass # Cost calculation may fail, but we're checking log levels - - # Find the cost calculation log records - cost_calc_records = [ - record for record in caplog.records - if "selected model name for cost calculation" in record.message - ] - - # Verify that cost calculation logs are at DEBUG level - assert len(cost_calc_records) > 0, "No cost calculation logs found" - - for record in cost_calc_records: - assert record.levelno == logging.DEBUG, \ - f"Cost calculation log should be DEBUG level, but was {record.levelname}" + + # Test that cost calculation logs are at DEBUG level + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + try: + cost = completion_cost( + completion_response=mock_response, + model="gpt-3.5-turbo" + ) + except Exception: + pass # Cost calculation may fail, but we're checking log levels + + # Find the cost calculation log records + cost_calc_records = [ + record for record in caplog.records + if "selected model name for cost calculation" in record.message + ] + + # Verify that cost calculation logs are at DEBUG level + assert len(cost_calc_records) > 0, "No cost calculation logs found" + + for record in cost_calc_records: + assert record.levelno == logging.DEBUG, \ + f"Cost calculation log should be DEBUG level, but was {record.levelname}" + finally: + # Restore original logger level + verbose_logger.setLevel(original_level) def test_batch_cost_calculation_uses_debug_level(caplog): @@ -65,29 +74,38 @@ def test_batch_cost_calculation_uses_debug_level(caplog): """ from litellm.cost_calculator import batch_cost_calculator from litellm.types.utils import Usage + from litellm._logging import verbose_logger - # Create a mock usage object - usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + # Ensure verbose_logger is set to DEBUG level to capture the debug logs + original_level = verbose_logger.level + verbose_logger.setLevel(logging.DEBUG) - # Test that batch cost calculation logs are at DEBUG level - with caplog.at_level(logging.DEBUG): - try: - batch_cost_calculator( - usage=usage, - model="gpt-3.5-turbo", - custom_llm_provider="openai" - ) - except Exception: - pass # May fail, but we're checking log levels - - # Find batch cost calculation log records - batch_cost_records = [ - record for record in caplog.records - if "Calculating batch cost per token" in record.message - ] - - # Verify logs exist and are at DEBUG level - if batch_cost_records: # May not always log depending on the code path - for record in batch_cost_records: - assert record.levelno == logging.DEBUG, \ - f"Batch cost calculation log should be DEBUG level, but was {record.levelname}" \ No newline at end of file + try: + # Create a mock usage object + usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + + # Test that batch cost calculation logs are at DEBUG level + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + try: + batch_cost_calculator( + usage=usage, + model="gpt-3.5-turbo", + custom_llm_provider="openai" + ) + except Exception: + pass # May fail, but we're checking log levels + + # Find batch cost calculation log records + batch_cost_records = [ + record for record in caplog.records + if "Calculating batch cost per token" in record.message + ] + + # Verify logs exist and are at DEBUG level + if batch_cost_records: # May not always log depending on the code path + for record in batch_cost_records: + assert record.levelno == logging.DEBUG, \ + f"Batch cost calculation log should be DEBUG level, but was {record.levelname}" + finally: + # Restore original logger level + verbose_logger.setLevel(original_level) \ No newline at end of file