Browse Source

Dry up test code that extracts TotalHits from SearchResponse (#103269)

We have the pattern of getting a response only for its `TotalHits` in a couple of spots,
some of them still leaking the `SearchResponse`.
This PR dries up those that I could quickly identify to use the new utilty method to simplify this pattern
or just moves the code to `assertHitCount` where possible.
Armin Braun 1 year ago
parent
commit
7a63a666e2
15 changed files with 97 additions and 175 deletions
  1. 16 26
      modules/reindex/src/internalClusterTest/java/org/elasticsearch/index/reindex/CrossClusterReindexIT.java
  2. 2 1
      modules/reindex/src/test/java/org/elasticsearch/reindex/DeleteByQueryBasicTests.java
  3. 3 2
      modules/repository-url/src/internalClusterTest/java/org/elasticsearch/repositories/url/URLSnapshotRestoreIT.java
  4. 2 1
      plugins/repository-hdfs/src/test/java/org/elasticsearch/repositories/hdfs/HdfsTests.java
  5. 2 8
      server/src/internalClusterTest/java/org/elasticsearch/index/shard/SearchIdleIT.java
  6. 15 27
      test/framework/src/main/java/org/elasticsearch/repositories/blobstore/ESBlobStoreRepositoryIntegTestCase.java
  7. 4 9
      test/framework/src/main/java/org/elasticsearch/snapshots/AbstractSnapshotIntegTestCase.java
  8. 8 10
      test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java
  9. 5 7
      x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/CloseFollowerIndexIT.java
  10. 19 40
      x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/IndexFollowingIT.java
  11. 6 19
      x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/LocalIndexFollowingIT.java
  12. 3 4
      x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/RestartIndexFollowingIT.java
  13. 4 8
      x-pack/plugin/ccr/src/test/java/org/elasticsearch/xpack/CcrIntegTestCase.java
  14. 6 12
      x-pack/plugin/watcher/src/internalClusterTest/java/org/elasticsearch/xpack/watcher/test/AbstractWatcherIntegrationTestCase.java
  15. 2 1
      x-pack/plugin/watcher/src/internalClusterTest/java/org/elasticsearch/xpack/watcher/test/integration/WatchAckTests.java

+ 16 - 26
modules/reindex/src/internalClusterTest/java/org/elasticsearch/index/reindex/CrossClusterReindexIT.java

@@ -9,12 +9,12 @@
 package org.elasticsearch.index.reindex;
 
 import org.apache.lucene.search.TotalHits;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.index.query.MatchAllQueryBuilder;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.reindex.ReindexPlugin;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.test.AbstractMultiClustersTestCase;
 
 import java.util.Collection;
@@ -60,11 +60,9 @@ public class CrossClusterReindexIT extends AbstractMultiClustersTestCase {
         new ReindexRequestBuilder(client(LOCAL_CLUSTER)).source(sourceIndexInRemote).destination("desc-index-001").get();
 
         assertTrue("Number of documents in source and desc indexes does not match", waitUntil(() -> {
-            SearchResponse resp = client(LOCAL_CLUSTER).prepareSearch("desc-index-001")
-                .setQuery(new MatchAllQueryBuilder())
-                .setSize(1000)
-                .get();
-            final TotalHits totalHits = resp.getHits().getTotalHits();
+            final TotalHits totalHits = SearchResponseUtils.getTotalHits(
+                client(LOCAL_CLUSTER).prepareSearch("desc-index-001").setQuery(new MatchAllQueryBuilder()).setSize(1000)
+            );
             return totalHits.relation == TotalHits.Relation.EQUAL_TO && totalHits.value == docsNumber;
         }));
     }
@@ -77,11 +75,9 @@ public class CrossClusterReindexIT extends AbstractMultiClustersTestCase {
         new ReindexRequestBuilder(client(LOCAL_CLUSTER)).source(sourceIndexInRemote).destination("test-index-001").get();
 
         assertTrue("Number of documents in source and desc indexes does not match", waitUntil(() -> {
-            SearchResponse resp = client(LOCAL_CLUSTER).prepareSearch("test-index-001")
-                .setQuery(new MatchAllQueryBuilder())
-                .setSize(1000)
-                .get();
-            final TotalHits totalHits = resp.getHits().getTotalHits();
+            final TotalHits totalHits = SearchResponseUtils.getTotalHits(
+                client(LOCAL_CLUSTER).prepareSearch("test-index-001").setQuery(new MatchAllQueryBuilder()).setSize(1000)
+            );
             return totalHits.relation == TotalHits.Relation.EQUAL_TO && totalHits.value == docsNumber;
         }));
     }
@@ -108,11 +104,9 @@ public class CrossClusterReindexIT extends AbstractMultiClustersTestCase {
             }
 
             assertTrue("Number of documents in source and desc indexes does not match", waitUntil(() -> {
-                SearchResponse resp = client(LOCAL_CLUSTER).prepareSearch("test-index-001")
-                    .setQuery(new MatchAllQueryBuilder())
-                    .setSize(1000)
-                    .get();
-                final TotalHits totalHits = resp.getHits().getTotalHits();
+                final TotalHits totalHits = SearchResponseUtils.getTotalHits(
+                    client(LOCAL_CLUSTER).prepareSearch("test-index-001").setQuery(new MatchAllQueryBuilder()).setSize(1000)
+                );
                 return totalHits.relation == TotalHits.Relation.EQUAL_TO && totalHits.value == docsNumber;
             }));
         }
