diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index 71d2a7e7833..6f7c20166c5 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -397,32 +397,45 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models( monkeypatch: pytest.MonkeyPatch, ) -> None: - from unittest.mock import AsyncMock, MagicMock, call + from unittest.mock import AsyncMock, MagicMock + import litellm import litellm.proxy.proxy_server as proxy_server + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + from litellm.types.utils import CredentialItem rows: Final = [MagicMock()] prisma_client: Final = MagicMock() prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) router: Final = MagicMock() router.get_model_list.return_value = [] - collaborators: Final = MagicMock() - collaborators.attach_mock(AsyncMock(), "get_credentials") - collaborators.attach_mock(MagicMock(), "_add_deployment") + installed: Final = MagicMock() + async def load_credentials_from_db(prisma_client: object) -> None: + CredentialAccessor.upsert_credentials( + [ + CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db"}, + credential_info={}, + ) + ] + ) + + def install_models(db_models: object) -> None: + installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(litellm, "credential_list", []) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "llm_model_list", None) - monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", collaborators.get_credentials) - monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", collaborators._add_deployment) + monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models) assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True - assert collaborators.mock_calls == [ - call.get_credentials(prisma_client=prisma_client), - call._add_deployment(db_models=rows), - ] + installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"}) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index b732ee02f2c..1ca636b4c2d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2901,24 +2901,27 @@ async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconcilin from litellm.utils import load_credentials_from_list pc = ProxyConfig() - fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db") + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) monkeypatch.setattr(proxy_server, "general_settings", {}) - collaborators = MagicMock() - collaborators.attach_mock(AsyncMock(return_value=[]), "_get_models_from_db") - collaborators.attach_mock(AsyncMock(wraps=pc.get_credentials), "get_credentials") - collaborators.attach_mock(AsyncMock(return_value=None), "_update_llm_router") - monkeypatch.setattr(pc, "_get_models_from_db", collaborators._get_models_from_db) - monkeypatch.setattr(pc, "get_credentials", collaborators.get_credentials) - monkeypatch.setattr(pc, "_update_llm_router", collaborators._update_llm_router) + installed = MagicMock() + + async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]: + fake_prisma.db.litellm_credentialstable.find_many.return_value = [ + _encrypted_credential_row("openai-cred", "sk-from-db") + ] + return [MagicMock()] + + async def install_models(new_models: object, proxy_logging_obj: object) -> None: + installed(credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands) + monkeypatch.setattr(pc, "_update_llm_router", install_models) await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) - assert [name for name, _, _ in collaborators.mock_calls] == [ - "_get_models_from_db", - "get_credentials", - "_update_llm_router", - ] + installed.assert_called_once_with(credential={"api_key": "sk-from-db"}) assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} request_kwargs = {"litellm_credential_name": "openai-cred"} load_credentials_from_list(request_kwargs)