google-genai 1.55.0__py3-none-any.whl → 1.57.0__py3-none-any.whl

google/genai/tests/client/test_client_requests.py

@@ -19,8 +19,6 @@
 from ... import _api_client as api_client
 from ... import Client
 from ... import types
-from unittest import mock
-import pytest
 
 
 def build_test_client(monkeypatch):
@@ -220,136 +218,4 @@ def test_build_request_with_custom_base_url_no_env_vars(monkeypatch):
       'test/path',
       {'key': 'value'},
   )
-  assert request.url == 'https://custom-base-url.com'
-
-
-def test_interactions_vertex_auth_header():
-  from ..._api_client import BaseApiClient
-  from ..._interactions._base_client import SyncAPIClient
-  from httpx import Client as HTTPClient
-
-  creds = mock.Mock()
-  creds.quota_project_id = "test-quota-project"
-  client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
-
-  with (
-      mock.patch.object(
-          BaseApiClient, "_access_token", return_value='fake-vertex-token'
-      ) as  mock_access_token,
-      mock.patch.object(
-          HTTPClient, "send",
-          return_value=mock.Mock(),
-      ) as mock_send,
-  ):
-
-    response = client.interactions.create(
-        model='gemini-2.5-flash',
-        input='What is the largest planet in our solar system?',
-    )
-
-    mock_send.assert_called_once()
-    mock_access_token.assert_called_once()
-    args, kwargs = mock_send.call_args
-    headers = args[0].headers
-    assert any(
-        key == "authorization" and value == 'Bearer fake-vertex-token'
-        for key, value in headers.items())
-    assert any(
-        key == "x-goog-user-project" and value == 'test-quota-project'
-        for key, value in headers.items())
-
-def test_interactions_vertex_key_no_auth_header():
-  from ..._api_client import BaseApiClient
-  from httpx import Client as HTTPClient
-
-  creds = mock.Mock()
-  client = Client(vertexai=True, api_key='test-api-key')
-
-  with (
-      mock.patch.object(
-          BaseApiClient, "_access_token", return_value='fake-vertex-token'
-      ) as  mock_access_token,
-      mock.patch.object(
-          HTTPClient, "send",
-          return_value=mock.Mock(),
-      ) as mock_send,
-  ):
-
-    response = client.interactions.create(
-        model='gemini-2.5-flash',
-        input='What is the largest planet in our solar system?',
-    )
-
-    mock_send.assert_called_once()
-    mock_access_token.assert_not_called()
-    args, kwargs = mock_send.call_args
-    headers = args[0].headers
-    assert any(
-        key == "x-goog-api-key" and value == 'test-api-key'
-        for key, value in headers.items())
-
-@pytest.mark.asyncio
-async def test_async_interactions_vertex_auth_header():
-  from ..._api_client import BaseApiClient
-  from ..._interactions._base_client import SyncAPIClient
-  from ..._api_client import AsyncHttpxClient
-
-  creds = mock.Mock()
-  creds.quota_project_id = "test-quota-project"
-  client = Client(vertexai=True, project='test-project', location='us-central1', credentials=creds)
-
-  with (
-      mock.patch.object(
-          BaseApiClient, "_async_access_token", return_value='fake-vertex-token'
-      ) as  mock_access_token,
-      mock.patch.object(
-          AsyncHttpxClient, "send",
-          return_value=mock.Mock(),
-      ) as mock_send,
-  ):
-
-    response = await client.aio.interactions.create(
-        model='gemini-2.5-flash',
-        input='What is the largest planet in our solar system?',
-    )
-
-    mock_send.assert_called_once()
-    mock_access_token.assert_called_once()
-    args, kwargs = mock_send.call_args
-    headers = args[0].headers
-    assert any(
-        key == "authorization" and value == 'Bearer fake-vertex-token'
-        for key, value in headers.items())
-    assert any(
-        key == "x-goog-user-project" and value == 'test-quota-project'
-        for key, value in headers.items())
-
-@pytest.mark.asyncio
-async def test_async_interactions_vertex_key_no_auth_header():
-  from ..._api_client import BaseApiClient
-  from ..._api_client import AsyncHttpxClient
-  creds = mock.Mock()
-  client = Client(vertexai=True, api_key='test-api-key')
-
-  with (
-      mock.patch.object(
-          BaseApiClient, "_async_access_token", return_value='fake-vertex-token'
-      ) as  mock_access_token,
-      mock.patch.object(
-          AsyncHttpxClient, "send",
-          return_value=mock.Mock(),
-      ) as mock_send,
-  ):
-
-    response = await client.aio.interactions.create(
-        model='gemini-2.5-flash',
-        input='What is the largest planet in our solar system?',
-    )
-
-    mock_send.assert_called_once()
-    mock_access_token.assert_not_called()
-    args, kwargs = mock_send.call_args
-    headers = args[0].headers
-    assert any(
-        key == "x-goog-api-key" and value == 'test-api-key'
-        for key, value in headers.items())
+  assert request.url == 'https://custom-base-url.com'

