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