@@ -142,11 +136,9 @@ public class CrossClusterReindexIT extends AbstractMultiClustersTestCase {
         new ReindexRequestBuilder(client(LOCAL_CLUSTER)).source(sourceIndexInRemote).destination("desc-index-001").get();
 
         assertTrue("Number of documents in source and desc indexes does not match", waitUntil(() -> {
-            SearchResponse resp = client(LOCAL_CLUSTER).prepareSearch("desc-index-001")
-                .setQuery(new MatchAllQueryBuilder())
-                .setSize(1000)
-                .get();
-            final TotalHits totalHits = resp.getHits().getTotalHits();
+            final TotalHits totalHits = SearchResponseUtils.getTotalHits(
+                client(LOCAL_CLUSTER).prepareSearch("desc-index-001").setQuery(new MatchAllQueryBuilder()).setSize(1000)
+            );
             return totalHits.relation == TotalHits.Relation.EQUAL_TO && totalHits.value == docsNumber;
         }));
     }
@@ -160,11 +152,9 @@ public class CrossClusterReindexIT extends AbstractMultiClustersTestCase {
         new ReindexRequestBuilder(client(LOCAL_CLUSTER)).source(sourceIndexInRemote).destination("desc-index-001").get();
 
         assertTrue("Number of documents in source and desc indexes does not match", waitUntil(() -> {
-            SearchResponse resp = client(LOCAL_CLUSTER).prepareSearch("desc-index-001")
-                .setQuery(new MatchAllQueryBuilder())
-                .setSize(1000)
-                .get();
-            final TotalHits totalHits = resp.getHits().getTotalHits();
+            final TotalHits totalHits = SearchResponseUtils.getTotalHits(
+                client(LOCAL_CLUSTER).prepareSearch("desc-index-001").setQuery(new MatchAllQueryBuilder()).setSize(1000)
+            );
             return totalHits.relation == TotalHits.Relation.EQUAL_TO && totalHits.value == docsNumber;
         }));
     }

+ 2 - 1
modules/reindex/src/test/java/org/elasticsearch/reindex/DeleteByQueryBasicTests.java

@@ -24,6 +24,7 @@ import org.elasticsearch.index.reindex.AbstractBulkByScrollRequest;
 import org.elasticsearch.index.reindex.BulkByScrollResponse;
 import org.elasticsearch.index.reindex.DeleteByQueryRequestBuilder;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.test.InternalSettingsPlugin;
 
