#208 to ensure tests are in place

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

@@ -45,6 +45,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -1230,6 +1231,31 @@ public void when_values_size_are_more_then_key_size(TestDataLoaderFactory factor
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get")
+    public void should_Support_loading_values_with_context(TestDataLoaderFactory factory) {
+        AtomicReference<BatchLoaderEnvironment> environmentREF = new AtomicReference<>();
+        DataLoader<Integer, Integer> identityLoader = factory.idLoaderWithContext(new DataLoaderOptions(), new ArrayList<>(), environmentREF);
+
+        identityLoader.load(1, "ctx1");
+        identityLoader.load(2, "ctx2");
+        identityLoader.loadMany(List.of(3, 4), List.of("ctx3", "ctx4"));
+
+        CompletableFuture<List<Integer>> cf = identityLoader.dispatch();
+        await().atMost(Duration.FIVE_SECONDS).until(cf::isDone);
+
+        assertThat(cf.toCompletableFuture().join(), equalTo(asList(1, 2, 3, 4)));
+
+        Map<Object, Object> keyContexts = environmentREF.get().getKeyContexts();
+        assertThat(keyContexts, equalTo(Map.of(
+                1, "ctx1",
+                2, "ctx2",
+                3, "ctx3",
+                4, "ctx4"
+        )));
+    }
+
+
     @Test
     public void can_call_a_loader_from_a_loader() throws Exception {
         List<Collection<String>> deepLoadCalls = new ArrayList<>();

src/test/java/org/dataloader/fixtures/parameterized/DelegatingDataLoaderFactory.java

@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.DelegatingDataLoader;
@@ -8,6 +9,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class DelegatingDataLoaderFactory implements TestDataLoaderFactory {
     // its delegates all the way down to the turtles
@@ -38,6 +40,11 @@ public <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<
         return mkDelegateDataLoader(delegateFactory.idLoader(options, loadCalls));
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderWithContext(options, loadCalls, environmentREF));
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
         return mkDelegateDataLoader(delegateFactory.idLoaderDelayed(options, loadCalls, delay));

src/test/java/org/dataloader/fixtures/parameterized/ListDataLoaderFactory.java

@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.fixtures.TestKit;
@@ -9,6 +10,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
@@ -23,6 +25,15 @@ public <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newDataLoader((keys, env) -> {
+            environmentREF.set(env);
+            loadCalls.add(new ArrayList<>(keys));
+            return completedFuture(keys);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
         return newDataLoader(keys -> CompletableFuture.supplyAsync(() -> {

src/test/java/org/dataloader/fixtures/parameterized/MappedDataLoaderFactory.java

@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.fixtures.TestKit;
@@ -11,10 +12,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.dataloader.DataLoaderFactory.newDataLoader;
 import static org.dataloader.DataLoaderFactory.newMappedDataLoader;
 import static org.dataloader.fixtures.TestKit.futureError;
 
@@ -31,6 +32,17 @@ public <K> DataLoader<K, K> idLoader(
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newMappedDataLoader((keys, environment) -> {
+            environmentREF.set(environment);
+            loadCalls.add(new ArrayList<>(keys));
+            Map<K, K> map = new HashMap<>();
+            keys.forEach(k -> map.put(k, k));
+            return completedFuture(map);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(
             DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {

src/test/java/org/dataloader/fixtures/parameterized/MappedPublisherDataLoaderFactory.java

@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.Try;
@@ -12,16 +13,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
-import java.util.stream.Collectors;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toList;
-import static java.util.stream.Collectors.toSet;
 import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoader;
 import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoaderWithTry;
-import static org.dataloader.DataLoaderFactory.newPublisherDataLoader;
 
 public class MappedPublisherDataLoaderFactory implements TestDataLoaderFactory, TestReactiveDataLoaderFactory {
 
@@ -36,6 +34,18 @@ public <K> DataLoader<K, K> idLoader(
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newMappedPublisherDataLoader((keys, subscriber, environment) -> {
+            environmentREF.set(environment);
+
+            loadCalls.add(new ArrayList<>(keys));
+            Map<K, K> map = new HashMap<>();
+            keys.forEach(k -> map.put(k, k));
+            Flux.fromIterable(map.entrySet()).subscribe(subscriber);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(
             DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {

src/test/java/org/dataloader/fixtures/parameterized/PublisherDataLoaderFactory.java

@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.Try;
@@ -11,9 +12,9 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
-import static org.dataloader.DataLoaderFactory.newDataLoader;
 import static org.dataloader.DataLoaderFactory.newPublisherDataLoader;
 import static org.dataloader.DataLoaderFactory.newPublisherDataLoaderWithTry;
 
@@ -28,6 +29,15 @@ public <K> DataLoader<K, K> idLoader(
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newPublisherDataLoader((keys, subscriber, environment) -> {
+            environmentREF.set(environment);
+            loadCalls.add(new ArrayList<>(keys));
+            Flux.fromIterable(keys).subscribe(subscriber);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
         return newPublisherDataLoader((keys, subscriber) -> {

src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java

@@ -1,12 +1,14 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 public interface TestDataLoaderFactory {
     <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<K>> loadCalls);
@@ -23,6 +25,11 @@ public interface TestDataLoaderFactory {
 
     DataLoader<String, String> idLoaderReturnsTooMany(int howManyMore, DataLoaderOptions options, ArrayList<Object> loadCalls);
 
+    // similar to above but batch loaders with context
+
+    <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF);
+
+
     // Convenience methods
 
     default <K> DataLoader<K, K> idLoader(DataLoaderOptions options) {