Fix Flaky Test with EndLSN (#45747)

* fix flaky test with endLSN

* fix flaky test with endLSN

Author: tvaron3

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosContainerChangeFeedTest.java

@@ -49,7 +49,6 @@
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Factory;
-import org.testng.annotations.Ignore;
 import org.testng.annotations.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -107,17 +106,17 @@ public static Object[][] changeFeedQueryCompleteAfterAvailableNowDataProvider()
     @DataProvider(name = "changeFeedQueryEndLSNDataProvider")
     public static Object[][] changeFeedQueryEndLSNDataProvider() {
         return new Object[][]{
-                // container RU, continuous ingest items, partition count
+                // container RU, continuous ingest items
                 // number of docs from cf, documents to write
 
                 // endLSN is less than number of documents
-                { 400, true, 1, 3, 6},
-                { 400, false, 1, 3, 6},
+                { 400, true, 3, 6},
+                { 400, false, 3, 6},
                 // endLSN is equal to number of documents
-                { 400, false, 1, 3, 3},
+                { 400, false, 3, 3},
                 // both partitions have more than the endLSN
-                { 11000, true, 5, 6, 30},
-                { 11000, false, 5, 6, 30},
+                { 11000, true, 6, 30},
+                { 11000, false, 6, 30},
         };
     }
 
@@ -947,12 +946,10 @@ public void split_only_notModified() throws Exception {
         assertThat(stateAfterLastDrainAttempt.getContinuation().getCompositeContinuationTokens()).hasSize(3);
     }
 
-    @Ignore
     @Test(groups = { "fast" }, dataProvider = "changeFeedQueryEndLSNDataProvider", timeOut = 100 * TIMEOUT)
     public void changeFeedQueryCompleteAfterEndLSN(
         int throughput,
         boolean shouldContinuouslyIngestItems,
-        int partitionCount,
         int expectedDocs,
         int docsToWrite) {
         String testContainerId = UUID.randomUUID().toString();
@@ -969,7 +966,7 @@ public void changeFeedQueryCompleteAfterEndLSN(
             List<FeedRange> feedRanges = testContainer.getFeedRanges().block();
             AtomicInteger currentPageCount = new AtomicInteger(0);
 
-            List<String> partitionKeys = insertDocumentsCore(partitionCount, docsToWrite, testContainer);
+            List<String> partitionKeys = insertDocumentsIntoTwoPartitions(docsToWrite, testContainer);
             CosmosChangeFeedRequestOptions cosmosChangeFeedRequestOptions =
                 CosmosChangeFeedRequestOptions.createForProcessingFromBeginning(FeedRange.forFullRange());
             ImplementationBridgeHelpers.CosmosChangeFeedRequestOptionsHelper.getCosmosChangeFeedRequestOptionsAccessor()
@@ -1146,6 +1143,42 @@ List<String> insertDocumentsCore(
         return partitionKeys;
     }
 
+    List<String> insertDocumentsIntoTwoPartitions(
+            int documentCount,
+            CosmosAsyncContainer container) {
+
+        List<ObjectNode> docs = new ArrayList<>();
+        List<String> partitionKeys = new ArrayList<>();
+
+        // these partition keys will land on two different partitions for hash v2
+        String partitionKey1 = "pk-1";
+        String partitionKey2 = "pk-8fbakldbas";
+        for (int j = 0; j < documentCount; j++) {
+            docs.add(getDocumentDefinition(partitionKey1));
+            docs.add(getDocumentDefinition(partitionKey2));
+        }
+        partitionKeys.add(partitionKey1);
+        partitionKeys.add(partitionKey2);
+
+        ArrayList<Mono<CosmosItemResponse<ObjectNode>>> result = new ArrayList<>();
+        for (ObjectNode jsonNodes : docs) {
+            result.add(container.createItem(jsonNodes));
+        }
+
+        List<ObjectNode> insertedDocs = Flux.merge(
+                        Flux.fromIterable(result),
+                        2)
+                .map(CosmosItemResponse::getItem).collectList().block();
+
+        for (ObjectNode doc : insertedDocs) {
+            partitionKeyToDocuments.put(
+                    doc.get(PARTITION_KEY_FIELD_NAME).textValue(),
+                    doc);
+        }
+        logger.info("FINISHED INSERT");
+        return partitionKeys;
+    }
+
     void deleteDocuments(
         int partitionCount,
         int documentCount) {