@@ -159,7 +160,7 @@ public class DeleteByQueryBasicTests extends ReindexTestCase {
         String routing = String.valueOf(randomIntBetween(2, docs));
 
         logger.info("--> counting documents with routing [{}]", routing);
-        long expected = prepareSearch().setSize(0).setRouting(routing).get().getHits().getTotalHits().value;
+        long expected = SearchResponseUtils.getTotalHitsValue(prepareSearch().setSize(0).setRouting(routing));
 
         logger.info("--> delete all documents with routing [{}] with a delete-by-query", routing);
         DeleteByQueryRequestBuilder delete = deleteByQuery().source("test").filter(QueryBuilders.matchAllQuery());

+ 3 - 2
modules/repository-url/src/internalClusterTest/java/org/elasticsearch/repositories/url/URLSnapshotRestoreIT.java

@@ -26,6 +26,7 @@ import java.util.Collection;
 import java.util.Collections;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 
@@ -63,7 +64,7 @@ public class URLSnapshotRestoreIT extends ESIntegTestCase {
             indexDoc("test-idx", Integer.toString(i), "foo", "bar" + i);
         }
         refresh();
-        assertThat(client.prepareSearch("test-idx").setSize(0).get().getHits().getTotalHits().value, equalTo(100L));
+        assertHitCount(client.prepareSearch("test-idx").setSize(0), 100);
 
         logger.info("--> snapshot");
         CreateSnapshotResponse createSnapshotResponse = client.admin()
@@ -110,7 +111,7 @@ public class URLSnapshotRestoreIT extends ESIntegTestCase {
             .get();
         assertThat(restoreSnapshotResponse.getRestoreInfo().totalShards(), greaterThan(0));
 
-        assertThat(client.prepareSearch("test-idx").setSize(0).get().getHits().getTotalHits().value, equalTo(100L));
+        assertHitCount(client.prepareSearch("test-idx").setSize(0), 100);
 
         logger.info("--> list available shapshots");
         GetSnapshotsResponse getSnapshotsResponse = client.admin().cluster().prepareGetSnapshots("url-repo").get();

+ 2 - 1
plugins/repository-hdfs/src/test/java/org/elasticsearch/repositories/hdfs/HdfsTests.java

@@ -19,6 +19,7 @@ import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.repositories.RepositoryException;
 import org.elasticsearch.repositories.blobstore.BlobStoreRepository;
 import org.elasticsearch.repositories.blobstore.BlobStoreTestUtil;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.snapshots.SnapshotState;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 
@@ -276,6 +277,6 @@ public class HdfsTests extends ESSingleNodeTestCase {
     }
 
     private long count(Client client, String index) {
-        return client.prepareSearch(index).setSize(0).get().getHits().getTotalHits().value;
+        return SearchResponseUtils.getTotalHitsValue(client.prepareSearch(index).setSize(0));
     }
 }

+ 2 - 8
server/src/internalClusterTest/java/org/elasticsearch/index/shard/SearchIdleIT.java

@@ -26,6 +26,7 @@ import org.elasticsearch.index.query.ExistsQueryBuilder;
 import org.elasticsearch.index.query.RangeQueryBuilder;
 import org.elasticsearch.index.refresh.RefreshStats;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.XContentType;
@@ -48,14 +49,7 @@ import static org.hamcrest.Matchers.equalTo;
 public class SearchIdleIT extends ESSingleNodeTestCase {
 
     public void testAutomaticRefreshSearch() throws InterruptedException {
-        runTestAutomaticRefresh(numDocs -> {
-            var resp = client().prepareSearch("test").get();
-            try {
-                return resp.getHits().getTotalHits().value;
-            } finally {
-                resp.decRef();
-            }
-        });
+        runTestAutomaticRefresh(numDocs -> SearchResponseUtils.getTotalHitsValue(client().prepareSearch("test")));
     }
 
     public void testAutomaticRefreshGet() throws InterruptedException {

+ 15 - 27
test/framework/src/main/java/org/elasticsearch/repositories/blobstore/ESBlobStoreRepositoryIntegTestCase.java

@@ -36,6 +36,7 @@ import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.repositories.Repository;
 import org.elasticsearch.repositories.RepositoryData;
 import org.elasticsearch.repositories.RepositoryMissingException;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.snapshots.AbstractSnapshotIntegTestCase;
 import org.elasticsearch.snapshots.SnapshotMissingException;
 import org.elasticsearch.snapshots.SnapshotRestoreException;
@@ -333,20 +334,8 @@ public abstract class ESBlobStoreRepositoryIntegTestCase extends ESIntegTestCase
                     logger.info("--> add random documents to {}", index);
                     addRandomDocuments(index, randomIntBetween(10, 1000));
                 } else {
-                    var resp = prepareSearch(index).setSize(0).get();
-                    final int docCount;
-                    try {
-                        docCount = (int) resp.getHits().getTotalHits().value;
-                    } finally {
-                        resp.decRef();
-                    }
-                    int deleteCount = randomIntBetween(1, docCount);
-                    logger.info("--> delete {} random documents from {}", deleteCount, index);
-                    for (int i = 0; i < deleteCount; i++) {
-                        int doc = randomIntBetween(0, docCount - 1);
-                        client().prepareDelete(index, Integer.toString(doc)).get();
-                    }
-                    client().admin().indices().prepareRefresh(index).get();
+                    final int docCount = (int) SearchResponseUtils.getTotalHitsValue(prepareSearch(index).setSize(0));
+                    deleteRandomDocs(index, docCount);
                 }
             }
 
@@ -395,13 +384,7 @@ public abstract class ESBlobStoreRepositoryIntegTestCase extends ESIntegTestCase
             if (randomBoolean() && i > 0) { // don't delete on the first iteration
                 int docCount = docCounts[i - 1];
                 if (docCount > 0) {
-                    int deleteCount = randomIntBetween(1, docCount);
-                    logger.info("--> delete {} random documents from {}", deleteCount, indexName);
-                    for (int j = 0; j < deleteCount; j++) {
-                        int doc = randomIntBetween(0, docCount - 1);
-                        client().prepareDelete(indexName, Integer.toString(doc)).get();
-                    }
-                    client().admin().indices().prepareRefresh(indexName).get();
+                    deleteRandomDocs(indexName, docCount);
                 }
             } else {
                 int docCount = randomIntBetween(10, 1000);
@@ -409,12 +392,7 @@ public abstract class ESBlobStoreRepositoryIntegTestCase extends ESIntegTestCase
                 addRandomDocuments(indexName, docCount);
             }
             // Check number of documents in this iteration
-            var resp = prepareSearch(indexName).setSize(0).get();
-            try {
-                docCounts[i] = (int) resp.getHits().getTotalHits().value;
-            } finally {
-                resp.decRef();
-            }
+            docCounts[i] = (int) SearchResponseUtils.getTotalHitsValue(prepareSearch(indexName).setSize(0));
             logger.info("-->  create snapshot {}:{} with {} documents", repoName, snapshotName + "-" + i, docCounts[i]);
             assertSuccessfulSnapshot(
                 clusterAdmin().prepareCreateSnapshot(repoName, snapshotName + "-" + i).setWaitForCompletion(true).setIndices(indexName)
@@ -446,6 +424,16 @@ public abstract class ESBlobStoreRepositoryIntegTestCase extends ESIntegTestCase
         }
     }
 
