diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/client.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/client.py index a28d03aaf..93519d00a 100644 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/client.py +++ b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/client.py @@ -14,7 +14,7 @@ # # SPDX-License-Identifier: Apache-2.0 -from google.auth import default, transport +from google import auth from openai import OpenAI as _OpenAI @@ -26,10 +26,10 @@ def __new__(cls, **openai_params) -> _OpenAI: location = openai_params.get('location') project_id = openai_params.get('project_id') if project_id: - credentials, _ = default() + credentials, _ = auth.default() else: - credentials, project_id = default() + credentials, project_id = auth.default() - credentials.refresh(transport.requests.Request()) + credentials.refresh(auth.transport.requests.Request()) base_url = f'https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project_id}/locations/{location}/endpoints/openapi' return _OpenAI(api_key=credentials.token, base_url=base_url) diff --git a/py/plugins/vertex-ai/tests/model_garden/test_client.py b/py/plugins/vertex-ai/tests/model_garden/test_client.py new file mode 100644 index 000000000..4c663e28f --- /dev/null +++ b/py/plugins/vertex-ai/tests/model_garden/test_client.py @@ -0,0 +1,75 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Unittests for VertexAI Model Garden OpenAI Client.""" + +from unittest.mock import MagicMock, patch + +from genkit.plugins.vertex_ai.model_garden.client import OpenAIClient + + +@patch('google.auth.default') +@patch('google.auth.transport.requests.Request') +@patch('openai.OpenAI') +def test_client_initialization_with_explicit_project_id( + mock_openai_cls, mock_request_cls, mock_default_auth +): + """Unittests for init client.""" + mock_location = "location" + mock_project_id = "project_id" + mock_token = "token" + + mock_credentials = MagicMock() + mock_credentials.token = mock_token + + mock_default_auth.return_value = (mock_credentials, "project_id") + + client_instance = OpenAIClient( + location=mock_location, + project_id=mock_project_id + ) + + mock_default_auth.assert_called_once() + mock_credentials.refresh.assert_called_once() + mock_request_cls.assert_called_once() + + assert client_instance is not None + + +@patch('google.auth.default') +@patch('google.auth.transport.requests.Request') +@patch('openai.OpenAI') +def test_client_initialization_without_explicit_project_id( + mock_openai_cls, mock_request_cls, mock_default_auth +): + """Unittests for init client.""" + mock_location = "location" + mock_token = "token" + + mock_credentials = MagicMock() + mock_credentials.token = mock_token + + mock_default_auth.return_value = (mock_credentials, "project_id") + + client_instance = OpenAIClient( + location=mock_location, + ) + + mock_default_auth.assert_called_once() + mock_credentials.refresh.assert_called_once() + mock_request_cls.assert_called_once() + + assert client_instance is not None diff --git a/py/plugins/vertex-ai/tests/model_garden/test_model_garden.py b/py/plugins/vertex-ai/tests/model_garden/test_model_garden.py new file mode 100644 index 000000000..d98df68e0 --- /dev/null +++ b/py/plugins/vertex-ai/tests/model_garden/test_model_garden.py @@ -0,0 +1,18 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Unittests for VertexAI Model Garden Models.""" +