google-genai 1.56.0__py3-none-any.whl → 1.58.0__py3-none-any.whl

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'

google/genai/tests/files/test_upload.py

@@ -0,0 +1,255 @@
+# 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 upload method."""
+
+
+import io
+import pathlib
+import pytest
+from ... import types
+from ... import errors
+from .. import pytest_helper
+
+# Upload method is not pydantic.
+test_table: list[pytest_helper.TestTableItem] = []
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method='files.upload',
+    test_table=test_table,
+)
+
+
+def test_image_png_upload(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(file='tests/data/google.png')
+    assert file.name.startswith('files/')
+
+def test_image_png_upload_with_path(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    p = pathlib.Path('tests/data/google.png')
+    file = client.files.upload(
+        file=p,
+        config=types.UploadFileConfig(display_name='test_image_png_path'),
+    )
+    assert file.name.startswith('files/')
+
+def test_image_png_upload_with_bytesio(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    with open('tests/data/google.png', 'rb') as f:
+      with io.BytesIO(f.read()) as buffer:
+        file = client.files.upload(
+            file=buffer,
+            config=types.UploadFileConfig(mime_type='image/png'),
+        )
+
+        assert file.name.startswith('files/')
+
+def test_image_png_upload_with_fd(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    with open('tests/data/google.png', 'rb') as f:
+      file = client.files.upload(
+          file=f,
+          config=types.UploadFileConfig(mime_type='image/png'),
+      )
+
+    assert file.name.startswith('files/')
+
+def test_image_png_upload_with_config(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/google.png',
+        config=types.UploadFileConfig(display_name='test_image_png'),
+    )
+    assert file.name.startswith('files/')
+
+
+def test_image_png_upload_with_config_dict(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/google.png', config={'display_name': 'test_image_png'}
+    )
+    assert file.name.startswith('files/')
+
+
+def test_image_jpg_upload(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(file='tests/data/google.jpg')
+    assert file.name.startswith('files/')
+
+
+def test_image_jpg_upload_with_config(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/google.jpg',
+        config=types.UploadFileConfig(display_name='test_image_jpg'),
+    )
+    assert file.name.startswith('files/')
+
+
+def test_image_jpg_upload_with_config_dict(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/google.jpg', config={'display_name': 'test_image_jpg'}
+    )
+    assert file.name.startswith('files/')
+
+
+def test_application_pdf_file_upload(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(file='tests/data/story.pdf')
+    assert file.name.startswith('files/')
+
+
+def test_application_pdf_upload_with_config(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/story.pdf',
+        config=types.UploadFileConfig(display_name='test_application_pdf'),
+    )
+    assert file.name.startswith('files/')
+
+
+def test_application_pdf_upload_with_config_dict(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/story.pdf',
+        config={'display_name': 'test_application_pdf'},
+    )
+    assert file.name.startswith('files/')
+
+
+def test_video_mp4_file_upload(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(file='tests/data/animal.mp4')
+    assert file.name.startswith('files/')
+
+
+def test_video_mp4_upload_with_config(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/animal.mp4',
+        config=types.UploadFileConfig(display_name='test_video_mp4'),
+    )
+    assert file.name.startswith('files/')
+
+
+def test_video_mp4_upload_with_config_dict(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/animal.mp4', config={'display_name': 'test_video_mp4'}
+    )
+    assert file.name.startswith('files/')
+
+
+def test_audio_m4a_file_upload(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/pixel.m4a',
+        config=types.UploadFileConfig(mime_type='audio/mp4'),
+    )
+    assert file.name.startswith('files/')
+
+
+def test_audio_m4a_upload_with_config(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/pixel.m4a',
+        config=types.UploadFileConfig(
+            display_name='test_audio_m4a', mime_type='audio/mp4'
+        ),
+    )
+    assert file.name.startswith('files/')
+
+
+def test_audio_m4a_upload_with_config_dict(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = client.files.upload(
+        file='tests/data/pixel.m4a',
+        config={'display_name': 'test_audio_m4a', 'mime_type': 'audio/mp4'},
+    )
+    assert file.name.startswith('files/')
+
+
+def test_bad_mime_type(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    with pytest.raises(errors.APIError, match="Unsupported MIME"):
+      file = client.files.upload(
+          file=io.BytesIO(b'test'),
+          config={'mime_type': 'bad/mime_type'},
+      )
+
+
+@pytest.mark.asyncio
+async def test_image_upload_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = await client.aio.files.upload(file='tests/data/google.png')
+    assert file.name.startswith('files/')
+
+
+@pytest.mark.asyncio
+async def test_image_upload_with_config_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = await client.aio.files.upload(
+        file='tests/data/google.png',
+        config=types.UploadFileConfig(display_name='test_image'),
+    )
+    assert file.name.startswith('files/')
+
+
+@pytest.mark.asyncio
+async def test_image_upload_with_config_dict_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    file = await client.aio.files.upload(
+        file='tests/data/google.png',
+        config={
+            'display_name': 'test_image',
+            'http_options': {'timeout': '8000'},
+        },
+    )
+    assert file.name.startswith('files/')
+
+
+@pytest.mark.asyncio
+async def test_image_upload_with_bytesio_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    with open('tests/data/google.png', 'rb') as f:
+      buffer = io.BytesIO(f.read())
+    file = await client.aio.files.upload(
+        file=buffer,
+        config=types.UploadFileConfig(
+            mime_type='image/png'),
+    )
+    assert file.name.startswith('files/')
+
+
+@pytest.mark.asyncio
+async def test_unknown_path_upload_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    try:
+      await client.aio.files.upload(file='unknown_path')
+    except FileNotFoundError as e:
+      assert 'is not a valid file path' in str(e)
+
+@pytest.mark.asyncio
+async def test_bad_mime_type_async(client):
+  with pytest_helper.exception_if_vertex(client, ValueError):
+    with pytest.raises(errors.APIError, match="Unsupported MIME"):
+      file = await client.aio.files.upload(
+          file=io.BytesIO(b'test'),
+          config={'mime_type': 'bad/mime_type'},
+      )

google/genai/tests/imports/test_no_optional_imports.py

@@ -0,0 +1,28 @@
+# 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.
+#
+
+import os
+import pytest
+
+IS_NOT_GITHUB_ACTIONS = os.getenv('GITHUB_ACTIONS') != 'true'
+
+
+@pytest.mark.skipif(IS_NOT_GITHUB_ACTIONS,
+                    reason='This test is only run on GitHub Actions.')
+def test_library_can_be_imported_without_optional_dependencies():
+  """Tests that the library can be imported without optional dependencies.
+  """
+  from google import genai
+  from google.genai import types