+    private void deleteRandomDocs(String indexName, int existingDocCount) {
+        int deleteCount = randomIntBetween(1, existingDocCount);
+        logger.info("--> delete {} random documents from {}", deleteCount, indexName);
+        for (int j = 0; j < deleteCount; j++) {
+            int doc = randomIntBetween(0, existingDocCount - 1);
+            client().prepareDelete(indexName, Integer.toString(doc)).get();
+        }
+        client().admin().indices().prepareRefresh(indexName).get();
+    }
+
     public void testIndicesDeletedFromRepository() throws Exception {
         final String repoName = createRepository(randomRepositoryName());
         Client client = client();

+ 4 - 9
test/framework/src/main/java/org/elasticsearch/snapshots/AbstractSnapshotIntegTestCase.java

@@ -16,7 +16,6 @@ import org.elasticsearch.action.admin.cluster.repositories.get.TransportGetRepos
 import org.elasticsearch.action.admin.cluster.snapshots.create.CreateSnapshotResponse;
 import org.elasticsearch.action.admin.cluster.snapshots.get.GetSnapshotsRequest;
 import org.elasticsearch.action.index.IndexRequestBuilder;
-import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.support.GroupedActionListener;
 import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
@@ -47,6 +46,7 @@ import org.elasticsearch.repositories.ShardGenerations;
 import org.elasticsearch.repositories.blobstore.BlobStoreRepository;
 import org.elasticsearch.repositories.blobstore.BlobStoreTestUtil;
 import org.elasticsearch.repositories.blobstore.ChecksumBlobStoreFormat;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.snapshots.mockstore.MockRepository;
@@ -502,14 +502,9 @@ public abstract class AbstractSnapshotIntegTestCase extends ESIntegTestCase {
     }
 
     protected long getCountForIndex(String indexName) {
-        var resp = client().search(
-            new SearchRequest(new SearchRequest(indexName).source(new SearchSourceBuilder().size(0).trackTotalHits(true)))
-        ).actionGet();
-        try {
-            return resp.getHits().getTotalHits().value;
-        } finally {
-            resp.decRef();
-        }
+        return SearchResponseUtils.getTotalHitsValue(
+            client().prepareSearch(indexName).setSource(new SearchSourceBuilder().size(0).trackTotalHits(true))
+        );
     }
 
     protected void assertDocCount(String index, long count) {

+ 8 - 10
test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java

@@ -129,6 +129,7 @@ import org.elasticsearch.script.MockScriptService;
 import org.elasticsearch.search.ConcurrentSearchTestPlugin;
 import org.elasticsearch.search.MockSearchService;
 import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.test.client.RandomizingClient;
 import org.elasticsearch.test.disruption.NetworkDisruption;
@@ -1064,17 +1065,14 @@ public abstract class ESIntegTestCase extends ESTestCase {
 
             if (lastKnownCount >= numDocs) {
                 try {
-                    var resp = prepareSearch().setTrackTotalHits(true).setSize(0).setQuery(matchAllQuery()).get();
-                    try {
-                        long count = resp.getHits().getTotalHits().value;
-                        if (count == lastKnownCount) {
-                            // no progress - try to refresh for the next time
-                            indicesAdmin().prepareRefresh().get();
-                        }
-                        lastKnownCount = count;
-                    } finally {
-                        resp.decRef();
+                    long count = SearchResponseUtils.getTotalHitsValue(
+                        prepareSearch().setTrackTotalHits(true).setSize(0).setQuery(matchAllQuery())
+                    );
+                    if (count == lastKnownCount) {
+                        // no progress - try to refresh for the next time
+                        indicesAdmin().prepareRefresh().get();
                     }
+                    lastKnownCount = count;
                 } catch (Exception e) { // count now acts like search and barfs if all shards failed...
                     logger.debug("failed to executed count", e);
                     throw e;

+ 5 - 7
x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/CloseFollowerIndexIT.java

@@ -8,7 +8,6 @@ package org.elasticsearch.xpack.ccr;
 
 import org.elasticsearch.action.admin.indices.close.CloseIndexRequest;
 import org.elasticsearch.action.admin.indices.open.OpenIndexRequest;
-import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.support.ActiveShardCount;
 import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.cluster.ClusterState;
@@ -19,6 +18,7 @@ import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.engine.ReadOnlyEngine;
 import org.elasticsearch.index.shard.CloseFollowerIndexErrorSuppressionHelper;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.CcrIntegTestCase;
 import org.elasticsearch.xpack.core.ccr.action.PutFollowAction;
@@ -125,14 +125,12 @@ public class CloseFollowerIndexIT extends CcrIntegTestCase {
         ensureFollowerGreen("index2");
 
         refresh(leaderClient(), "index1");
-        SearchRequest leaderSearchRequest = new SearchRequest("index1");
-        leaderSearchRequest.source().trackTotalHits(true);
-        long leaderIndexDocs = leaderClient().search(leaderSearchRequest).actionGet().getHits().getTotalHits().value;
+        long leaderIndexDocs = SearchResponseUtils.getTotalHitsValue(leaderClient().prepareSearch("index1").setTrackTotalHits(true));
         assertBusy(() -> {
             refresh(followerClient(), "index2");
-            SearchRequest followerSearchRequest = new SearchRequest("index2");
-            followerSearchRequest.source().trackTotalHits(true);
-            long followerIndexDocs = followerClient().search(followerSearchRequest).actionGet().getHits().getTotalHits().value;
+            long followerIndexDocs = SearchResponseUtils.getTotalHitsValue(
+                followerClient().prepareSearch("index2").setTrackTotalHits(true)
+            );
             assertThat(followerIndexDocs, equalTo(leaderIndexDocs));
         }, 30L, TimeUnit.SECONDS);
     }

+ 19 - 40
x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/IndexFollowingIT.java

@@ -119,6 +119,7 @@ import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.elasticsearch.xpack.ccr.CcrRetentionLeases.retentionLeaseId;
 import static org.hamcrest.Matchers.containsString;
@@ -369,9 +370,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             leaderClient().prepareIndex("index1").setId(Long.toString(i)).setSource(source, XContentType.JSON).get();
         }
 
-        assertBusy(
-            () -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs))
-        );
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), firstBatchNumDocs));
         MappingMetadata mappingMetadata = followerClient().admin().indices().prepareGetMappings("index2").get().getMappings().get("index2");
         assertThat(XContentMapValues.extractValue("properties.f.type", mappingMetadata.sourceAsMap()), equalTo("integer"));
         assertThat(XContentMapValues.extractValue("properties.k", mappingMetadata.sourceAsMap()), nullValue());
