Copy/tweak original/ DataLoader tests for publisher equivalents

Given the large number of existing tests, we copy across this existing
set for our publisher tests.

What this really indicates is that we should invest in parameterised
testing, but this is a bit painful in JUnit 4 - so we'll bump to JUnit 5
independently and parameterise when we have this available.

This is important because re-using the existing test suite reveals a
failure that we'll need to address.

Author: AlexandreCarlton

src/test/java/org/dataloader/DataLoaderMappedPublisherBatchLoaderTest.java

@@ -3,55 +3,65 @@
 import org.junit.Test;
 import reactor.core.publisher.Flux;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
 import static org.awaitility.Awaitility.await;
 import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoader;
+import static org.dataloader.DataLoaderOptions.newOptions;
+import static org.dataloader.fixtures.TestKit.listFrom;
+import static org.dataloader.impl.CompletableFutureKit.cause;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertThat;
 
 public class DataLoaderMappedPublisherBatchLoaderTest {
 
-    @Test
-    public void should_Build_a_really_really_simple_data_loader() {
-        AtomicBoolean success = new AtomicBoolean();
-        DataLoader<Integer, Integer> identityLoader = newMappedPublisherDataLoader(keysAsValues(), DataLoaderOptions.newOptions());
-
-        CompletionStage<Integer> future1 = identityLoader.load(1);
+    MappedPublisherBatchLoader<String, String> evensOnlyMappedBatchLoader = (keys, subscriber) -> {
+        Map<String, String> mapOfResults = new HashMap<>();
 
-        future1.thenAccept(value -> {
-            assertThat(value, equalTo(1));
-            success.set(true);
+        AtomicInteger index = new AtomicInteger();
+        keys.forEach(k -> {
+            int i = index.getAndIncrement();
+            if (i % 2 == 0) {
+                mapOfResults.put(k, k);
+            }
         });
-        identityLoader.dispatch();
-        await().untilAtomic(success, is(true));
+        Flux.fromIterable(mapOfResults.entrySet()).subscribe(subscriber);
+    };
+
+    private static <K, V> DataLoader<K, V> idMapLoader(DataLoaderOptions options, List<Collection<K>> loadCalls) {
+        MappedPublisherBatchLoader<K, V> kvBatchLoader = (keys, subscriber) -> {
+            loadCalls.add(new ArrayList<>(keys));
+            Map<K, V> map = new HashMap<>();
+            //noinspection unchecked
+            keys.forEach(k -> map.put(k, (V) k));
+            Flux.fromIterable(map.entrySet()).subscribe(subscriber);
+        };
+        return DataLoaderFactory.newMappedPublisherDataLoader(kvBatchLoader, options);
     }
 
-    @Test
-    public void should_Support_loading_multiple_keys_in_one_call() {
-        AtomicBoolean success = new AtomicBoolean();
-        DataLoader<Integer, Integer> identityLoader = newMappedPublisherDataLoader(keysAsValues(), DataLoaderOptions.newOptions());
-
-        CompletionStage<List<Integer>> futureAll = identityLoader.loadMany(asList(1, 2));
-        futureAll.thenAccept(promisedValues -> {
-            assertThat(promisedValues.size(), is(2));
-            success.set(true);
-        });
-        identityLoader.dispatch();
-        await().untilAtomic(success, is(true));
-        assertThat(futureAll.toCompletableFuture().join(), equalTo(asList(1, 2)));
+    private static <K, V> DataLoader<K, V> idMapLoaderBlowsUps(
+        DataLoaderOptions options, List<Collection<K>> loadCalls) {
+        return newMappedPublisherDataLoader((MappedPublisherBatchLoader<K, V>) (keys, subscriber) -> {
+            loadCalls.add(new ArrayList<>(keys));
+            Flux.<Map.Entry<K, V>>error(new IllegalStateException("Error")).subscribe(subscriber);
+        }, options);
     }
 
+
     @Test
-    public void simple_dataloader() {
-        DataLoader<String, String> loader = newMappedPublisherDataLoader(keysAsValues(), DataLoaderOptions.newOptions());
+    public void basic_map_batch_loading() {
+        DataLoader<String, String> loader = DataLoaderFactory.newMappedPublisherDataLoader(evensOnlyMappedBatchLoader);
 
         loader.load("A");
         loader.load("B");
@@ -60,12 +70,13 @@ public void simple_dataloader() {
         List<String> results = loader.dispatchAndJoin();
 
         assertThat(results.size(), equalTo(4));
-        assertThat(results, equalTo(asList("A", "B", "C", "D")));
+        assertThat(results, equalTo(asList("A", null, "C", null)));
     }
 
     @Test
-    public void should_observer_batch_multiple_requests() throws ExecutionException, InterruptedException {
-        DataLoader<Integer, Integer> identityLoader = newMappedPublisherDataLoader(keysAsValues(), new DataLoaderOptions());
+    public void should_map_Batch_multiple_requests() throws ExecutionException, InterruptedException {
+        List<Collection<Integer>> loadCalls = new ArrayList<>();
+        DataLoader<Integer, Integer> identityLoader = idMapLoader(new DataLoaderOptions(), loadCalls);
 
         CompletableFuture<Integer> future1 = identityLoader.load(1);
         CompletableFuture<Integer> future2 = identityLoader.load(2);
@@ -74,11 +85,91 @@ public void should_observer_batch_multiple_requests() throws ExecutionException,
         await().until(() -> future1.isDone() && future2.isDone());
         assertThat(future1.get(), equalTo(1));
         assertThat(future2.get(), equalTo(2));
+        assertThat(loadCalls, equalTo(singletonList(asList(1, 2))));
+    }
+
+    @Test
+    public void can_split_max_batch_sizes_correctly() {
+        List<Collection<Integer>> loadCalls = new ArrayList<>();
+        DataLoader<Integer, Integer> identityLoader = idMapLoader(newOptions().setMaxBatchSize(5), loadCalls);
+
+        for (int i = 0; i < 21; i++) {
+            identityLoader.load(i);
+        }
+        List<Collection<Integer>> expectedCalls = new ArrayList<>();
+        expectedCalls.add(listFrom(0, 5));
+        expectedCalls.add(listFrom(5, 10));
+        expectedCalls.add(listFrom(10, 15));
+        expectedCalls.add(listFrom(15, 20));
+        expectedCalls.add(listFrom(20, 21));
+
+        List<Integer> result = identityLoader.dispatch().join();
+
+        assertThat(result, equalTo(listFrom(0, 21)));
+        assertThat(loadCalls, equalTo(expectedCalls));
+    }
+
+    @Test
+    public void should_Propagate_error_to_all_loads() {
+        List<Collection<Integer>> loadCalls = new ArrayList<>();
+        DataLoader<Integer, Integer> errorLoader = idMapLoaderBlowsUps(new DataLoaderOptions(), loadCalls);
+
+        CompletableFuture<Integer> future1 = errorLoader.load(1);
+        CompletableFuture<Integer> future2 = errorLoader.load(2);
+        errorLoader.dispatch();
+
+        await().until(future1::isDone);
+
+        assertThat(future1.isCompletedExceptionally(), is(true));
+        Throwable cause = cause(future1);
+        assert cause != null;
+        assertThat(cause, instanceOf(IllegalStateException.class));
+        assertThat(cause.getMessage(), equalTo("Error"));
+
+        await().until(future2::isDone);
+        cause = cause(future2);
+        assert cause != null;
+        assertThat(cause.getMessage(), equalTo(cause.getMessage()));
+
+        assertThat(loadCalls, equalTo(singletonList(asList(1, 2))));
+    }
+
+    @Test
+    public void should_work_with_duplicate_keys_when_caching_disabled() throws ExecutionException, InterruptedException {
+        List<Collection<String>> loadCalls = new ArrayList<>();
+        DataLoader<String, String> identityLoader =
+            idMapLoader(newOptions().setCachingEnabled(false), loadCalls);
+
+        CompletableFuture<String> future1 = identityLoader.load("A");
+        CompletableFuture<String> future2 = identityLoader.load("B");
+        CompletableFuture<String> future3 = identityLoader.load("A");
+        identityLoader.dispatch();
+
+        await().until(() -> future1.isDone() && future2.isDone() && future3.isDone());
+        assertThat(future1.get(), equalTo("A"));
+        assertThat(future2.get(), equalTo("B"));
+        assertThat(future3.get(), equalTo("A"));
+
+        // the map batch functions use a set of keys as input and hence remove duplicates unlike list variant
+        assertThat(loadCalls, equalTo(singletonList(asList("A", "B"))));
     }
 
-    private static <K> MappedPublisherBatchLoader<K, K> keysAsValues() {
-        return (keys, subscriber) -> Flux
-            .fromStream(keys.stream().map(k -> Map.entry(k, k)))
-            .subscribe(subscriber);
+    @Test
+    public void should_work_with_duplicate_keys_when_caching_enabled() throws ExecutionException, InterruptedException {
+        List<Collection<String>> loadCalls = new ArrayList<>();
+        DataLoader<String, String> identityLoader =
+            idMapLoader(newOptions().setCachingEnabled(true), loadCalls);
+
+        CompletableFuture<String> future1 = identityLoader.load("A");
+        CompletableFuture<String> future2 = identityLoader.load("B");
+        CompletableFuture<String> future3 = identityLoader.load("A");
+        identityLoader.dispatch();
+
+        await().until(() -> future1.isDone() && future2.isDone() && future3.isDone());
+        assertThat(future1.get(), equalTo("A"));
+        assertThat(future2.get(), equalTo("B"));
+        assertThat(future3.get(), equalTo("A"));
+        assertThat(loadCalls, equalTo(singletonList(asList("A", "B"))));
     }
+
 }