@@ -1134,10 +1134,14 @@ public void testBulkRequestExecutionWithFailures() throws Exception {
1134
1134
Exception error = new RuntimeException ();
1135
1135
doAnswer (args -> {
1136
1136
@ SuppressWarnings ("unchecked" )
1137
- BiConsumer <IngestDocument , Exception > handler = (BiConsumer ) args .getArguments ()[1 ];
1138
- handler .accept (null , error );
1137
+ List <IngestDocumentWrapper > ingestDocumentWrappers = (List ) args .getArguments ()[0 ];
1138
+ Consumer <List <IngestDocumentWrapper >> handler = (Consumer ) args .getArguments ()[1 ];
1139
+ for (IngestDocumentWrapper wrapper : ingestDocumentWrappers ) {
1140
+ wrapper .update (wrapper .getIngestDocument (), error );
1141
+ }
1142
+ handler .accept (ingestDocumentWrappers );
1139
1143
return null ;
1140
- }).when (processor ).execute (any (), any ());
1144
+ }).when (processor ).batchExecute (any (), any ());
1141
1145
IngestService ingestService = createWithProcessors (
1142
1146
Collections .singletonMap ("mock" , (factories , tag , description , config ) -> processor )
1143
1147
);
@@ -1192,10 +1196,11 @@ public void testBulkRequestExecution() throws Exception {
1192
1196
when (processor .getTag ()).thenReturn ("mockTag" );
1193
1197
doAnswer (args -> {
1194
1198
@ SuppressWarnings ("unchecked" )
1195
- BiConsumer <IngestDocument , Exception > handler = (BiConsumer ) args .getArguments ()[1 ];
1196
- handler .accept (RandomDocumentPicks .randomIngestDocument (random ()), null );
1199
+ List <IngestDocumentWrapper > ingestDocumentWrappers = (List ) args .getArguments ()[0 ];
1200
+ Consumer <List <IngestDocumentWrapper >> handler = (Consumer ) args .getArguments ()[1 ];
1201
+ handler .accept (ingestDocumentWrappers );
1197
1202
return null ;
1198
- }).when (processor ).execute (any (), any ());
1203
+ }).when (processor ).batchExecute (any (), any ());
1199
1204
Map <String , Processor .Factory > map = new HashMap <>(2 );
1200
1205
map .put ("mock" , (factories , tag , description , config ) -> processor );
1201
1206
@@ -1957,6 +1962,42 @@ public void testExecuteBulkRequestInBatchWithExceptionAndDropInCallback() {
1957
1962
verify (mockCompoundProcessor , never ()).execute (any (), any ());
1958
1963
}
1959
1964
1965
+ public void testExecuteBulkRequestInBatchWithDefaultBatchSize () {
1966
+ CompoundProcessor mockCompoundProcessor = mockCompoundProcessor ();
1967
+ IngestService ingestService = createWithProcessors (
1968
+ Collections .singletonMap ("mock" , (factories , tag , description , config ) -> mockCompoundProcessor )
1969
+ );
1970
+ createPipeline ("_id" , ingestService );
1971
+ BulkRequest bulkRequest = new BulkRequest ();
1972
+ IndexRequest indexRequest1 = new IndexRequest ("_index" ).id ("_id1" ).source (emptyMap ()).setPipeline ("_id" ).setFinalPipeline ("_none" );
1973
+ bulkRequest .add (indexRequest1 );
1974
+ IndexRequest indexRequest2 = new IndexRequest ("_index" ).id ("_id2" ).source (emptyMap ()).setPipeline ("_id" ).setFinalPipeline ("_none" );
1975
+ bulkRequest .add (indexRequest2 );
1976
+ IndexRequest indexRequest3 = new IndexRequest ("_index" ).id ("_id3" ).source (emptyMap ()).setPipeline ("_none" ).setFinalPipeline ("_id" );
1977
+ bulkRequest .add (indexRequest3 );
1978
+ IndexRequest indexRequest4 = new IndexRequest ("_index" ).id ("_id4" ).source (emptyMap ()).setPipeline ("_id" ).setFinalPipeline ("_none" );
1979
+ bulkRequest .add (indexRequest4 );
1980
+ @ SuppressWarnings ("unchecked" )
1981
+ final Map <Integer , Exception > failureHandler = new HashMap <>();
1982
+ final Map <Thread , Exception > completionHandler = new HashMap <>();
1983
+ final List <Integer > dropHandler = new ArrayList <>();
1984
+ ingestService .executeBulkRequest (
1985
+ 4 ,
1986
+ bulkRequest .requests (),
1987
+ failureHandler ::put ,
1988
+ completionHandler ::put ,
1989
+ dropHandler ::add ,
1990
+ Names .WRITE ,
1991
+ bulkRequest
1992
+ );
1993
+ assertTrue (failureHandler .isEmpty ());
1994
+ assertTrue (dropHandler .isEmpty ());
1995
+ assertEquals (1 , completionHandler .size ());
1996
+ assertNull (completionHandler .get (Thread .currentThread ()));
1997
+ verify (mockCompoundProcessor , times (1 )).batchExecute (any (), any ());
1998
+ verify (mockCompoundProcessor , never ()).execute (any (), any ());
1999
+ }
2000
+
1960
2001
public void testPrepareBatches_same_index_pipeline () {
1961
2002
IngestService .IndexRequestWrapper wrapper1 = createIndexRequestWrapper ("index1" , Collections .singletonList ("p1" ));
1962
2003
IngestService .IndexRequestWrapper wrapper2 = createIndexRequestWrapper ("index1" , Collections .singletonList ("p1" ));
0 commit comments