Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 30 additions & 8 deletions google/genai/_api_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1894,17 +1894,40 @@ def close(self) -> None:
"""Closes the API client."""
# Let users close the custom client explicitly by themselves. Otherwise,
# close the client when the object is garbage collected.
if not self._http_options.httpx_client:
self._httpx_client.close()
# Guard against partial initialization if __init__ failed.

try:
options, client = self._http_options, self._httpx_client
except AttributeError:
return

if not options.httpx_client:
client.close()

async def aclose(self) -> None:
"""Closes the API async client."""
# Let users close the custom client explicitly by themselves. Otherwise,
# close the client when the object is garbage collected.
if not self._http_options.httpx_async_client:
await self._async_httpx_client.aclose()
if self._aiohttp_session and not self._http_options.aiohttp_client:
await self._aiohttp_session.close()
# Guard against partial initialization if __init__ failed.

try:
options = self._http_options
except AttributeError:
return

try:
async_http_client = self._async_httpx_client
except AttributeError:
async_http_client = None
if async_http_client and not options.httpx_async_client:
await async_http_client.aclose()

try:
aio_http_session = self._aiohttp_session
except AttributeError:
aio_http_session = None
if aio_http_session and not options.aiohttp_client:
await aio_http_session.close()

def __del__(self) -> None:
"""Closes the API client when the object is garbage collected.
Expand All @@ -1914,8 +1937,7 @@ def __del__(self) -> None:
"""

try:
if not self._http_options.httpx_client:
self.close()
self.close()
except Exception: # pylint: disable=broad-except
pass

Expand Down
25 changes: 24 additions & 1 deletion google/genai/tests/client/test_client_close.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
"""Tests for closing the clients and context managers."""
import asyncio
from unittest import mock

import gc
from google.oauth2 import credentials
import pytest
try:
Expand Down Expand Up @@ -194,3 +194,26 @@ async def run():
assert async_client._api_client._aiohttp_session.closed

asyncio.run(run())


def test_partial_init_validation_error_sync_safety():
"""Verifies that BaseApiClient.close() does not crash on a partially initialized client."""

with pytest.raises(ValueError):
api_client.BaseApiClient(api_key=None, vertexai=False)

# force GC to trigger __del__
gc.collect()

@pytest.mark.asyncio
async def test_partial_init_validation_error_async_safety():
"""Verifies that BaseApiClient.aclose() does not crash the event loop on a partially initialized client."""

with pytest.raises(ValueError):
api_client.BaseApiClient(api_key=None, vertexai=False)

# force GC to trigger __del__
gc.collect()

# make sure the event loop runs to completion
await asyncio.sleep(0)