google/genai/tests/files/test_register.py

@@ -0,0 +1,272 @@
+# 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.
+#
+
+
+"""Test files register method."""
+
+import json
+from unittest import mock
+
+from google.auth import credentials
+import httpx
+import pytest
+
+from ... import _api_client
+from ... import Client
+from ... import types
+from .. import pytest_helper
+
+
+class FakeCredentials(credentials.Credentials):
+
+  def __init__(self, token="fake_token", expired=False, quota_project_id=None):
+    super().__init__()
+    self.token = token
+    self._expired = expired
+    self._quota_project_id = quota_project_id
+    self.refresh_count = 0
+
+  @property
+  def expired(self):
+    return self._expired
+
+  @property
+  def quota_project_id(self):
+    return self._quota_project_id
+
+  def refresh(self, request):
+    self.refresh_count += 1
+    self.token = "refreshed_token"
+    self._expired = False
+
+
+@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
+def test_simple_token(mock_request):
+  client = Client(api_key="dummy_key")
+  captured_request = None
+
+  def side_effect(self, http_request, stream=False):
+    nonlocal captured_request
+    captured_request = http_request
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    response = client.files.register_files(
+        auth=FakeCredentials(token="test_token"),
+        uris=["gs://test-bucket/test-file-1.txt"],
+    )
+
+    assert len(response.files) == 1
+    assert response.files[0].uri == "files/abc"
+    assert captured_request.headers["authorization"] == "Bearer test_token"
+
+
+@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
+def test_token_refresh(mock_request):
+  client = Client(api_key="dummy_key")
+  captured_request = None
+
+  def side_effect(self, http_request, stream=False):
+    nonlocal captured_request
+    captured_request = http_request
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    creds = FakeCredentials(expired=True)
+    response = client.files.register_files(
+        auth=creds,
+        uris=["gs://test-bucket/test-file-1.txt"],
+    )
+    assert creds.refresh_count == 1
+    assert len(response.files) == 1
+    assert response.files[0].uri == "files/abc"
+    assert captured_request.headers["authorization"] == "Bearer refreshed_token"
+
+
+@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
+def test_quota_project(mock_request):
+  client = Client(api_key="dummy_key")
+  captured_request = None
+
+  def side_effect(self, http_request, stream=False):
+    nonlocal captured_request
+    captured_request = http_request
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    creds = FakeCredentials(quota_project_id="test_project")
+    response = client.files.register_files(
+        auth=creds,
+        uris=["gs://test-bucket/test-file-1.txt"],
+    )
+    assert len(response.files) == 1
+    assert response.files[0].uri == "files/abc"
+    assert captured_request.headers["x-goog-user-project"] == "test_project"
+
+
+@mock.patch.object(_api_client.BaseApiClient, "_request_once", autospec=True)
+def test_multiple_uris(mock_request):
+  client = Client(api_key="dummy_key")
+
+  def side_effect(self, http_request, stream=False):
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[
+            json.dumps({"files": [{"uri": "files/abc"}, {"uri": "files/def"}]})
+        ],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    response = client.files.register_files(
+        auth=FakeCredentials(),
+        uris=[
+            "gs://test-bucket/test-file-1.txt",
+            "gs://test-bucket/test-file-2.txt",
+        ],
+    )
+    assert len(response.files) == 2
+    assert response.files[0].uri == "files/abc"
+    assert response.files[1].uri == "files/def"
+
+
+@pytest.mark.asyncio
+@mock.patch.object(
+    _api_client.BaseApiClient, "_async_request_once", autospec=True
+)
+async def test_async_single(mock_request):
+  client = Client(api_key="dummy_key")
+
+  async def side_effect(self, http_request, stream=False):
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    response = await client.aio.files.register_files(
+        auth=FakeCredentials(),
+        uris=["gs://test-bucket/test-file-1.txt"],
+    )
+
+    assert len(response.files) == 1
+    assert response.files[0].uri == "files/abc"
+
+
+@pytest.mark.asyncio
+@mock.patch.object(
+    _api_client.BaseApiClient, "_async_request_once", autospec=True
+)
+async def test_async_token_refresh(mock_request):
+  client = Client(api_key="dummy_key")
+  captured_request = None
+
+  async def side_effect(self, http_request, stream=False):
+    nonlocal captured_request
+    captured_request = http_request
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    creds = FakeCredentials(expired=True)
+    response = await client.aio.files.register_files(
+        auth=creds,
+        uris=["gs://test-bucket/test-file-1.txt"],
+    )
+    assert creds.refresh_count == 1
+    assert len(response.files) == 1
+    assert response.files[0].uri == "files/abc"
+    assert captured_request.headers["authorization"] == "Bearer refreshed_token"
+
+
+@pytest.mark.asyncio
+@mock.patch.object(
+    _api_client.BaseApiClient, "_async_request_once", autospec=True
+)
+async def test_async_quota_project(mock_request):
+  client = Client(api_key="dummy_key")
+  captured_request = None
+
+  async def side_effect(self, http_request, stream=False):
+    nonlocal captured_request
+    captured_request = http_request
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[json.dumps({"files": [{"uri": "files/abc"}]})],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    creds = FakeCredentials(quota_project_id="test_project")
+    response = await client.aio.files.register_files(
+        auth=creds,
+        uris=["gs://test-bucket/test-file-1.txt"],
+    )
+    assert len(response.files) == 1
+    assert response.files[0].uri == "files/abc"
+    assert captured_request.headers["x-goog-user-project"] == "test_project"
+
+
+@pytest.mark.asyncio
+@mock.patch.object(
+    _api_client.BaseApiClient, "_async_request_once", autospec=True
+)
+async def test_async_multiple_uris(mock_request):
+  client = Client(api_key="dummy_key")
+
+  async def side_effect(self, http_request, stream=False):
+    return _api_client.HttpResponse(
+        headers={},
+        response_stream=[
+            json.dumps({"files": [{"uri": "files/abc"}, {"uri": "files/def"}]})
+        ],
+    )
+
+  mock_request.side_effect = side_effect
+
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    response = await client.aio.files.register_files(
+        auth=FakeCredentials(),
+        uris=[
+            "gs://test-bucket/test-file-1.txt",
+            "gs://test-bucket/test-file-2.txt",
+        ],
+    )
+    assert len(response.files) == 2
+    assert response.files[0].uri == "files/abc"
+    assert response.files[1].uri == "files/def"