@@ -382,12 +381,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             leaderClient().prepareIndex("index1").setId(Long.toString(i)).setSource(source, XContentType.JSON).get();
         }
 
-        assertBusy(
-            () -> assertThat(
-                followerClient().prepareSearch("index2").get().getHits().getTotalHits().value,
-                equalTo(firstBatchNumDocs + secondBatchNumDocs)
-            )
-        );
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), firstBatchNumDocs + secondBatchNumDocs));
         mappingMetadata = followerClient().admin().indices().prepareGetMappings("index2").get().getMappings().get("index2");
         assertThat(XContentMapValues.extractValue("properties.f.type", mappingMetadata.sourceAsMap()), equalTo("integer"));
         assertThat(XContentMapValues.extractValue("properties.k.type", mappingMetadata.sourceAsMap()), equalTo("long"));
@@ -413,7 +407,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
 
         leaderClient().prepareIndex("index1").setId("1").setSource("{\"f\":1}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
         pauseFollow("index2");
 
         MappingMetadata mappingMetadata = followerClient().admin().indices().prepareGetMappings("index2").get().getMappings().get("index2");
@@ -711,7 +705,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
 
         leaderClient().prepareIndex("index1").setId("1").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
 
         leaderClient().admin().indices().close(new CloseIndexRequest("index1")).actionGet();
         assertBusy(() -> {
@@ -735,7 +729,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
 
         leaderClient().admin().indices().open(new OpenIndexRequest("index1")).actionGet();
         leaderClient().prepareIndex("index1").setId("2").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(2L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 2));
 
         pauseFollow("index2");
     }
@@ -757,7 +751,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
 
         leaderClient().prepareIndex("index1").setId("1").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
 
         followerClient().admin().indices().close(new CloseIndexRequest("index2").masterNodeTimeout(TimeValue.MAX_VALUE)).actionGet();
         leaderClient().prepareIndex("index1").setId("2").setSource("{}", XContentType.JSON).get();
@@ -769,7 +763,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             assertThat(response.getStatsResponses().get(0).status().failedWriteRequests(), greaterThanOrEqualTo(1L));
         });
         followerClient().admin().indices().open(new OpenIndexRequest("index2").masterNodeTimeout(TimeValue.MAX_VALUE)).actionGet();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(2L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 2));
 
         pauseFollow("index2");
     }
