diff --git a/tests/litellm/proxy/test_proxy_server.py b/tests/litellm/proxy/test_proxy_server.py new file mode 100644 index 0000000000..3f0a57e284 --- /dev/null +++ b/tests/litellm/proxy/test_proxy_server.py @@ -0,0 +1,74 @@ +import importlib +import json +import os +import socket +import subprocess +import sys +from unittest.mock import AsyncMock, MagicMock, mock_open, patch + +import click +import httpx +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system-path + +import litellm + + +@pytest.mark.asyncio +async def test_initialize_scheduled_jobs_credentials(): + """ + Test that get_credentials is only called when store_model_in_db is True + """ + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + # Mock dependencies + mock_prisma_client = MagicMock() + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + + with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch( + "litellm.proxy.proxy_server.store_model_in_db", False + ): # set store_model_in_db to False + + # Test when store_model_in_db is False + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + # Verify get_credentials was not called + mock_proxy_config.get_credentials.assert_not_called() + + # Now test with store_model_in_db = True + with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True): + + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + # Verify get_credentials was called both directly and scheduled + assert mock_proxy_config.get_credentials.call_count == 1 # Direct call + + # Verify a scheduled job was added for get_credentials + mock_scheduler_calls = [ + call[0] for call in mock_proxy_config.get_credentials.mock_calls + ] + assert len(mock_scheduler_calls) > 0