feat: Adding API read/write timeout.

.github/workflows/integration-test.yml

@@ -57,7 +57,7 @@ jobs:
           kubectl get svc -A
           curl $(minikube service chromadb --url)/api/v1
       - name: Test with Maven
-        run: mvn clean test
+        run: mvn --batch-mode clean test
         env:
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}

pom.xml

@@ -115,6 +115,12 @@
             <version>${junit-version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>com.github.tomakehurst</groupId>
+            <artifactId>wiremock-jre8</artifactId>
+            <version>2.35.1</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

src/main/java/tech/amikos/chromadb/Client.java

@@ -16,12 +16,29 @@
  */
 public class Client {
     final ApiClient apiClient = new ApiClient();
-
+    private int timeout = 60;
     DefaultApi api;
 
     public Client(String basePath) {
         apiClient.setBasePath(basePath);
         api = new DefaultApi(apiClient);
+        apiClient.setHttpClient(apiClient.getHttpClient().newBuilder()
+                .readTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
+                .writeTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
+                .build());
+        api.getApiClient().setUserAgent("Chroma-JavaClient/0.1.x");
+    }
+
+    /**
+     * Set the timeout for the client
+     * @param timeout timeout in seconds
+     */
+    public void setTimeout(int timeout) {
+        this.timeout = timeout;
+        apiClient.setHttpClient(apiClient.getHttpClient().newBuilder()
+                .readTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
+                .writeTimeout(this.timeout, java.util.concurrent.TimeUnit.SECONDS)
+                .build());
     }
 
     public Collection getCollection(String collectionName, EmbeddingFunction embeddingFunction) throws ApiException {

src/test/java/TestAPI.java

@@ -1,4 +1,6 @@
+import com.github.tomakehurst.wiremock.junit.WireMockRule;
 import com.google.gson.internal.LinkedTreeMap;
+import org.junit.Rule;
 import org.junit.Test;
 import tech.amikos.chromadb.*;
 import tech.amikos.chromadb.Collection;
@@ -8,11 +10,16 @@
 import java.math.BigDecimal;
 import java.util.*;
 
+import static com.github.tomakehurst.wiremock.client.WireMock.*;
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
 import static org.junit.Assert.*;
 import static org.junit.Assume.*;
 
 public class TestAPI {
 
+    @Rule
+    public WireMockRule wireMockRule = new WireMockRule(8001);
+
 
     @Test
     public void testHeartbeat() throws ApiException, IOException {
@@ -396,4 +403,36 @@ public void testQueryExampleHF() throws ApiException {
     }
 
 
+    @Test
+    public void testTimeoutOk() throws ApiException, IOException {
+        stubFor(get(urlEqualTo("/api/v1/heartbeat"))
+                .willReturn(aResponse()
+                        .withHeader("Content-Type", "application/json")
+                        .withBody("{\"nanosecond heartbeat\": 123456789}").withFixedDelay(2000)));
+
+        Utils.loadEnvFile(".env");
+        Client client = new Client("http://127.0.0.1:8001");
+        client.setTimeout(3);
+        Map<String, BigDecimal> hb = client.heartbeat();
+        assertTrue(hb.containsKey("nanosecond heartbeat"));
+    }
+
+    @Test(expected = ApiException.class)
+    public void testTimeoutExpires() throws ApiException, IOException{
+        stubFor(get(urlEqualTo("/api/v1/heartbeat"))
+                .willReturn(aResponse()
+                        .withHeader("Content-Type", "application/json")
+                        .withBody("{\"nanosecond heartbeat\": 123456789}").withFixedDelay(2000)));
+
+        Utils.loadEnvFile(".env");
+        Client client = new Client("http://127.0.0.1:8001");
+        client.setTimeout(1);
+        try {
+            client.heartbeat();
+        } catch (ApiException e) {
+            assertTrue(e.getMessage().contains("Read timed out") || e.getMessage().contains("timeout"));
+            throw e;
+        }
+
+    }
 }