@@ -791,7 +785,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
 
         leaderClient().prepareIndex("index1").setId("1").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
 
         leaderClient().admin().indices().delete(new DeleteIndexRequest("index1")).actionGet();
         assertBusy(() -> {
@@ -872,7 +866,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
 
         leaderClient().prepareIndex("index1").setId("1").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
 
         followerClient().admin().indices().delete(new DeleteIndexRequest("index2").masterNodeTimeout(TimeValue.MAX_VALUE)).actionGet();
         leaderClient().prepareIndex("index1").setId("2").setSource("{}", XContentType.JSON).get();
@@ -935,7 +929,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         PutFollowAction.Request followRequest = putFollow("index1", "index2");
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
         leaderClient().prepareIndex("index1").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> { assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)); });
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
 
         // Indexing directly into index2 would fail now, because index2 is a follow index.
         // We can't test this here because an assertion trips before an actual error is thrown and then index call hangs.
@@ -952,7 +946,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             .setSource("{}", XContentType.JSON)
             .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
             .get();
-        assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(2L));
+        assertHitCount(followerClient().prepareSearch("index2"), 2);
     }
 
     public void testUnknownClusterAlias() throws Exception {
@@ -1024,9 +1018,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         for (long i = 0; i < firstBatchNumDocs; i++) {
             leaderClient().prepareIndex("leader").setSource("{}", XContentType.JSON).get();
         }
-        assertBusy(
-            () -> assertThat(followerClient().prepareSearch("follower").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs))
-        );
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("follower"), firstBatchNumDocs));
 
         // Sanity check that the setting has not been set in follower index:
         {
@@ -1053,10 +1045,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             assertThat(getFollowTaskSettingsVersion("follower"), equalTo(2L));
 
             try {
-                assertThat(
-                    followerClient().prepareSearch("follower").get().getHits().getTotalHits().value,
-                    equalTo(firstBatchNumDocs + secondBatchNumDocs)
-                );
+                assertHitCount(followerClient().prepareSearch("follower"), firstBatchNumDocs + secondBatchNumDocs);
             } catch (Exception e) {
                 throw new AssertionError("error while searching", e);
             }
@@ -1080,9 +1069,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         for (long i = 0; i < firstBatchNumDocs; i++) {
             leaderClient().prepareIndex("leader").setSource("{}", XContentType.JSON).get();
         }
-        assertBusy(
-            () -> assertThat(followerClient().prepareSearch("follower").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs))
-        );
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("follower"), firstBatchNumDocs));
 
         // Sanity check that the setting has not been set in follower index:
         {
@@ -1108,10 +1095,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             assertThat(getFollowTaskSettingsVersion("follower"), equalTo(2L));
 
             try {
-                assertThat(
-                    followerClient().prepareSearch("follower").get().getHits().getTotalHits().value,
-                    equalTo(firstBatchNumDocs + secondBatchNumDocs)
-                );
+                assertHitCount(followerClient().prepareSearch("follower"), firstBatchNumDocs + secondBatchNumDocs);
             } catch (Exception e) {
                 throw new AssertionError("error while searching", e);
             }
@@ -1133,9 +1117,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             leaderClient().prepareIndex("leader").setSource("{}", XContentType.JSON).get();
         }
 
-        assertBusy(
-            () -> assertThat(followerClient().prepareSearch("follower").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs))
-        );
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("follower"), firstBatchNumDocs));
         assertThat(getFollowTaskSettingsVersion("follower"), equalTo(1L));
         assertThat(getFollowTaskMappingVersion("follower"), equalTo(1L));
 
@@ -1185,10 +1167,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
             );
 
             try {
-                assertThat(
-                    followerClient().prepareSearch("follower").get().getHits().getTotalHits().value,
-                    equalTo(firstBatchNumDocs + secondBatchNumDocs)
-                );
+                assertHitCount(followerClient().prepareSearch("follower"), firstBatchNumDocs + secondBatchNumDocs);
             } catch (Exception e) {
                 throw new AssertionError("error while searching", e);
             }
