diff --git a/pom.xml b/pom.xml index 67d1776e3c5..84f0204e039 100644 --- a/pom.xml +++ b/pom.xml @@ -158,7 +158,6 @@ jspecify 1.0.0 - org.jetbrains.kotlin kotlin-stdlib-jdk8 diff --git a/src/main/java/com/google/genai/ApiClient.java b/src/main/java/com/google/genai/ApiClient.java index 888180b70fa..3e2b11fc0e7 100644 --- a/src/main/java/com/google/genai/ApiClient.java +++ b/src/main/java/com/google/genai/ApiClient.java @@ -52,7 +52,7 @@ /** Interface for an API client which issues HTTP requests to the GenAI APIs. */ @InternalApi -public abstract class ApiClient { +public abstract class ApiClient implements AutoCloseable { // {x-version-update-start:google-genai:released} private static final String SDK_VERSION = "1.40.0"; @@ -726,4 +726,17 @@ public static void setDefaultBaseUrls( ApiClient.geminiBaseUrl = geminiBaseUrl; ApiClient.vertexBaseUrl = vertexBaseUrl; } + + @Override + public void close() { + try { + httpClient().dispatcher().executorService().shutdown(); + httpClient().connectionPool().evictAll(); + if (httpClient().cache() != null) { + httpClient().cache().close(); + } + } catch (IOException e) { + throw new GenAiIOException("Failed to close the client.", e); + } + } } diff --git a/src/main/java/com/google/genai/Client.java b/src/main/java/com/google/genai/Client.java index 0bc250b550b..e82f7270f88 100644 --- a/src/main/java/com/google/genai/Client.java +++ b/src/main/java/com/google/genai/Client.java @@ -330,15 +330,7 @@ Optional baseUrl() { /** Closes the Client instance together with its instantiated http client. */ @Override public void close() { - try { - apiClient.httpClient().dispatcher().executorService().shutdown(); - apiClient.httpClient().connectionPool().evictAll(); - if (apiClient.httpClient().cache() != null) { - apiClient.httpClient().cache().close(); - } - } catch (IOException e) { - throw new GenAiIOException("Failed to close the client.", e); - } + apiClient.close(); } /** diff --git a/src/test/java/com/google/genai/ClientTest.java b/src/test/java/com/google/genai/ClientTest.java index ee48cb33793..b84e8026314 100644 --- a/src/test/java/com/google/genai/ClientTest.java +++ b/src/test/java/com/google/genai/ClientTest.java @@ -188,24 +188,4 @@ public void testSetDefaultBaseUrls() { // Reset the base URLs after the test. Client.setDefaultBaseUrls(Optional.empty(), Optional.empty()); } - - @Test - public void testCloseClient() throws Exception { - // Arrange - ApiClient apiClient = mock(ApiClient.class); - OkHttpClient httpClient = new OkHttpClient(); - when(apiClient.httpClient()).thenReturn(httpClient); - - Client client = Client.builder().apiKey(API_KEY).vertexAI(false).build(); - Field apiClientField = Client.class.getDeclaredField("apiClient"); - apiClientField.setAccessible(true); - apiClientField.set(client, apiClient); - - // Act - client.close(); - - // Assert - assertTrue(httpClient.dispatcher().executorService().isShutdown()); - assertEquals(0, httpClient.connectionPool().connectionCount()); - } } diff --git a/src/test/java/com/google/genai/HttpApiClientTest.java b/src/test/java/com/google/genai/HttpApiClientTest.java index 68a7020e00f..63616accb9a 100644 --- a/src/test/java/com/google/genai/HttpApiClientTest.java +++ b/src/test/java/com/google/genai/HttpApiClientTest.java @@ -23,9 +23,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; +import static org.mockito.Mockito.*; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; @@ -46,6 +44,7 @@ import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; + import okhttp3.Call; import okhttp3.Callback; import okhttp3.Dispatcher; @@ -1483,4 +1482,21 @@ public void testNoDefaultLocationWhenUsingApiKeyOnlyMode( assertNull(client.location()); assertTrue(client.vertexAI()); } + + @Test + public void testCloseClient() { + HttpApiClient client = + new HttpApiClient( + Optional.empty(), + Optional.of(PROJECT), + Optional.of(LOCATION), + Optional.of(CREDENTIALS), + Optional.empty(), + Optional.empty()); + + client.close(); + + assertTrue(client.httpClient().dispatcher().executorService().isShutdown()); + assertEquals(0, client.httpClient().connectionPool().connectionCount()); + } }