Fix up tests due to chaining of whenComplete

Author: zsmoore

src/main/java/org/dataloader/DataLoaderHelper.java

@@ -189,17 +189,22 @@ CompletableFuture<V> load(K key, Object loadContext) {
 
         ctx.onDispatched();
         loaderOptions.getDispatchStrategy().loadCalled(dataLoader, key, loadContext, loadCallFuture);
-        loadCallFuture = loadCallFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted(dataLoader, result, error));
-        loadCallFuture.whenComplete(ctx::onCompleted);
+        loadCallFuture = loadCallFuture.whenComplete((result, error) -> {
+                    loaderOptions.getDispatchStrategy().loadCompleted(dataLoader, result, error);
+                    ctx.onCompleted(result, error);
+                }
+        );
         return loadCallFuture;
     }
 
     private CompletableFuture<V> incrementCacheHitAndReturnCF(DataLoaderInstrumentationContext<Object> ctx, K key, Object loadContext, CompletableFuture<V> cachedFuture) {
         stats.incrementCacheHitCount(new IncrementCacheHitCountStatisticsContext<>(key, loadContext));
         ctx.onDispatched();
         loaderOptions.getDispatchStrategy().loadCalled(dataLoader, key, loadContext, cachedFuture);
-        cachedFuture = cachedFuture.whenComplete((result, error) -> loaderOptions.getDispatchStrategy().loadCompleted(dataLoader, result, error));
-        cachedFuture.whenComplete(ctx::onCompleted);
+        cachedFuture = cachedFuture.whenComplete((result, error) -> {
+            loaderOptions.getDispatchStrategy().loadCompleted(dataLoader, result, error);
+            ctx.onCompleted(result, error);
+        });
         return cachedFuture;
     }
 

src/test/java/org/dataloader/DataLoaderCacheMapTest.java

@@ -25,18 +25,15 @@ private <T> BatchLoader<T, T> keysAsValues() {
     public void should_provide_all_futures_from_cache() {
         DataLoader<Integer, Integer> dataLoader = newDataLoader(keysAsValues());
 
-        CompletableFuture<Integer> cf1 = dataLoader.load(1);
-        CompletableFuture<Integer> cf2 = dataLoader.load(2);
-        CompletableFuture<Integer> cf3 = dataLoader.load(3);
+        dataLoader.load(1);
+        dataLoader.load(2);
+        dataLoader.load(3);
 
         CacheMap<Object, Integer> cacheMap = dataLoader.getCacheMap();
         Collection<CompletableFuture<Integer>> futures = cacheMap.getAll();
         assertThat(futures.size(), equalTo(3));
 
 
-        assertThat(cacheMap.get(1), equalTo(cf1));
-        assertThat(cacheMap.get(2), equalTo(cf2));
-        assertThat(cacheMap.get(3), equalTo(cf3));
         assertThat(cacheMap.containsKey(1), equalTo(true));
         assertThat(cacheMap.containsKey(2), equalTo(true));
         assertThat(cacheMap.containsKey(3), equalTo(true));
@@ -65,8 +62,7 @@ public void should_access_to_future_dependants() {
         Collection<CompletableFuture<Integer>> futures = dataLoader.getCacheMap().getAll();
 
         List<CompletableFuture<Integer>> futuresList = new ArrayList<>(futures);
-        // instrumentation is depending on the CF completing + dispatch strategy
-        assertThat(futuresList.get(0).getNumberOfDependents(), equalTo(6));
-        assertThat(futuresList.get(1).getNumberOfDependents(), equalTo(3));
+        assertThat(futuresList.get(0).getNumberOfDependents(), equalTo(2)); // instrumentation is depending on the CF completing
+        assertThat(futuresList.get(1).getNumberOfDependents(), equalTo(1));
     }
 }

src/test/java/org/dataloader/DataLoaderIfPresentTest.java

@@ -32,8 +32,6 @@ public void should_detect_if_present_cf() {
         cachedPromise = dataLoader.getIfPresent(1);
         assertThat(cachedPromise.isPresent(), equalTo(true));
 
-        assertThat(cachedPromise.get(), sameInstance(future1));
-
         // but it's not done!
         assertThat(cachedPromise.get().isDone(), equalTo(false));
         //

src/test/java/org/dataloader/DataLoaderTest.java

@@ -282,7 +282,6 @@ public void should_Coalesce_identical_requests(TestDataLoaderFactory factory) th
 
         CompletableFuture<Integer> future1a = identityLoader.load(1);
         CompletableFuture<Integer> future1b = identityLoader.load(1);
-        assertThat(future1a, equalTo(future1b));
         identityLoader.dispatch();
 
         await().until(future1a::isDone);

src/test/java/org/dataloader/DataLoaderValueCacheTest.java

@@ -13,6 +13,8 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -340,10 +342,17 @@ public CompletableFuture<List<Try<Object>>> getValues(List<String> keys) {
     }
 
     private boolean isAssertionException(CompletableFuture<String> fA) {
-        Throwable throwable = Try.tryFuture(fA).join().getThrowable();
+        Throwable throwable = maybeUnwrapCompletionException(Try.tryFuture(fA).join().getThrowable());
         return throwable instanceof DataLoaderAssertionException;
     }
 
+    private Throwable maybeUnwrapCompletionException(Throwable throwable) {
+        if (throwable instanceof CompletionException) {
+            return maybeUnwrapCompletionException(throwable.getCause());
+        }
+        return throwable;
+    }
+
 
     @ParameterizedTest
     @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get")

src/test/java/org/dataloader/strategy/GreedyLevelByLevelChainedDispatchStrategyStressTest.java

@@ -4,6 +4,7 @@
 import org.dataloader.DataLoaderFactory;
 import org.dataloader.DataLoaderRegistry;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
 
 import java.time.Duration;
@@ -266,7 +267,7 @@ void verifyExecutionOrder() throws Exception {
                 List.of("E", "F")
         )));
 
-        assertThat(completionOrder, equalTo(List.of(
+        assertThat(flatCompletionOrder, equalTo(List.of(
                 "H",
                 "J",
                 "I",