@@ -1574,7 +1553,7 @@ public class IndexFollowingIT extends CcrIntegTestCase {
         followerClient().execute(PutFollowAction.INSTANCE, followRequest).get();
 
         leaderClient().prepareIndex("index1").setId("1").setSource("{}", XContentType.JSON).get();
-        assertBusy(() -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(1L)));
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), 1));
 
         assertBusy(() -> {
             String action = ShardFollowTask.NAME + "[c]";

+ 6 - 19
x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/LocalIndexFollowingIT.java

@@ -32,6 +32,7 @@ import java.util.concurrent.CountDownLatch;
 import java.util.stream.StreamSupport;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
@@ -53,21 +54,14 @@ public class LocalIndexFollowingIT extends CcrSingleNodeTestCase {
         final PutFollowAction.Request followRequest = getPutFollowRequest("leader", "follower");
         client().execute(PutFollowAction.INSTANCE, followRequest).get();
 
-        assertBusy(
-            () -> { assertThat(client().prepareSearch("follower").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs)); }
-        );
+        assertBusy(() -> assertHitCount(client().prepareSearch("follower"), firstBatchNumDocs));
 
         final long secondBatchNumDocs = randomIntBetween(2, 64);
         for (int i = 0; i < secondBatchNumDocs; i++) {
             prepareIndex("leader").setSource("{}", XContentType.JSON).get();
         }
 
-        assertBusy(() -> {
-            assertThat(
-                client().prepareSearch("follower").get().getHits().getTotalHits().value,
-                equalTo(firstBatchNumDocs + secondBatchNumDocs)
-            );
-        });
+        assertBusy(() -> assertHitCount(client().prepareSearch("follower"), firstBatchNumDocs + secondBatchNumDocs));
 
         PauseFollowAction.Request pauseRequest = new PauseFollowAction.Request("follower");
         client().execute(PauseFollowAction.INSTANCE, pauseRequest);
@@ -78,12 +72,7 @@ public class LocalIndexFollowingIT extends CcrSingleNodeTestCase {
         }
 
         client().execute(ResumeFollowAction.INSTANCE, getResumeFollowRequest("follower")).get();
-        assertBusy(() -> {
-            assertThat(
-                client().prepareSearch("follower").get().getHits().getTotalHits().value,
-                equalTo(firstBatchNumDocs + secondBatchNumDocs + thirdBatchNumDocs)
-            );
-        });
+        assertBusy(() -> assertHitCount(client().prepareSearch("follower"), firstBatchNumDocs + secondBatchNumDocs + thirdBatchNumDocs));
         ensureEmptyWriteBuffers();
     }
 
@@ -136,9 +125,7 @@ public class LocalIndexFollowingIT extends CcrSingleNodeTestCase {
                 assertEquals(firstBatchNumDocs, indexingPressure.stats().getCurrentPrimaryOps());
             });
             blocker.countDown();
-            assertBusy(
-                () -> { assertThat(client().prepareSearch("follower").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs)); }
-            );
+            assertBusy(() -> assertHitCount(client().prepareSearch("follower"), firstBatchNumDocs));
             ensureEmptyWriteBuffers();
         } finally {
             if (blocker.getCount() > 0) {
@@ -210,7 +197,7 @@ public class LocalIndexFollowingIT extends CcrSingleNodeTestCase {
             prepareIndex("index-1").setSource("{}", XContentType.JSON).get();
         }
         client().execute(PutFollowAction.INSTANCE, getPutFollowRequest("index-1", "index-2")).get();
-        assertBusy(() -> assertThat(client().prepareSearch("index-2").get().getHits().getTotalHits().value, equalTo((long) numDocs)));
+        assertBusy(() -> assertHitCount(client().prepareSearch("index-2"), numDocs));
 
         // Then switch index-1 to be a follower of index-0
         assertAcked(client().admin().indices().prepareCreate("index-0").setSource(settings, XContentType.JSON));

+ 3 - 4
x-pack/plugin/ccr/src/internalClusterTest/java/org/elasticsearch/xpack/ccr/RestartIndexFollowingIT.java

@@ -34,6 +34,7 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
@@ -78,9 +79,7 @@ public class RestartIndexFollowingIT extends CcrIntegTestCase {
             leaderClient().prepareIndex("index1").setId(Integer.toString(i)).setSource(source, XContentType.JSON).get();
         }
 
-        assertBusy(
-            () -> assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(firstBatchNumDocs))
-        );
+        assertBusy(() -> assertHitCount(followerClient().prepareSearch("index2"), firstBatchNumDocs));
 
         getFollowerCluster().fullRestart();
         ensureFollowerGreen("index2");
@@ -115,7 +114,7 @@ public class RestartIndexFollowingIT extends CcrIntegTestCase {
                     followerClient().execute(PutFollowAction.INSTANCE, putFollow("index1", "index2", ActiveShardCount.ALL)).actionGet();
                 }
             }
