Fix batchCheck error handling (#183)

* Fix batchCheck latch handling on errors

* Apply Spotless formatting

Author: aaguiarz

src/main/java/dev/openfga/sdk/api/client/OpenFgaClient.java

@@ -25,6 +25,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -715,6 +716,7 @@ public CompletableFuture<ClientBatchCheckResponse> batchCheck(
         var latch = new CountDownLatch(batchedChecks.size());
 
         var responses = new ConcurrentLinkedQueue<ClientBatchCheckSingleResponse>();
+        var failure = new AtomicReference<Throwable>();
 
         var override = new ConfigurationOverride().addHeaders(options);
 
@@ -735,26 +737,36 @@ public CompletableFuture<ClientBatchCheckResponse> batchCheck(
 
                     return api.batchCheck(configuration.getStoreId(), body, override);
                 })
-                .handleAsync((batchCheckResponseApiResponse, throwable) -> {
-                    Map<String, BatchCheckSingleResult> response =
-                            batchCheckResponseApiResponse.getData().getResult();
-
-                    List<ClientBatchCheckSingleResponse> batchResults = new ArrayList<>();
-                    response.forEach((key, result) -> {
-                        boolean allowed = Boolean.TRUE.equals(result.getAllowed());
-                        ClientBatchCheckItem checkItem = correlationIdToCheck.get(key);
-                        var singleResponse =
-                                new ClientBatchCheckSingleResponse(allowed, checkItem, key, result.getError());
-                        batchResults.add(singleResponse);
-                    });
-                    return batchResults;
-                })
-                .thenAccept(responses::addAll)
-                .thenRun(latch::countDown);
+                .whenComplete((batchCheckResponseApiResponse, throwable) -> {
+                    try {
+                        if (throwable != null) {
+                            failure.compareAndSet(null, throwable);
+                            return;
+                        }
+
+                        Map<String, BatchCheckSingleResult> response =
+                                batchCheckResponseApiResponse.getData().getResult();
+
+                        List<ClientBatchCheckSingleResponse> batchResults = new ArrayList<>();
+                        response.forEach((key, result) -> {
+                            boolean allowed = Boolean.TRUE.equals(result.getAllowed());
+                            ClientBatchCheckItem checkItem = correlationIdToCheck.get(key);
+                            var singleResponse =
+                                    new ClientBatchCheckSingleResponse(allowed, checkItem, key, result.getError());
+                            batchResults.add(singleResponse);
+                        });
+                        responses.addAll(batchResults);
+                    } finally {
+                        latch.countDown();
+                    }
+                });
 
         try {
             batchedChecks.forEach(batch -> executor.execute(() -> singleBatchCheckRequest.accept(batch)));
             latch.await();
+            if (failure.get() != null) {
+                return CompletableFuture.failedFuture(failure.get());
+            }
             return CompletableFuture.completedFuture(new ClientBatchCheckResponse(new ArrayList<>(responses)));
         } catch (Exception e) {
             return CompletableFuture.failedFuture(e);

src/test/java/dev/openfga/sdk/api/client/OpenFgaClientTest.java

@@ -2072,6 +2072,30 @@ public void batchCheck_withOptions() throws Exception {
         assertTrue(response.getResult().isEmpty());
     }
 
+    @Test
+    public void batchCheck_rateLimited() {
+        // Given
+        String postUrl = String.format("https://api.fga.example/stores/%s/batch-check", DEFAULT_STORE_ID);
+        mockHttpClient.onPost(postUrl).doReturn(429, "{\"code\":\"rate_limited\",\"message\":\"Too Many Requests\"}");
+
+        ClientBatchCheckItem item = new ClientBatchCheckItem()
+                .user(DEFAULT_USER)
+                .relation(DEFAULT_RELATION)
+                ._object(DEFAULT_OBJECT)
+                .correlationId("cor-1");
+        ClientBatchCheckRequest request = new ClientBatchCheckRequest().checks(List.of(item));
+
+        // When
+        ExecutionException execException = assertThrows(
+                ExecutionException.class, () -> fga.batchCheck(request).get());
+
+        // Then
+        mockHttpClient.verify().post(postUrl).called(1 + DEFAULT_MAX_RETRIES);
+        var exception = assertInstanceOf(FgaApiRateLimitExceededError.class, execException.getCause());
+        assertEquals(429, exception.getStatusCode());
+        assertEquals("{\"code\":\"rate_limited\",\"message\":\"Too Many Requests\"}", exception.getResponseData());
+    }
+
     /**
      * Expand all relationships in userset tree format, and following userset rewrite rules.  Useful to reason
      * about and debug a certain relationship.