first vommit
This commit is contained in:
commit
a04e68bc95
4926 files changed
+764440
No files matched your search
Whitespace-only changes.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,486 @@
|
||||
# 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 Interactions API."""
|
||||
|
||||
from ... import Client
|
||||
from ... import _base_url
|
||||
from unittest import mock
|
||||
import pytest
|
||||
from httpx import Request, Response
|
||||
from ..._api_client import AsyncHttpxClient, BaseApiClient
|
||||
from httpx import Client as HTTPClient
|
||||
import os
|
||||
|
||||
ENV_VARS = [
|
||||
"GOOGLE_CLOUD_PROJECT",
|
||||
"GEMINI_API_KEY",
|
||||
"GOOGLE_API_KEY",
|
||||
"GOOGLE_CLOUD_LOCATION",
|
||||
]
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_env_vars(monkeypatch):
|
||||
for var in ENV_VARS:
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
|
||||
def test_interactions_gemini_url(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url).endswith('/v1beta/interactions')
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
|
||||
def test_interactions_gemini_no_vertex_auth(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_access_token") as mock_access_token,
|
||||
mock.patch.object(HTTPClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_access_token.assert_not_called()
|
||||
|
||||
def test_interactions_gemini_retry(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
client._api_client.max_retries = 2
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.side_effect = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
]
|
||||
client.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
assert mock_send.call_count == 3
|
||||
|
||||
def test_interactions_gemini_extra_headers(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'X-Custom-Header': 'TestValue'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert request.headers['x-custom-header'] == 'TestValue'
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
|
||||
def test_interactions_vertex_auth_header():
|
||||
from ..._api_client import BaseApiClient
|
||||
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=Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
) 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=Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
) 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())
|
||||
|
||||
def test_interactions_vertex_url():
|
||||
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-token") as mock_access_token,
|
||||
mock.patch("httpx.Client.send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == 'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/interactions'
|
||||
|
||||
def test_interactions_vertex_auth_refresh_on_retry():
|
||||
from ..._api_client import BaseApiClient
|
||||
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)
|
||||
client._api_client.max_retries = 2
|
||||
|
||||
token_values = ['token1', 'token2', 'token3']
|
||||
token_iter = iter(token_values)
|
||||
|
||||
def get_token():
|
||||
return next(token_iter)
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_access_token", side_effect=get_token) as mock_access_token,
|
||||
mock.patch.object(HTTPClient, "send") as mock_send,
|
||||
):
|
||||
responses = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
]
|
||||
response_iter = iter(responses)
|
||||
recorded_headers = []
|
||||
def my_send(req, *args, **kwargs):
|
||||
recorded_headers.append(dict(req.headers))
|
||||
return next(response_iter)
|
||||
mock_send.side_effect = my_send
|
||||
|
||||
client.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
|
||||
assert mock_access_token.call_count == 3
|
||||
assert mock_send.call_count == 3
|
||||
# Check headers of each call
|
||||
for i in range(3):
|
||||
assert recorded_headers[i]['authorization'] == f'Bearer {token_values[i]}'
|
||||
|
||||
|
||||
def test_interactions_vertex_extra_headers_override():
|
||||
from ..._api_client import BaseApiClient
|
||||
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='default-token') as mock_access_token,
|
||||
mock.patch.object(HTTPClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
|
||||
# Override Authorization
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'Authorization': 'Bearer manual-token'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['authorization'] == 'Bearer manual-token'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_access_token.reset_mock()
|
||||
|
||||
# Provide API Key
|
||||
client.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'x-goog-api-key': 'manual-key'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['x-goog-api-key'] == 'manual-key'
|
||||
assert 'authorization' not in headers
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_url(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url).endswith('/v1beta/interactions')
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_no_vertex_auth(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_async_access_token") as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_access_token.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_retry(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
client.aio._api_client.max_retries = 2
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.side_effect = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
]
|
||||
await client.aio.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
assert mock_send.call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_gemini_extra_headers(monkeypatch):
|
||||
monkeypatch.setenv('GOOGLE_API_KEY', 'test-api-key')
|
||||
client = Client()
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'X-Custom-Header': 'TestValue'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert request.headers['x-custom-header'] == 'TestValue'
|
||||
assert request.headers['x-goog-api-key'] == 'test-api-key'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_auth_header():
|
||||
from ..._api_client import BaseApiClient
|
||||
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, "_access_token", return_value='fake-vertex-token'
|
||||
) as mock_access_token,
|
||||
mock.patch.object(
|
||||
AsyncHttpxClient, "send",
|
||||
return_value=Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
) 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
|
||||
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=Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
) 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())
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_url():
|
||||
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, "_access_token", return_value="fake-token") as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == 'https://us-central1-aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/interactions'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_auth_refresh_on_retry():
|
||||
from ..._api_client import BaseApiClient
|
||||
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)
|
||||
client.aio._api_client.max_retries = 2
|
||||
|
||||
token_values = ['token1', 'token2', 'token3']
|
||||
token_iter = iter(token_values)
|
||||
|
||||
def get_token():
|
||||
return next(token_iter)
|
||||
|
||||
with (
|
||||
mock.patch.object(BaseApiClient, "_access_token", side_effect=get_token) as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
responses = [
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(500, request=Request('POST', ''), headers={"retry-after-ms": "1", "content-type": "application/json"}, content='{"status": "completed"}'),
|
||||
Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}'),
|
||||
]
|
||||
response_iter = iter(responses)
|
||||
recorded_headers = []
|
||||
async def my_send(req, *args, **kwargs):
|
||||
recorded_headers.append(dict(req.headers))
|
||||
return next(response_iter)
|
||||
mock_send.side_effect = my_send
|
||||
|
||||
await client.aio.interactions.create(model='gemini-1.5-flash', input='Hello')
|
||||
|
||||
assert mock_access_token.call_count == 3
|
||||
assert mock_send.call_count == 3
|
||||
for i in range(3):
|
||||
assert recorded_headers[i]['authorization'] == f'Bearer {token_values[i]}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_interactions_vertex_extra_headers_override():
|
||||
from ..._api_client import BaseApiClient
|
||||
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, "_access_token", return_value='default-token') as mock_access_token,
|
||||
mock.patch.object(AsyncHttpxClient, "send") as mock_send,
|
||||
):
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
|
||||
# Override Authorization
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'Authorization': 'Bearer manual-token'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['authorization'] == 'Bearer manual-token'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_access_token.reset_mock()
|
||||
|
||||
# Provide API Key
|
||||
await client.aio.interactions.create(
|
||||
model='gemini-1.5-flash',
|
||||
input='Hello',
|
||||
extra_headers={'x-goog-api-key': 'manual-key'}
|
||||
)
|
||||
mock_send.assert_called_once()
|
||||
headers = mock_send.call_args[0][0].headers
|
||||
assert headers['x-goog-api-key'] == 'manual-key'
|
||||
assert 'authorization' not in headers
|
||||
+99
@@ -0,0 +1,99 @@
|
||||
# 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.
|
||||
|
||||
from unittest import mock
|
||||
|
||||
from ..._gaos import google_genai as gaos_google_genai
|
||||
import pytest
|
||||
|
||||
from ... import client as client_lib
|
||||
|
||||
|
||||
pytest_plugins = ("pytest_asyncio",)
|
||||
|
||||
|
||||
|
||||
def test_client_timeout():
|
||||
with mock.patch.object(
|
||||
gaos_google_genai, "GenAI", spec_set=True
|
||||
) as mock_genai:
|
||||
mock_client = mock.Mock()
|
||||
mock_genai.return_value = mock_client
|
||||
|
||||
client = client_lib.Client(
|
||||
api_key="placeholder",
|
||||
http_options={"api_version": "v1alpha", "timeout": 5000},
|
||||
)
|
||||
|
||||
# Trigger client build
|
||||
_ = client.interactions.with_raw_response
|
||||
|
||||
mock_genai.assert_called_once_with(
|
||||
security=mock.ANY,
|
||||
api_version=mock.ANY,
|
||||
user_project=mock.ANY,
|
||||
server_url=mock.ANY,
|
||||
client=mock.ANY,
|
||||
timeout_ms=5000,
|
||||
retry_config=mock.ANY,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_client_timeout():
|
||||
with mock.patch.object(
|
||||
gaos_google_genai, "AsyncGenAI", spec_set=True
|
||||
) as mock_async_genai:
|
||||
mock_client = mock.Mock()
|
||||
mock_async_genai.return_value = mock_client
|
||||
|
||||
client = client_lib.Client(
|
||||
api_key="placeholder",
|
||||
http_options={"api_version": "v1alpha", "timeout": 5000},
|
||||
)
|
||||
|
||||
# Trigger client build
|
||||
_ = client.aio.interactions.with_raw_response
|
||||
|
||||
mock_async_genai.assert_called_once_with(
|
||||
security=mock.ANY,
|
||||
api_version=mock.ANY,
|
||||
user_project=mock.ANY,
|
||||
server_url=mock.ANY,
|
||||
async_client=mock.ANY,
|
||||
timeout_ms=5000,
|
||||
retry_config=mock.ANY,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("error")
|
||||
def test_unrecognized_model_serialization():
|
||||
from ..._gaos.types.interactions.createmodelinteraction import CreateModelInteraction
|
||||
# This shouldn't raise a Pydantic serialization error due to UnrecognizedStr
|
||||
obj = CreateModelInteraction(model="gemini-3.5-flash", input="hello")
|
||||
dumped = obj.model_dump()
|
||||
assert dumped["model"] == "gemini-3.5-flash"
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("error")
|
||||
def test_unrecognized_model_request_serialization():
|
||||
from ..._gaos.models.createinteraction import CreateInteractionRequest
|
||||
from ..._gaos.types.interactions.createmodelinteraction import CreateModelInteraction
|
||||
body = CreateModelInteraction(model="gemini-3.5-flash", input="hello")
|
||||
req = CreateInteractionRequest(body=body)
|
||||
dumped = req.model_dump()
|
||||
assert dumped["body"]["model"] == "gemini-3.5-flash"
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
|
||||
# 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 Interactions API URL paths."""
|
||||
|
||||
from unittest import mock
|
||||
import pytest
|
||||
from httpx import Request, Response
|
||||
from ..._api_client import AsyncHttpxClient
|
||||
from httpx import Client as HTTPClient
|
||||
from .. import pytest_helper
|
||||
import google.auth
|
||||
|
||||
@mock.patch.object(google.auth, "default", autospec=True)
|
||||
def test_interactions_paths(mock_auth_default, client):
|
||||
interaction_id = "test-interaction-id"
|
||||
|
||||
mock_creds = mock.Mock()
|
||||
mock_creds.token = "test-token"
|
||||
mock_creds.expired = False
|
||||
mock_creds.quota_project_id = "test-quota-project"
|
||||
mock_auth_default.return_value = (mock_creds, "test-project")
|
||||
|
||||
if client._api_client.vertexai:
|
||||
expected_base_url = f'https://{client._api_client.location}-aiplatform.googleapis.com/v1beta1/projects/{client._api_client.project}/locations/{client._api_client.location}'
|
||||
else:
|
||||
expected_base_url = "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
with mock.patch.object(HTTPClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('GET', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
client.interactions.get(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}?stream=false'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
client.interactions.cancel(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}/cancel'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('DELETE', ''))
|
||||
client.interactions.delete(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@mock.patch.object(google.auth, "default", autospec=True)
|
||||
async def test_async_interactions_paths(mock_auth_default, client):
|
||||
interaction_id = "test-interaction-id"
|
||||
|
||||
mock_creds = mock.Mock()
|
||||
mock_creds.token = "test-token"
|
||||
mock_creds.expired = False
|
||||
mock_creds.quota_project_id = "test-quota-project"
|
||||
mock_auth_default.return_value = (mock_creds, "test-project")
|
||||
|
||||
if client._api_client.vertexai:
|
||||
expected_base_url = f'https://{client._api_client.location}-aiplatform.googleapis.com/v1beta1/projects/{client._api_client.project}/locations/{client._api_client.location}'
|
||||
else:
|
||||
expected_base_url = "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
with mock.patch.object(AsyncHttpxClient, "send") as mock_send:
|
||||
mock_send.return_value = Response(200, request=Request('GET', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
await client.aio.interactions.get(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}?stream=false'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('POST', ''), headers={'content-type': 'application/json'}, content='{"status": "completed"}')
|
||||
await client.aio.interactions.cancel(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}/cancel'
|
||||
|
||||
mock_send.reset_mock()
|
||||
mock_send.return_value = Response(200, request=Request('DELETE', ''))
|
||||
await client.aio.interactions.delete(id=interaction_id)
|
||||
mock_send.assert_called_once()
|
||||
request = mock_send.call_args[0][0]
|
||||
assert str(request.url) == f'{expected_base_url}/interactions/{interaction_id}'
|
||||
|
||||
pytestmark = pytest_helper.setup(
|
||||
file=__file__,
|
||||
globals_for_file=globals(),
|
||||
test_table=[],
|
||||
)
|
||||
Reference in new issue
Block a user