diff --git a/tests/batches_tests/test_fine_tuning_api.py b/tests/batches_tests/test_fine_tuning_api.py index 5dfbc0de71..cb570ff3c3 100644 --- a/tests/batches_tests/test_fine_tuning_api.py +++ b/tests/batches_tests/test_fine_tuning_api.py @@ -125,7 +125,6 @@ async def test_create_fine_tune_jobs_async(): @pytest.mark.asyncio() async def test_create_vertex_fine_tune_jobs_mocked(): - load_vertex_ai_credentials() # Define reusable variables for the test project_id = "633608382793" location = "us-central1" @@ -164,7 +163,10 @@ async def test_create_vertex_fine_tune_jobs_mocked(): with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=mock_response, - ) as mock_post: + ) as mock_post, patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token", + return_value=("fake-token", project_id), + ): create_fine_tuning_response = await litellm.acreate_fine_tuning_job( model=base_model, custom_llm_provider="vertex_ai", @@ -218,7 +220,6 @@ async def test_create_vertex_fine_tune_jobs_mocked(): @pytest.mark.asyncio() async def test_create_vertex_fine_tune_jobs_mocked_with_hyperparameters(): - load_vertex_ai_credentials() # Define reusable variables for the test project_id = "633608382793" location = "us-central1" @@ -257,7 +258,10 @@ async def test_create_vertex_fine_tune_jobs_mocked_with_hyperparameters(): with patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", return_value=mock_response, - ) as mock_post: + ) as mock_post, patch( + "litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token", + return_value=("fake-token", project_id), + ): create_fine_tuning_response = await litellm.acreate_fine_tuning_job( model=base_model, custom_llm_provider="vertex_ai", diff --git a/tests/litellm_utils_tests/test_secret_manager.py b/tests/litellm_utils_tests/test_secret_manager.py index a3e7a8496f..a1190193ea 100644 --- a/tests/litellm_utils_tests/test_secret_manager.py +++ b/tests/litellm_utils_tests/test_secret_manager.py @@ -1,3 +1,4 @@ +import base64 import os import sys import time @@ -24,7 +25,7 @@ from litellm.secret_managers.main import ( get_secret, _should_read_secret_from_secret_manager, ) -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch, MagicMock def load_vertex_ai_credentials(): @@ -221,7 +222,6 @@ def test_oidc_env_path(): del os.environ[env_var_name] -@pytest.mark.flaky(retries=6, delay=1) def test_google_secret_manager(): """ Test that we can get a secret from Google Secret Manager @@ -230,44 +230,71 @@ def test_google_secret_manager(): from litellm.secret_managers.google_secret_manager import GoogleSecretManager - load_vertex_ai_credentials() - secret_manager = GoogleSecretManager() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "payload": { + "data": base64.b64encode(b"anything").decode("utf-8"), + } + } - secret_val = secret_manager.get_secret_from_google_secret_manager( - secret_name="OPENAI_API_KEY" - ) - print("secret_val: {}".format(secret_val)) + with patch( + "litellm.proxy.proxy_server.premium_user", True + ), patch.object( + GoogleSecretManager, + "sync_construct_request_headers", + return_value={"Authorization": "Bearer mock_token"}, + ): + secret_manager = GoogleSecretManager() + secret_manager.sync_httpx_client = MagicMock() + secret_manager.sync_httpx_client.get.return_value = mock_response - assert ( - secret_val == "anything" - ), "did not get expected secret value. expect 'anything', got '{}'".format( - secret_val - ) + secret_val = secret_manager.get_secret_from_google_secret_manager( + secret_name="OPENAI_API_KEY" + ) + print("secret_val: {}".format(secret_val)) + + assert ( + secret_val == "anything" + ), "did not get expected secret value. expect 'anything', got '{}'".format( + secret_val + ) + + secret_manager.sync_httpx_client.get.assert_called_once() + call_url = secret_manager.sync_httpx_client.get.call_args[1]["url"] + assert "projects/litellm-ci-cd/secrets/OPENAI_API_KEY" in call_url def test_google_secret_manager_read_in_memory(): """ - Test that Google Secret manager returs in memory value when it exists + Test that Google Secret manager returns in memory value when it exists """ from litellm.secret_managers.google_secret_manager import GoogleSecretManager - load_vertex_ai_credentials() os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = "litellm-ci-cd" - secret_manager = GoogleSecretManager() - secret_manager.cache.cache_dict["UNIQUE_KEY"] = None - secret_manager.cache.cache_dict["UNIQUE_KEY_2"] = "lite-llm" - secret_val = secret_manager.get_secret_from_google_secret_manager( - secret_name="UNIQUE_KEY" - ) - print("secret_val: {}".format(secret_val)) - assert secret_val == None + with patch( + "litellm.proxy.proxy_server.premium_user", True + ), patch.object( + GoogleSecretManager, + "sync_construct_request_headers", + return_value={"Authorization": "Bearer mock_token"}, + ): + secret_manager = GoogleSecretManager() + secret_manager.cache.cache_dict["UNIQUE_KEY"] = None + secret_manager.cache.cache_dict["UNIQUE_KEY_2"] = "lite-llm" - secret_val = secret_manager.get_secret_from_google_secret_manager( - secret_name="UNIQUE_KEY_2" - ) - print("secret_val: {}".format(secret_val)) - assert secret_val == "lite-llm" + secret_val = secret_manager.get_secret_from_google_secret_manager( + secret_name="UNIQUE_KEY" + ) + print("secret_val: {}".format(secret_val)) + assert secret_val is None + + secret_val = secret_manager.get_secret_from_google_secret_manager( + secret_name="UNIQUE_KEY_2" + ) + print("secret_val: {}".format(secret_val)) + assert secret_val == "lite-llm" def test_should_read_secret_from_secret_manager():