google/genai/tests/files/test_register_table.py

@@ -0,0 +1,70 @@
+# 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.
+#
+
+
+"""Test files get method."""
+
+import pytest
+from ... import types
+from ... import Client
+from ... import _api_client
+from .. import pytest_helper
+import google.auth
+
+
+# $ gcloud config set project vertex-sdk-dev
+# $ gcloud auth application-default login --no-launch-browser --scopes="https://www.googleapis.com/auth/cloud-platform,https://www.googleapis.com/auth/devstorage.read_only"
+def get_headers():
+  try:
+    credentials, _ = google.auth.default()
+    token = _api_client.get_token_from_credentials(None, credentials)
+    headers = {
+      "Authorization": f"Bearer {token}",}
+    if credentials.quota_project_id:
+      headers["x-goog-user-project"] = credentials.quota_project_id
+  except google.auth.exceptions.DefaultCredentialsError:
+    # So this can run in replay mode without credentials.
+    headers = {}
+
+
+test_table: list[pytest_helper.TestTableItem] = [
+    pytest_helper.TestTableItem(
+        name='test_register',
+        parameters=types._RegisterFilesParameters(uris=['gs://unified-genai-dev/image.jpg']),
+        exception_if_vertex='only supported in the Gemini Developer client',
+        skip_in_api_mode=(
+            'The files have a TTL, they cannot be reliably retrieved for a long'
+            ' time.'
+        ),
+    ),
+]
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='files._register_files',
+    test_table=test_table,
+    http_options={
+        'headers': get_headers(),
+    },
+)
+
+
+@pytest.mark.asyncio
+async def test_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    files = await client.aio.files._register_files(uris=['gs://unified-genai-dev/image.jpg'])
+    assert files.files
+    assert files.files[0].mime_type == 'image/jpeg'