Fix ShardBulkInferenceActionFilterIT#testRestart (#127671) (#127740)

Simplify and fix the wrong assumptions in these tests.

Closes #127592
Closes #127593
Closes #127594

Author: jimczi

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilterIT.java

@@ -9,14 +9,14 @@
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
-import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
 import org.elasticsearch.action.bulk.BulkItemResponse;
 import org.elasticsearch.action.bulk.BulkRequestBuilder;
 import org.elasticsearch.action.bulk.BulkResponse;
 import org.elasticsearch.action.delete.DeleteRequestBuilder;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.settings.Settings;
@@ -243,12 +243,10 @@ public void testRestart() throws Exception {
 
     private void assertRandomBulkOperations(String indexName, Function<Boolean, Map<String, Object>> sourceSupplier) throws Exception {
         int numHits = numHits(indexName);
-        int totalBulkReqs = randomIntBetween(2, 100);
-        long totalDocs = numHits;
+        int totalBulkReqs = randomIntBetween(2, 10);
         Set<String> ids = new HashSet<>();
-
-        for (int bulkReqs = numHits; bulkReqs < totalBulkReqs; bulkReqs++) {
-            BulkRequestBuilder bulkReqBuilder = client().prepareBulk();
+        for (int bulkReqs = 0; bulkReqs < totalBulkReqs; bulkReqs++) {
+            BulkRequestBuilder bulkReqBuilder = client().prepareBulk().setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
             int totalBulkSize = randomIntBetween(1, 100);
             for (int bulkSize = 0; bulkSize < totalBulkSize; bulkSize++) {
                 if (ids.size() > 0 && rarely(random())) {
@@ -258,24 +256,15 @@ private void assertRandomBulkOperations(String indexName, Function<Boolean, Map<
                     bulkReqBuilder.add(request);
                     continue;
                 }
-                String id = Long.toString(totalDocs++);
-                boolean isIndexRequest = randomBoolean();
+                boolean isIndexRequest = ids.size() == 0 || randomBoolean();
                 Map<String, Object> source = sourceSupplier.apply(isIndexRequest);
                 if (isIndexRequest) {
+                    String id = randomAlphaOfLength(20);
                     bulkReqBuilder.add(new IndexRequestBuilder(client()).setIndex(indexName).setId(id).setSource(source));
                     ids.add(id);
                 } else {
-                    boolean isUpsert = randomBoolean();
-                    UpdateRequestBuilder request = new UpdateRequestBuilder(client()).setIndex(indexName).setDoc(source);
-                    if (isUpsert || ids.size() == 0) {
-                        request.setDocAsUpsert(true);
-                    } else {
-                        // Update already existing document
-                        id = randomFrom(ids);
-                    }
-                    request.setId(id);
-                    bulkReqBuilder.add(request);
-                    ids.add(id);
+                    String id = randomFrom(ids);
+                    bulkReqBuilder.add(new UpdateRequestBuilder(client()).setIndex(indexName).setId(id).setDoc(source));
                 }
             }
             BulkResponse bulkResponse = bulkReqBuilder.get();
@@ -294,8 +283,7 @@ private void assertRandomBulkOperations(String indexName, Function<Boolean, Map<
             }
             assertFalse(bulkResponse.hasFailures());
         }
-        client().admin().indices().refresh(new RefreshRequest(indexName)).get();
-        assertThat(numHits(indexName), equalTo(ids.size() + numHits));
+        assertThat(numHits(indexName), equalTo(numHits + ids.size()));
     }
 
     private int numHits(String indexName) throws Exception {