#208 to ensure tests are in place by bbakerman · Pull Request #211 · graphql-java/java-dataloader

26 changes: 26 additions & 0 deletions src/test/java/org/dataloader/DataLoaderTest.java

Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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<>();
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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));
Expand Down

11 changes: 11 additions & 0 deletions src/test/java/org/dataloader/fixtures/parameterized/ListDataLoaderFactory.java

Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand All @@ -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(() -> {
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;

Expand All @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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 {

Expand All @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;

Expand All @@ -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) -> {
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
Expand All @@ -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) {
Expand Down