-            assertThat(followerClient().prepareSearch("index2").get().getHits().getTotalHits().value, equalTo(totalDocs));
+            assertHitCount(followerClient().prepareSearch("index2"), totalDocs);
         }, 30L, TimeUnit.SECONDS);
 
         cleanRemoteCluster();

+ 4 - 8
x-pack/plugin/ccr/src/test/java/org/elasticsearch/xpack/CcrIntegTestCase.java

@@ -61,6 +61,7 @@ import org.elasticsearch.license.LicenseSettings;
 import org.elasticsearch.license.LicensesMetadata;
 import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.snapshots.RestoreInfo;
 import org.elasticsearch.snapshots.RestoreService;
@@ -794,14 +795,9 @@ public abstract class CcrIntegTestCase extends ESTestCase {
 
             if (lastKnownCount >= numDocs) {
                 try {
-                    long count = indexer.getClient()
-                        .prepareSearch()
-                        .setTrackTotalHits(true)
-                        .setSize(0)
-                        .setQuery(QueryBuilders.matchAllQuery())
-                        .get()
-                        .getHits()
-                        .getTotalHits().value;
+                    long count = SearchResponseUtils.getTotalHitsValue(
+                        indexer.getClient().prepareSearch().setTrackTotalHits(true).setSize(0).setQuery(QueryBuilders.matchAllQuery())
+                    );
 
                     if (count == lastKnownCount) {
                         // no progress - try to refresh for the next time

+ 6 - 12
x-pack/plugin/watcher/src/internalClusterTest/java/org/elasticsearch/xpack/watcher/test/AbstractWatcherIntegrationTestCase.java

@@ -34,6 +34,7 @@ import org.elasticsearch.license.XPackLicenseState;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.script.MockMustacheScriptEngine;
 import org.elasticsearch.search.SearchHit;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.ESIntegTestCase.ClusterScope;
@@ -312,8 +313,7 @@ public abstract class AbstractWatcherIntegrationTestCase extends ESIntegTestCase
     }
 
     protected long docCount(String index, SearchSourceBuilder source) {
-        SearchRequestBuilder builder = prepareSearch(index).setSource(source).setSize(0);
-        return builder.get().getHits().getTotalHits().value;
+        return SearchResponseUtils.getTotalHitsValue(prepareSearch(index).setSource(source).setSize(0));
     }
 
     protected SearchResponse searchHistory(SearchSourceBuilder builder) {
@@ -416,16 +416,10 @@ public abstract class AbstractWatcherIntegrationTestCase extends ESIntegTestCase
 
     protected long findNumberOfPerformedActions(String watchName) {
         refresh();
-        SearchResponse searchResponse = prepareSearch(HistoryStoreField.DATA_STREAM + "*").setIndicesOptions(
-            IndicesOptions.lenientExpandOpen()
-        ).setQuery(boolQuery().must(matchQuery("watch_id", watchName)).must(matchQuery("state", ExecutionState.EXECUTED.id()))).get();
-        long totalHistsValue;
-        try {
-            totalHistsValue = searchResponse.getHits().getTotalHits().value;
-        } finally {
-            searchResponse.decRef();
-        }
-        return totalHistsValue;
+        return SearchResponseUtils.getTotalHitsValue(
+            prepareSearch(HistoryStoreField.DATA_STREAM + "*").setIndicesOptions(IndicesOptions.lenientExpandOpen())
+                .setQuery(boolQuery().must(matchQuery("watch_id", watchName)).must(matchQuery("state", ExecutionState.EXECUTED.id())))
+        );
     }
 
     protected void assertWatchWithNoActionNeeded(final String watchName, final long expectedWatchActionsWithNoActionNeeded)

+ 2 - 1
x-pack/plugin/watcher/src/internalClusterTest/java/org/elasticsearch/xpack/watcher/test/integration/WatchAckTests.java

@@ -13,6 +13,7 @@ import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.protocol.xpack.watcher.PutWatchResponse;
+import org.elasticsearch.search.SearchResponseUtils;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.watcher.actions.ActionStatus;
 import org.elasticsearch.xpack.core.watcher.execution.ExecutionState;
@@ -234,7 +235,7 @@ public class WatchAckTests extends AbstractWatcherIntegrationTestCase {
         assertThat(ackResponse.getStatus().actionStatus("_id").ackStatus().state(), is(ActionStatus.AckStatus.State.ACKED));
 
         refresh("actions");
-        long countAfterAck = prepareSearch("actions").setQuery(matchAllQuery()).get().getHits().getTotalHits().value;
+        long countAfterAck = SearchResponseUtils.getTotalHitsValue(prepareSearch("actions").setQuery(matchAllQuery()));
         assertThat(countAfterAck, greaterThanOrEqualTo(1L));
 
         restartWatcherRandomly();