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