TermsAggregatorTests flaky test fix

Signed-off-by: Sandesh Kumar <[email protected]>
sandeshkr419 · May 6, 2024 · 6b95eeb · 6b95eeb
1 parent bd00518
commit 6b95eeb
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/...r/src/test/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorTests.java b/...r/src/test/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorTests.java
@@ -353,33 +353,40 @@ private void testSimple(
                     newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE)
                 )
             ) {
+                List<Document> documents = new ArrayList<>();
                 Document document = new Document();
                 addFieldConsumer.apply(document, "string", "a");
                 addFieldConsumer.apply(document, "string", "b");
-                indexWriter.addDocument(document);
+                documents.add(document);
+
                 document = new Document();
                 addFieldConsumer.apply(document, "string", "");
                 addFieldConsumer.apply(document, "string", "c");
                 addFieldConsumer.apply(document, "string", "a");
-                indexWriter.addDocument(document);
+                documents.add(document);
+
                 document = new Document();
                 addFieldConsumer.apply(document, "string", "b");
                 addFieldConsumer.apply(document, "string", "d");
-                indexWriter.addDocument(document);
+                documents.add(document);
+
                 document = new Document();
                 addFieldConsumer.apply(document, "string", "");
                 if (includeDocCountField) {
                     // Adding _doc_count to one document
                     document.add(new NumericDocValuesField("_doc_count", 10));
                 }
-                indexWriter.addDocument(document);
+                documents.add(document);
 
                 if (includeDeletedDocumentsInSegment) {
                     document = new Document();
                     ADD_SORTED_SET_FIELD_INDEXED.apply(document, "string", "e");
-                    indexWriter.addDocument(document);
+                    documents.add(document);
+                    indexWriter.addDocuments(documents);
                     indexWriter.deleteDocuments(new Term("string", "e"));
                     assertEquals(5, indexWriter.getDocStats().maxDoc);  // deleted document still in segment
+                } else {
+                    indexWriter.addDocuments(documents);
                 }
 
                 try (IndexReader indexReader = maybeWrapReaderEs(indexWriter.getReader())) {