first vommit
This commit is contained in:
@@ -0,0 +1,17 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
"""Tests for the Google GenAI SDK's files module."""
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,46 @@
|
||||
# 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 delete method."""
|
||||
|
||||
import pytest
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_delete',
|
||||
parameters=types._DeleteFileParameters(name='files/1g583ke2xdsn'),
|
||||
exception_if_vertex='only supported in Gemini Developer API mode',
|
||||
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.delete',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.get(name='files/n1gls7dyh90q')
|
||||
@@ -0,0 +1,85 @@
|
||||
# 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 pathlib
|
||||
import pytest
|
||||
from ... import _transformers as t
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = []
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='t.t_file_name',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
pytest_plugins = ('pytest_asyncio',)
|
||||
|
||||
|
||||
def test_name_transform_name(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
for file in client.files.list():
|
||||
if file.download_uri is not None:
|
||||
break
|
||||
else:
|
||||
raise ValueError('No files found with a `download_uri`.')
|
||||
|
||||
file_id = file.name.split('/')[-1]
|
||||
video = types.Video(uri=file.download_uri)
|
||||
generated_video = types.GeneratedVideo(video=video)
|
||||
for f in [
|
||||
file,
|
||||
file_id,
|
||||
file.name,
|
||||
file.uri,
|
||||
file.download_uri,
|
||||
video,
|
||||
generated_video,
|
||||
]:
|
||||
name = t.t_file_name(f)
|
||||
assert name == file_id
|
||||
|
||||
|
||||
def test_basic_download(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
for file in client.files.list():
|
||||
if file.download_uri is not None:
|
||||
break
|
||||
else:
|
||||
raise ValueError('No files found with a `download_uri`.')
|
||||
|
||||
content = client.files.download(file=file)
|
||||
assert content[4:8] == b'ftyp'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_download_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
async for file in await client.aio.files.list():
|
||||
if file.download_uri is not None:
|
||||
break
|
||||
else:
|
||||
raise ValueError('No files found with a `download_uri`.')
|
||||
|
||||
content = await client.aio.files.download(file=file)
|
||||
assert content[4:8] == b'ftyp'
|
||||
@@ -0,0 +1,46 @@
|
||||
# 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 pytest_helper
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_get',
|
||||
parameters=types._GetFileParameters(name='files/vjvu9fwk2qj8'),
|
||||
exception_if_vertex='only supported in Gemini Developer API mode',
|
||||
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.get',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
file = await client.aio.files.get(name='files/vjvu9fwk2qj8')
|
||||
@@ -0,0 +1,72 @@
|
||||
# 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 list method."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
test_table: list[pytest_helper.TestTableItem] = [
|
||||
pytest_helper.TestTableItem(
|
||||
name='test_not_empty',
|
||||
exception_if_vertex='only supported in Gemini Developer API mode',
|
||||
parameters=types._ListFilesParameters(
|
||||
config=types.ListFilesConfig(
|
||||
page_size=2,
|
||||
),
|
||||
),
|
||||
),
|
||||
]
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_method='files.list',
|
||||
test_table=test_table,
|
||||
)
|
||||
|
||||
|
||||
def test_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
files = client.files.list(config={'page_size': 2})
|
||||
assert 'content-type' in files.sdk_http_response.headers
|
||||
assert files.name == 'files'
|
||||
assert files.page_size == 2
|
||||
assert len(files) <= 2
|
||||
|
||||
# Iterate through all the pages. Then next_page() should raise an exception.
|
||||
for _ in files:
|
||||
pass
|
||||
with pytest.raises(IndexError, match='No more pages to fetch.'):
|
||||
files.next_page()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pager(client):
|
||||
with pytest_helper.exception_if_vertex(client, ValueError):
|
||||
files = await client.aio.files.list(config={'page_size': 2})
|
||||
|
||||
assert 'Content-Type' in files.sdk_http_response.headers
|
||||
assert files.name == 'files'
|
||||
assert files.page_size == 2
|
||||
assert len(files) <= 2
|
||||
|
||||
# Iterate through all the pages. Then next_page() should raise an exception.
|
||||
async for _ in files:
|
||||
pass
|
||||
with pytest.raises(IndexError, match='No more pages to fetch.'):
|
||||
await files.next_page()
|
||||
@@ -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"
|
||||
@@ -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._InternalRegisterFilesParameters(uris=['gs://unified-genai-dev/image.jpg']),
|
||||
exception_if_vertex='only supported in Gemini Developer API mode',
|
||||
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'
|
||||
@@ -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'},
|
||||
)
|
||||
Reference in New Issue
Block a user