From a358230d2f3fe612083b467298e8f4dc748ab2b1 Mon Sep 17 00:00:00 2001
From: Sandesh Kumar <sandeshkr419@gmail.com>
Date: Mon, 6 May 2024 10:38:38 -0700
Subject: [PATCH] TermsAggregatorTests flaky test fix

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
---
 .../bucket/terms/TermsAggregatorTests.java      | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorTests.java
index 6d105c27a692f..e37d50e53a160 100644
--- a/server/src/test/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorTests.java
+++ b/server/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())) {