3535import org .junit .jupiter .params .provider .Arguments ;
3636import org .junit .jupiter .params .provider .MethodSource ;
3737
38- import java .util .ArrayList ;
39- import java .util .Collection ;
40- import java .util .HashMap ;
41- import java .util .List ;
42- import java .util .Map ;
43- import java .util .Optional ;
38+ import java .util .*;
4439import java .util .concurrent .CompletableFuture ;
4540import java .util .concurrent .CompletionStage ;
4641import java .util .concurrent .ExecutionException ;
4742import java .util .concurrent .atomic .AtomicBoolean ;
4843import java .util .concurrent .atomic .AtomicInteger ;
44+ import java .util .function .Function ;
4945import java .util .function .Supplier ;
5046import java .util .stream .Collectors ;
5147import java .util .stream .Stream ;
5248
5349import static java .util .Arrays .asList ;
54- import static java .util .Collections .emptyList ;
55- import static java .util .Collections .singletonList ;
50+ import static java .util .Collections .*;
5651import static java .util .concurrent .CompletableFuture .*;
5752import static org .awaitility .Awaitility .await ;
5853import static org .dataloader .DataLoaderFactory .newDataLoader ;
6661import static org .dataloader .fixtures .TestKit .listFrom ;
6762import static org .dataloader .impl .CompletableFutureKit .cause ;
6863import static org .hamcrest .MatcherAssert .assertThat ;
69- import static org .hamcrest .Matchers .empty ;
70- import static org .hamcrest .Matchers .equalTo ;
71- import static org .hamcrest .Matchers .instanceOf ;
72- import static org .hamcrest .Matchers .is ;
64+ import static org .hamcrest .Matchers .*;
7365import static org .junit .jupiter .api .Assertions .assertArrayEquals ;
7466import static org .junit .jupiter .api .Assertions .fail ;
7567
@@ -116,19 +108,25 @@ public void basic_map_batch_loading() {
116108 };
117109 DataLoader <String , String > loader = DataLoaderFactory .newMappedDataLoader (evensOnlyMappedBatchLoader );
118110
111+ final List <String > keys = asList ("C" , "D" );
112+ final Map <String , ?> keysAndContexts = new LinkedHashMap <>();
113+ keysAndContexts .put ("E" , null );
114+ keysAndContexts .put ("F" , null );
115+
119116 loader .load ("A" );
120117 loader .load ("B" );
121- loader .loadMany (asList ("C" , "D" ));
118+ loader .loadMany (keys );
119+ loader .loadMany (keysAndContexts );
122120
123121 List <String > results = loader .dispatchAndJoin ();
124122
125- assertThat (results .size (), equalTo (4 ));
126- assertThat (results , equalTo (asList ("A" , null , "C" , null )));
123+ assertThat (results .size (), equalTo (6 ));
124+ assertThat (results , equalTo (asList ("A" , null , "C" , null , "E" , null )));
127125 }
128126
129127 @ ParameterizedTest
130128 @ MethodSource ("dataLoaderFactories" )
131- public void should_Support_loading_multiple_keys_in_one_call (TestDataLoaderFactory factory ) {
129+ public void should_Support_loading_multiple_keys_in_one_call_via_list (TestDataLoaderFactory factory ) {
132130 AtomicBoolean success = new AtomicBoolean ();
133131 DataLoader <Integer , Integer > identityLoader = factory .idLoader (new DataLoaderOptions (), new ArrayList <>());
134132
@@ -142,6 +140,26 @@ public void should_Support_loading_multiple_keys_in_one_call(TestDataLoaderFacto
142140 assertThat (futureAll .toCompletableFuture ().join (), equalTo (asList (1 , 2 )));
143141 }
144142
143+ @ ParameterizedTest
144+ @ MethodSource ("dataLoaderFactories" )
145+ public void should_Support_loading_multiple_keys_in_one_call_via_map (TestDataLoaderFactory factory ) {
146+ AtomicBoolean success = new AtomicBoolean ();
147+ DataLoader <Integer , Integer > identityLoader = factory .idLoader (new DataLoaderOptions (), new ArrayList <>());
148+
149+ final Map <Integer , ?> keysAndContexts = new LinkedHashMap <>();
150+ keysAndContexts .put (1 , null );
151+ keysAndContexts .put (2 , null );
152+
153+ CompletionStage <Map <Integer , Integer >> futureAll = identityLoader .loadMany (keysAndContexts );
154+ futureAll .thenAccept (promisedValues -> {
155+ assertThat (promisedValues .size (), is (2 ));
156+ success .set (true );
157+ });
158+ identityLoader .dispatch ();
159+ await ().untilAtomic (success , is (true ));
160+ assertThat (futureAll .toCompletableFuture ().join (), equalTo (Map .of (1 , 1 , 2 , 2 )));
161+ }
162+
145163 @ ParameterizedTest
146164 @ MethodSource ("dataLoaderFactories" )
147165 public void should_Resolve_to_empty_list_when_no_keys_supplied (TestDataLoaderFactory factory ) {
@@ -159,7 +177,22 @@ public void should_Resolve_to_empty_list_when_no_keys_supplied(TestDataLoaderFac
159177
160178 @ ParameterizedTest
161179 @ MethodSource ("dataLoaderFactories" )
162- public void should_Return_zero_entries_dispatched_when_no_keys_supplied (TestDataLoaderFactory factory ) {
180+ public void should_Resolve_to_empty_map_when_no_keys_supplied (TestDataLoaderFactory factory ) {
181+ AtomicBoolean success = new AtomicBoolean ();
182+ DataLoader <Integer , Integer > identityLoader = factory .idLoader (new DataLoaderOptions (), new ArrayList <>());
183+ CompletableFuture <Map <Integer , Integer >> futureEmpty = identityLoader .loadMany (emptyMap ());
184+ futureEmpty .thenAccept (promisedValues -> {
185+ assertThat (promisedValues .size (), is (0 ));
186+ success .set (true );
187+ });
188+ identityLoader .dispatch ();
189+ await ().untilAtomic (success , is (true ));
190+ assertThat (futureEmpty .join (), anEmptyMap ());
191+ }
192+
193+ @ ParameterizedTest
194+ @ MethodSource ("dataLoaderFactories" )
195+ public void should_Return_zero_entries_dispatched_when_no_keys_supplied_via_list (TestDataLoaderFactory factory ) {
163196 AtomicBoolean success = new AtomicBoolean ();
164197 DataLoader <Integer , Integer > identityLoader = factory .idLoader (new DataLoaderOptions (), new ArrayList <>());
165198 CompletableFuture <List <Integer >> futureEmpty = identityLoader .loadMany (emptyList ());
@@ -172,6 +205,21 @@ public void should_Return_zero_entries_dispatched_when_no_keys_supplied(TestData
172205 assertThat (dispatchResult .getKeysCount (), equalTo (0 ));
173206 }
174207
208+ @ ParameterizedTest
209+ @ MethodSource ("dataLoaderFactories" )
210+ public void should_Return_zero_entries_dispatched_when_no_keys_supplied_via_map (TestDataLoaderFactory factory ) {
211+ AtomicBoolean success = new AtomicBoolean ();
212+ DataLoader <Integer , Integer > identityLoader = factory .idLoader (new DataLoaderOptions (), new ArrayList <>());
213+ CompletableFuture <Map <Integer , Integer >> futureEmpty = identityLoader .loadMany (emptyMap ());
214+ futureEmpty .thenAccept (promisedValues -> {
215+ assertThat (promisedValues .size (), is (0 ));
216+ success .set (true );
217+ });
218+ DispatchResult <Integer > dispatchResult = identityLoader .dispatchWithCounts ();
219+ await ().untilAtomic (success , is (true ));
220+ assertThat (dispatchResult .getKeysCount (), equalTo (0 ));
221+ }
222+
175223 @ ParameterizedTest
176224 @ MethodSource ("dataLoaderFactories" )
177225 public void should_Batch_multiple_requests (TestDataLoaderFactory factory ) throws ExecutionException , InterruptedException {
@@ -286,10 +334,17 @@ public void should_Cache_on_redispatch(TestDataLoaderFactory factory) throws Exe
286334 CompletableFuture <List <String >> future2 = identityLoader .loadMany (asList ("A" , "B" ));
287335 identityLoader .dispatch ();
288336
289- await ().until (() -> future1 .isDone () && future2 .isDone ());
337+ Map <String , ?> keysAndContexts = new LinkedHashMap <>();
338+ keysAndContexts .put ("A" , null );
339+ keysAndContexts .put ("C" , null );
340+ CompletableFuture <Map <String , String >> future3 = identityLoader .loadMany (keysAndContexts );
341+ identityLoader .dispatch ();
342+
343+ await ().until (() -> future1 .isDone () && future2 .isDone () && future3 .isDone ());
290344 assertThat (future1 .get (), equalTo ("A" ));
291345 assertThat (future2 .get (), equalTo (asList ("A" , "B" )));
292- assertThat (loadCalls , equalTo (asList (singletonList ("A" ), singletonList ("B" ))));
346+ assertThat (future3 .get (), equalTo (keysAndContexts .entrySet ().stream ().collect (Collectors .toMap (Map .Entry ::getKey , Map .Entry ::getKey ))));
347+ assertThat (loadCalls , equalTo (asList (singletonList ("A" ), singletonList ("B" ), singletonList ("C" ))));
293348 }
294349
295350 @ ParameterizedTest
0 commit comments