test_create_vertex_fine_tune_jobs_mocked

This commit is contained in:
Ishaan Jaffer 2026-03-30 18:01:23 -07:00
parent 5596728cae
commit facb230fee
2 changed files with 63 additions and 32 deletions

View File

@ -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",

View File

@ -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():