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

google/genai/tests/client/test_client_close.py

@@ -0,0 +1,197 @@
+# 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 closing the clients and context managers."""
+import asyncio
+from unittest import mock
+
+from google.oauth2 import credentials
+import pytest
+try:
+  import aiohttp
+  AIOHTTP_NOT_INSTALLED = False
+except ImportError:
+  AIOHTTP_NOT_INSTALLED = True
+  aiohttp = mock.MagicMock()
+
+
+from ... import _api_client as api_client
+from ... import Client
+
+
+requires_aiohttp = pytest.mark.skipif(
+    AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.'
+)
+
+
+def test_close_httpx_client():
+  """Tests that the httpx client is closed when the client is closed."""
+  api_client.has_aiohttp = False
+  client = Client(
+      vertexai=True,
+      project='test_project',
+      location='global',
+  )
+  client.close()
+  assert client._api_client._httpx_client.is_closed
+
+
+def test_httpx_client_context_manager():
+  """Tests that the httpx client is closed when the client is closed."""
+  api_client.has_aiohttp = False
+  with Client(
+      vertexai=True,
+      project='test_project',
+      location='global',
+  ) as client:
+    pass
+    assert not client._api_client._httpx_client.is_closed
+
+  assert client._api_client._httpx_client.is_closed
+
+
+@pytest.mark.asyncio
+async def test_aclose_httpx_client():
+  """Tests that the httpx async client is closed when the client is closed."""
+  api_client.has_aiohttp = False
+  async_client = Client(
+      vertexai=True,
+      project='test_project',
+      location='global',
+  ).aio
+  await async_client.aclose()
+  assert async_client._api_client._async_httpx_client.is_closed
+
+
+@pytest.mark.asyncio
+async def test_async_httpx_client_context_manager():
+  """Tests that the httpx async client is closed when the client is closed."""
+  api_client.has_aiohttp = False
+  async with Client(
+      vertexai=True,
+      project='test_project',
+      location='global',
+  ).aio as async_client:
+    pass
+    assert not async_client._api_client._async_httpx_client.is_closed
+
+  assert async_client._api_client._async_httpx_client.is_closed
+
+
+@requires_aiohttp
+@pytest.fixture
+def mock_request():
+  mock_aiohttp_response = mock.Mock(spec=aiohttp.ClientSession.request)
+  mock_aiohttp_response.return_value = mock_aiohttp_response
+  yield mock_aiohttp_response
+
+
+def _patch_auth_default():
+  return mock.patch(
+      'google.auth.default',
+      return_value=(credentials.Credentials('magic_token'), 'test_project'),
+      autospec=True,
+  )
+
+
+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_aclose_aiohttp_session(mock_request):
+  """Tests that the aiohttp session is closed when the client is closed."""
+  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),
+    )
+    with _patch_auth_default():
+      async_client = Client(
+          vertexai=True,
+          project='test_project',
+          location='global',
+      ).aio
+      # aiohttp session is created in the first request instead of client
+      # initialization.
+      _ = await async_client._api_client._async_request_once(
+          api_client.HttpRequest(
+              method='GET',
+              url='https://example.com',
+              headers={},
+              data=None,
+              timeout=None,
+          )
+      )
+      assert async_client._api_client._aiohttp_session is not None
+      assert not async_client._api_client._aiohttp_session.closed
+      # Close the client and check that the session is closed.
+      await async_client.aclose()
+      assert async_client._api_client._aiohttp_session.closed
+
+  asyncio.run(run())
+
+
+@requires_aiohttp
+@mock.patch.object(aiohttp.ClientSession, 'request', autospec=True)
+def test_aiohttp_session_context_manager(mock_request):
+  """Tests that the aiohttp session is closed when the client is closed."""
+  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),
+    )
+    with _patch_auth_default():
+      async with Client(
+          vertexai=True,
+          project='test_project',
+          location='global',
+      ).aio as async_client:
+        # aiohttp session is created in the first request instead of client
+        # initialization.
+        _ = await async_client._api_client._async_request_once(
+            api_client.HttpRequest(
+                method='GET',
+                url='https://example.com',
+                headers={},
+                data=None,
+                timeout=None,
+            )
+        )
+        assert async_client._api_client._aiohttp_session is not None
+        assert not async_client._api_client._aiohttp_session.closed
+
+      assert async_client._api_client._aiohttp_session.closed
+
+  asyncio.run(run())