Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -442,34 +442,36 @@ def test_validate_environment_with_optional_params(self, mock_ensure_access_toke
"""Test that validate_environment accepts and uses optional_params for credentials."""
# Mock the authentication
mock_ensure_access_token.return_value = ("test-access-token", "test-project-123")

optional_params = {
"vertex_credentials": "path/to/credentials.json",
"vertex_project": "custom-project-id",
"query": "test query",
"documents": ["doc1"]
}

headers = self.config.validate_environment(
headers={},
model=self.model,
api_key=None,
optional_params=optional_params
)

# Verify that _ensure_access_token was called with the credentials from optional_params
mock_ensure_access_token.assert_called_once()
call_args = mock_ensure_access_token.call_args
# The first call argument should be credentials (which will be the value from optional_params)
# We can't check the exact value easily due to how get_vertex_ai_credentials pops values,
# but we can verify the headers were set correctly

expected_headers = {
"Authorization": "Bearer test-access-token",
"Content-Type": "application/json",
"X-Goog-User-Project": "test-project-123"
}
assert headers == expected_headers

# Mock get_vertex_ai_credentials to prevent it from trying to load the fake credential file
with patch.object(self.config, 'get_vertex_ai_credentials', return_value=None):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mock targets wrong method

validate_environment (in transformation.py:80) calls self.safe_get_vertex_ai_credentials(litellm_params), not self.get_vertex_ai_credentials(...). Since safe_get_vertex_ai_credentials is the method actually invoked, this patch.object for get_vertex_ai_credentials has no effect — it mocks a method that is never called in this code path.

The test already works because _ensure_access_token is mocked via the @patch decorator on the test method, which intercepts the call before any real credential loading occurs.

If the intent is to prevent safe_get_vertex_ai_credentials from returning the fake credential string, the mock should target that method instead:

Suggested change
with patch.object(self.config, 'get_vertex_ai_credentials', return_value=None):
with patch.object(self.config, 'safe_get_vertex_ai_credentials', return_value=None):

Note: the same pattern may need to be applied to test_validate_environment_preserves_optional_params_for_get_complete_url (line 476) which also passes fake credentials without mocking.

headers = self.config.validate_environment(
headers={},
model=self.model,
api_key=None,
optional_params=optional_params
)

# Verify that _ensure_access_token was called with the credentials from optional_params
mock_ensure_access_token.assert_called_once()
call_args = mock_ensure_access_token.call_args
# The first call argument should be credentials (which will be the value from optional_params)
# We can't check the exact value easily due to how get_vertex_ai_credentials pops values,
# but we can verify the headers were set correctly

expected_headers = {
"Authorization": "Bearer test-access-token",
"Content-Type": "application/json",
"X-Goog-User-Project": "test-project-123"
}
assert headers == expected_headers

@patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token')
def test_validate_environment_preserves_optional_params_for_get_complete_url(
Expand All @@ -487,17 +489,19 @@ def test_validate_environment_preserves_optional_params_for_get_complete_url(
"vertex_project": "custom-project-id",
}

# Call validate_environment first – this previously popped the values in-place
self.config.validate_environment(
headers={},
model=self.model,
api_key=None,
optional_params=optional_params,
)
# Mock get_vertex_ai_credentials to prevent it from trying to load the fake credential file
with patch.object(self.config, 'get_vertex_ai_credentials', return_value=None):
# Call validate_environment first – this previously popped the values in-place
self.config.validate_environment(
headers={},
model=self.model,
api_key=None,
optional_params=optional_params,
)

# Ensure the original optional_params dict still retains the vertex keys
assert optional_params["vertex_credentials"] == "path/to/credentials.json"
assert optional_params["vertex_project"] == "custom-project-id"
# Ensure the original optional_params dict still retains the vertex keys
assert optional_params["vertex_credentials"] == "path/to/credentials.json"
assert optional_params["vertex_project"] == "custom-project-id"

# get_complete_url should still be able to access the vertex params
with patch('litellm.llms.vertex_ai.rerank.transformation.get_secret_str', return_value=None):
Expand Down
Loading