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

google/genai/tests/client/test_retries.py

@@ -0,0 +1,846 @@
+# 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 http retries."""
+
+import asyncio
+from collections.abc import Sequence
+import datetime
+from unittest import mock
+import pytest
+try:
+    import aiohttp
+    AIOHTTP_NOT_INSTALLED = False
+except ImportError:
+    AIOHTTP_NOT_INSTALLED = True
+    aiohttp = mock.MagicMock()
+
+from google.oauth2 import credentials
+import httpx
+import tenacity
+
+from ... import _api_client as api_client
+from ... import errors
+from ... import types
+
+
+requires_aiohttp = pytest.mark.skipif(
+    AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
+)
+
+
+_RETRIED_CODES = (
+    408,  # Request timeout.
+    429,  # Too many requests.
+    500,  # Internal server error.
+    502,  # Bad gateway.
+    503,  # Service unavailable.
+    504,  # Gateway timeout.
+)
+
+
+@pytest.fixture(autouse=True)
+def reset_has_aiohttp():
+  yield
+  api_client.has_aiohttp = False
+
+
+def _final_codes(retried_codes: Sequence[int] = _RETRIED_CODES):
+  return [code for code in range(100, 600) if code not in retried_codes]
+
+
+def _httpx_response(code: int):
+  return httpx.Response(
+      status_code=code,
+      headers={'status-code': str(code)},
+      content=b'',
+  )
+
+
+# Args
+
+
+def test_retry_args_disabled():
+  args = api_client.retry_args(None)
+
+  assert set(args.keys()) == {'stop', 'reraise'}
+  assert args['stop'].max_attempt_number == 1
+  assert args['reraise']
+
+
+def test_retry_args_enabled_with_defaults():
+  # Empty options means use the default values whereas None means no retries.
+  args = api_client.retry_args(types.HttpRetryOptions())
+
+  assert set(args.keys()) == {
+      'stop',
+      'retry',
+      'wait',
+      'reraise',
+      'before_sleep',
+  }
+
+  assert args['stop'].max_attempt_number == 5
+
+  wait = args['wait']
+  assert wait.exp_base == 2
+  assert wait.initial == 1
+  assert wait.jitter == 1
+  assert wait.max == 60
+
+  retry = args['retry']
+  for code in _RETRIED_CODES:
+    try:
+      errors.APIError.raise_for_response(_httpx_response(code))
+      assert False, 'Expected APIError to be raised.'
+    except errors.APIError as e:
+      assert retry.predicate(e)
+
+  for code in _final_codes():
+    try:
+      errors.APIError.raise_for_response(_httpx_response(code))
+      # Does not raise for some codes.
+    except errors.APIError as e:
+      # Does not retry for error codes outside of the retried codes list.
+      assert not retry.predicate(e)
+
+    assert args['reraise']
+
+
+def test_retry_wait():
+  timestamps = []
+
+  def fn():
+    now = datetime.datetime.now()
+    timestamps.append(now)
+    raise errors.APIError.raise_for_response(_httpx_response(429))
+
+  retrying = tenacity.Retrying(
+      **api_client.retry_args(types.HttpRetryOptions())
+  )
+
+  try:
+    retrying(fn)
+    assert False, 'Expected APIError to be raised.'
+  except errors.APIError:
+    pass
+
+  assert len(timestamps) == 5
+  assert timestamps[1] - timestamps[0] >= datetime.timedelta(seconds=1)
+  assert timestamps[2] - timestamps[1] >= datetime.timedelta(seconds=2)
+  assert timestamps[3] - timestamps[2] >= datetime.timedelta(seconds=4)
+  assert timestamps[4] - timestamps[3] >= datetime.timedelta(seconds=8)
+
+
+def test_retry_args_enabled_with_custom_values_are_not_overridden():
+  options = types.HttpRetryOptions(
+      attempts=10,
+      initial_delay=10,
+      max_delay=100,
+      exp_base=1.5,
+      jitter=0.5,
+      http_status_codes=[408, 429],
+  )
+  retry_args = api_client.retry_args(options)
+  assert retry_args['stop'].max_attempt_number == 10
+
+  wait = retry_args['wait']
+  assert wait.initial == 10
+  assert wait.max == 100
+  assert wait.exp_base == 1.5
+  assert wait.jitter == 0.5
+
+  retry = retry_args['retry']
+  for code in [408, 429]:
+    try:
+      errors.APIError.raise_for_response(_httpx_response(code))
+      assert False, 'Expected APIError to be raised.'
+    except errors.APIError as e:
+      assert retry.predicate(e)
+
+  for code in _final_codes([408, 429]):
+    try:
+      errors.APIError.raise_for_response(_httpx_response(code))
+      # Does not raise for some codes.
+    except errors.APIError as e:
+      # Does not retry for error codes outside of the retried codes list.
+      assert not retry.predicate(e)
+
+
+def _patch_auth_default():
+  return mock.patch(
+      'google.auth.default',
+      return_value=(credentials.Credentials('magic_token'), 'test_project'),
+      autospec=True,
+  )
+
+
+def _transport_options(http_options=None, transport=None, async_transport=None):
+  http_options = http_options or types.HttpOptions()
+  http_options.client_args = {'transport': transport}
+  http_options.async_client_args = {'transport': async_transport}
+  return http_options
+
+
+# Sync
+
+
+def test_disabled_retries_successful_request_executes_once():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.return_value = _httpx_response(200)
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(transport=mock_transport),
+  )
+
+  with _patch_auth_default():
+    response = client.request(http_method='GET', path='path', request_dict={})
+    mock_transport.handle_request.assert_called_once()
+    assert response.headers['status-code'] == '200'
+
+
+def test_disabled_retries_failed_request_executes_once():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.return_value = _httpx_response(429)
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(transport=mock_transport),
+  )
+
+  with _patch_auth_default():
+    try:
+      client.request(http_method='GET', path='path', request_dict={})
+      assert False, 'Expected APIError to be raised.'
+    except errors.APIError as e:
+      assert e.code == 429
+    mock_transport.handle_request.assert_called_once()
+
+
+_RETRY_OPTIONS = types.HttpRetryOptions(
+    attempts=2,
+    initial_delay=0,
+    max_delay=1,
+    exp_base=0.1,
+    jitter=0.1,
+    http_status_codes=[429, 504],
+)
+
+
+def test_retries_successful_request_executes_once():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.return_value = _httpx_response(200)
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    response = client.request(http_method='GET', path='path', request_dict={})
+    mock_transport.handle_request.assert_called_once()
+    assert response.headers['status-code'] == '200'
+
+
+def test_retries_failed_request_retries_successfully():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.side_effect = (
+      _httpx_response(429),
+      _httpx_response(200),
+  )
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    response = client.request(http_method='GET', path='path', request_dict={})
+    mock_transport.handle_request.assert_called()
+    assert response.headers['status-code'] == '200'
+
+
+def test_retries_failed_request_retries_successfully_at_request_level():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.side_effect = (
+      _httpx_response(429),
+      _httpx_response(200),
+  )
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    response = client.request(
+        http_method='GET',
+        path='path',
+        request_dict={},
+        http_options=types.HttpOptions(
+            retry_options=_RETRY_OPTIONS
+        ),  # At request level.
+    )
+    mock_transport.handle_request.assert_called()
+    assert response.headers['status-code'] == '200'
+
+
+def test_retries_failed_request_retries_unsuccessfully():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.side_effect = (
+      _httpx_response(429),
+      _httpx_response(504),
+  )
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    try:
+      client.request(http_method='GET', path='path', request_dict={})
+      assert False, 'Expected APIError to be raised.'
+    except errors.APIError as e:
+      assert e.code == 504
+    mock_transport.handle_request.assert_called()
+
+
+def test_retries_failed_request_retries_unsuccessfully_at_request_level():
+  mock_transport = mock.Mock(spec=httpx.BaseTransport)
+  mock_transport.handle_request.side_effect = (
+      _httpx_response(429),
+      _httpx_response(504),
+  )
+
+  client = api_client.BaseApiClient(
+      vertexai=True,
+      project='test_project',
+      location='global',
+      http_options=_transport_options(
+          transport=mock_transport,
+      ),
+  )
+
+  with _patch_auth_default():
+    try:
+      client.request(
+          http_method='GET',
+          path='path',
+          request_dict={},
+          http_options={'retry_options': _RETRY_OPTIONS},  # At request level.
+      )
+      assert False, 'Expected APIError to be raised.'
+    except errors.APIError as e:
+      assert e.code == 504
+    mock_transport.handle_request.assert_called()
+
+
+# Async httpx
+
+
+def test_async_disabled_retries_successful_request_executes_once():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.return_value = _httpx_response(200)
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(async_transport=mock_transport),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_transport.handle_async_request.assert_called_once()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+def test_async_disabled_retries_failed_request_executes_once():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.return_value = _httpx_response(429)
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(async_transport=mock_transport),
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET', path='path', request_dict={}
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 429
+      mock_transport.handle_async_request.assert_called_once()
+
+  asyncio.run(run())
+
+
+def test_async_retries_successful_request_executes_once():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.return_value = _httpx_response(200)
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_transport.handle_async_request.assert_called_once()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+def test_async_retries_failed_request_retries_successfully():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.side_effect = (
+        _httpx_response(429),
+        _httpx_response(200),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_transport.handle_async_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+def test_async_retries_failed_request_retries_successfully_at_request_level():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.side_effect = (
+        _httpx_response(429),
+        _httpx_response(200),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET',
+          path='path',
+          request_dict={},
+          http_options=types.HttpOptions(
+              retry_options=_RETRY_OPTIONS
+          ),  # At request level.
+      )
+      mock_transport.handle_async_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+def test_async_retries_failed_request_retries_unsuccessfully():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.side_effect = (
+        _httpx_response(429),
+        _httpx_response(504),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET', path='path', request_dict={}
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 504
+      mock_transport.handle_async_request.assert_called()
+
+  asyncio.run(run())
+
+
+def test_async_retries_failed_request_retries_unsuccessfully_at_request_level():
+  api_client.has_aiohttp = False
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.side_effect = (
+        _httpx_response(429),
+        _httpx_response(504),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            async_transport=mock_transport,
+        ),
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET',
+            path='path',
+            request_dict={},
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 504
+      mock_transport.handle_async_request.assert_called()
+
+  asyncio.run(run())
+
+
+# Async aiohttp
+
+
+async def _aiohttp_async_response(status: int):
+  """Has to return a coroutine hence async."""
+  response = mock.Mock(spec=aiohttp.ClientResponse)
+  response.status = status
+  response.headers = {'status-code': str(status)}
+  response.json.return_value = {}
+  response.text.return_value = 'test'
+  return response
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_disabled_retries_successful_request_executes_once(
+    mock_request,
+):
+  api_client.has_aiohttp = True  # Force aiohttp
+
+  async def run():
+    mock_request.return_value = _aiohttp_async_response(200)
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_request.assert_called_once()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_disabled_retries_failed_request_executes_once(mock_request):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.return_value = _aiohttp_async_response(429)
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET', path='path', request_dict={}
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 429
+      mock_request.assert_called_once()
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_successful_request_executes_once(mock_request):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.return_value = _aiohttp_async_response(200)
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+        ),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_request.assert_called_once()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_failed_request_retries_successfully(mock_request):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        _aiohttp_async_response(429),
+        _aiohttp_async_response(200),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+        ),
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_failed_request_retries_successfully_at_request_level(
+    mock_request,
+):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        _aiohttp_async_response(429),
+        _aiohttp_async_response(200),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET',
+          path='path',
+          request_dict={},
+          http_options=types.HttpOptions(
+              retry_options=_RETRY_OPTIONS
+          ),  # At request level.
+      )
+      mock_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_failed_request_retries_unsuccessfully(mock_request):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        _aiohttp_async_response(429),
+        _aiohttp_async_response(504),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+        ),
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET', path='path', request_dict={}
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 504
+      mock_request.assert_called()
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_failed_request_retries_unsuccessfully_at_request_level(
+    mock_request,
+):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        _aiohttp_async_response(429),
+        _aiohttp_async_response(504),
+    )
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      try:
+        await client.async_request(
+            http_method='GET',
+            path='path',
+            request_dict={},
+            http_options={'retry_options': _RETRY_OPTIONS},  # At request level.
+        )
+        assert False, 'Expected APIError to be raised.'
+      except errors.APIError as e:
+        assert e.code == 504
+      mock_request.assert_called()
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_retries_client_connector_error_retries_successfully(
+    mock_request,
+):
+  api_client.has_aiohttp = True
+
+  async def run():
+    mock_request.side_effect = (
+        aiohttp.ClientConnectorError(
+            connection_key=aiohttp.client_reqrep.ConnectionKey(
+                'localhost', 80, False, True, None, None, None
+            ),
+            os_error=OSError,
+        ),
+        _aiohttp_async_response(200),
+    )
+    # The request will be automatically retried once, if catching the
+    # ClientConnectorError.
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+    )
+
+    with _patch_auth_default():
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+      mock_request.assert_called()
+      assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())