Fix a few more explicit SearchResponse uses in the server package

For elastic#102030
original-brownbear · Nov 10, 2023 · 1ce52a1 · 1ce52a1
1 parent 8964b87
commit 1ce52a1
Show file tree

Hide file tree

Showing 11 changed files with 161 additions and 137 deletions.
diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/IndicesRequestIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/IndicesRequestIT.java
@@ -55,7 +55,6 @@
 import org.elasticsearch.action.get.MultiGetRequest;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.search.SearchRequest;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchTransportService;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.action.support.replication.TransportReplicationActionTests;
@@ -102,7 +101,7 @@
 import java.util.function.Function;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
-import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.hamcrest.Matchers.emptyIterable;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
@@ -560,9 +559,10 @@ public void testSearchQueryThenFetch() throws Exception {
         refresh();
 
         SearchRequest searchRequest = new SearchRequest(randomIndicesOrAliases).searchType(SearchType.QUERY_THEN_FETCH);
-        SearchResponse searchResponse = internalCluster().coordOnlyNodeClient().search(searchRequest).actionGet();
-        assertNoFailures(searchResponse);
-        assertThat(searchResponse.getHits().getTotalHits().value, greaterThan(0L));
+        assertNoFailuresAndResponse(
+            internalCluster().coordOnlyNodeClient().search(searchRequest),
+            searchResponse -> assertThat(searchResponse.getHits().getTotalHits().value, greaterThan(0L))
+        );
 
         clearInterceptedActions();
         assertIndicesSubset(
@@ -589,9 +589,10 @@ public void testSearchDfsQueryThenFetch() throws Exception {
         refresh();
 
         SearchRequest searchRequest = new SearchRequest(randomIndicesOrAliases).searchType(SearchType.DFS_QUERY_THEN_FETCH);
-        SearchResponse searchResponse = internalCluster().coordOnlyNodeClient().search(searchRequest).actionGet();
-        assertNoFailures(searchResponse);
-        assertThat(searchResponse.getHits().getTotalHits().value, greaterThan(0L));
+        assertNoFailuresAndResponse(
+            internalCluster().coordOnlyNodeClient().search(searchRequest),
+            searchResponse -> assertThat(searchResponse.getHits().getTotalHits().value, greaterThan(0L))
+        );
 
         clearInterceptedActions();
         assertIndicesSubset(

diff --git a/...c/internalClusterTest/java/org/elasticsearch/action/admin/cluster/node/tasks/TasksIT.java b/...c/internalClusterTest/java/org/elasticsearch/action/admin/cluster/node/tasks/TasksIT.java
@@ -27,7 +27,6 @@
 import org.elasticsearch.action.bulk.BulkAction;
 import org.elasticsearch.action.index.IndexAction;
 import org.elasticsearch.action.search.SearchAction;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchTransportService;
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.action.support.replication.ReplicationResponse;
@@ -83,6 +82,7 @@
 import static org.elasticsearch.core.TimeValue.timeValueSeconds;
 import static org.elasticsearch.http.HttpTransportSettings.SETTING_HTTP_MAX_HEADER_SIZE;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertFutureThrows;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.containsString;
@@ -771,17 +771,19 @@ public void testTaskStoringSuccessfulResult() throws Exception {
 
         assertNoFailures(indicesAdmin().prepareRefresh(TaskResultsService.TASK_INDEX).get());
 
-        SearchResponse searchResponse = prepareSearch(TaskResultsService.TASK_INDEX).setSource(
-            SearchSourceBuilder.searchSource().query(QueryBuilders.termQuery("task.action", taskInfo.action()))
-        ).get();
-
-        assertEquals(1L, searchResponse.getHits().getTotalHits().value);
-
-        searchResponse = prepareSearch(TaskResultsService.TASK_INDEX).setSource(
-            SearchSourceBuilder.searchSource().query(QueryBuilders.termQuery("task.node", taskInfo.taskId().getNodeId()))
-        ).get();
+        assertHitCount(
+            prepareSearch(TaskResultsService.TASK_INDEX).setSource(
+                SearchSourceBuilder.searchSource().query(QueryBuilders.termQuery("task.action", taskInfo.action()))
+            ),
+            1L
+        );
 
-        assertEquals(1L, searchResponse.getHits().getTotalHits().value);
+        assertHitCount(
+            prepareSearch(TaskResultsService.TASK_INDEX).setSource(
+                SearchSourceBuilder.searchSource().query(QueryBuilders.termQuery("task.node", taskInfo.taskId().getNodeId()))
+            ),
+            1L
+        );
 
         GetTaskResponse getResponse = expectFinishedTask(taskId);
         assertEquals(result, getResponse.getTask().getResponseAsMap());

diff --git a/.../internalClusterTest/java/org/elasticsearch/action/admin/indices/create/SplitIndexIT.java b/.../internalClusterTest/java/org/elasticsearch/action/admin/indices/create/SplitIndexIT.java
@@ -24,7 +24,7 @@
 import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.index.IndexRequestBuilder;
-import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.action.search.SearchRequestBuilder;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
@@ -60,7 +60,8 @@
 import static org.elasticsearch.index.query.QueryBuilders.termQuery;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
-import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
@@ -243,27 +244,28 @@ private void splitToN(int sourceShards, int firstSplitShards, int secondSplitSha
             assertNested("first_split", numDocs);
             assertNested("second_split", numDocs);
         }
-        assertAllUniqueDocs(prepareSearch("second_split").setSize(100).setQuery(new TermsQueryBuilder("foo", "bar")).get(), numDocs);
-        assertAllUniqueDocs(prepareSearch("first_split").setSize(100).setQuery(new TermsQueryBuilder("foo", "bar")).get(), numDocs);
-        assertAllUniqueDocs(prepareSearch("source").setSize(100).setQuery(new TermsQueryBuilder("foo", "bar")).get(), numDocs);
+        assertAllUniqueDocs(prepareSearch("second_split").setSize(100).setQuery(new TermsQueryBuilder("foo", "bar")), numDocs);
+        assertAllUniqueDocs(prepareSearch("first_split").setSize(100).setQuery(new TermsQueryBuilder("foo", "bar")), numDocs);
+        assertAllUniqueDocs(prepareSearch("source").setSize(100).setQuery(new TermsQueryBuilder("foo", "bar")), numDocs);
     }
 
     public void assertNested(String index, int numDocs) {
         // now, do a nested query
-        SearchResponse searchResponse = prepareSearch(index).setQuery(
-            nestedQuery("nested1", termQuery("nested1.n_field1", "n_value1_1"), ScoreMode.Avg)
-        ).get();
-        assertNoFailures(searchResponse);
-        assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) numDocs));
+        assertNoFailuresAndResponse(
+            prepareSearch(index).setQuery(nestedQuery("nested1", termQuery("nested1.n_field1", "n_value1_1"), ScoreMode.Avg)),
+            searchResponse -> assertThat(searchResponse.getHits().getTotalHits().value, equalTo((long) numDocs))
+        );
     }
 
-    public void assertAllUniqueDocs(SearchResponse response, int numDocs) {
-        Set<String> ids = new HashSet<>();
-        for (int i = 0; i < response.getHits().getHits().length; i++) {
-            String id = response.getHits().getHits()[i].getId();
-            assertTrue("found ID " + id + " more than once", ids.add(id));
-        }
-        assertEquals(numDocs, ids.size());
+    public void assertAllUniqueDocs(SearchRequestBuilder request, int numDocs) {
+        assertResponse(request, response -> {
+            Set<String> ids = new HashSet<>();
+            for (int i = 0; i < response.getHits().getHits().length; i++) {
+                String id = response.getHits().getHits()[i].getId();
+                assertTrue("found ID " + id + " more than once", ids.add(id));
+            }
+            assertEquals(numDocs, ids.size());
+        });
     }
 
     public void testSplitIndexPrimaryTerm() throws Exception {

diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkProcessor2RetryIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkProcessor2RetryIT.java
@@ -11,7 +11,6 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.DocWriteRequest;
 import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.query.QueryBuilders;
@@ -26,6 +25,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
@@ -136,15 +136,17 @@ public void afterBulk(long executionId, BulkRequest request, Exception failure)
 
         indicesAdmin().refresh(new RefreshRequest()).get();
 
-        SearchResponse results = prepareSearch(INDEX_NAME).setQuery(QueryBuilders.matchAllQuery()).setSize(0).get();
-        assertThat(bulkProcessor.getTotalBytesInFlight(), equalTo(0L));
-        if (rejectedExecutionExpected) {
-            assertThat((int) results.getHits().getTotalHits().value, lessThanOrEqualTo(numberOfAsyncOps));
-        } else if (rejectedAfterAllRetries) {
-            assertThat((int) results.getHits().getTotalHits().value, lessThan(numberOfAsyncOps));
-        } else {
-            assertThat((int) results.getHits().getTotalHits().value, equalTo(numberOfAsyncOps));
-        }
+        final boolean finalRejectedAfterAllRetries = rejectedAfterAllRetries;
+        assertResponse(prepareSearch(INDEX_NAME).setQuery(QueryBuilders.matchAllQuery()).setSize(0), results -> {
+            assertThat(bulkProcessor.getTotalBytesInFlight(), equalTo(0L));
+            if (rejectedExecutionExpected) {
+                assertThat((int) results.getHits().getTotalHits().value, lessThanOrEqualTo(numberOfAsyncOps));
+            } else if (finalRejectedAfterAllRetries) {
+                assertThat((int) results.getHits().getTotalHits().value, lessThan(numberOfAsyncOps));
+            } else {
+                assertThat((int) results.getHits().getTotalHits().value, equalTo(numberOfAsyncOps));
+            }
+        });
     }
 
     private static void indexDocs(BulkProcessor2 processor, int numDocs) {

diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkProcessorRetryIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/BulkProcessorRetryIT.java
@@ -9,7 +9,6 @@
 
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.query.QueryBuilders;
@@ -25,6 +24,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
@@ -131,15 +131,16 @@ public void afterBulk(long executionId, BulkRequest request, Throwable failure)
 
         indicesAdmin().refresh(new RefreshRequest()).get();
 
-        SearchResponse results = prepareSearch(INDEX_NAME).setQuery(QueryBuilders.matchAllQuery()).setSize(0).get();
-
-        if (rejectedExecutionExpected) {
-            assertThat((int) results.getHits().getTotalHits().value, lessThanOrEqualTo(numberOfAsyncOps));
-        } else if (rejectedAfterAllRetries) {
-            assertThat((int) results.getHits().getTotalHits().value, lessThan(numberOfAsyncOps));
-        } else {
-            assertThat((int) results.getHits().getTotalHits().value, equalTo(numberOfAsyncOps));
-        }
+        final boolean finalRejectedAfterAllRetries = rejectedAfterAllRetries;
+        assertResponse(prepareSearch(INDEX_NAME).setQuery(QueryBuilders.matchAllQuery()).setSize(0), results -> {
+            if (rejectedExecutionExpected) {
+                assertThat((int) results.getHits().getTotalHits().value, lessThanOrEqualTo(numberOfAsyncOps));
+            } else if (finalRejectedAfterAllRetries) {
+                assertThat((int) results.getHits().getTotalHits().value, lessThan(numberOfAsyncOps));
+            } else {
+                assertThat((int) results.getHits().getTotalHits().value, equalTo(numberOfAsyncOps));
+            }
+        });
     }
 
     private void assertRetriedCorrectly(CorrelatingBackoffPolicy internalPolicy, Object bulkResponse, Throwable failure) {

diff --git a/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/WriteAckDelayIT.java b/server/src/internalClusterTest/java/org/elasticsearch/action/bulk/WriteAckDelayIT.java
@@ -9,7 +9,6 @@
 package org.elasticsearch.action.bulk;
 
 import org.elasticsearch.action.index.IndexRequestBuilder;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ESIntegTestCase;
@@ -18,6 +17,8 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
+
 public class WriteAckDelayIT extends ESIntegTestCase {
 
     @Override
@@ -42,17 +43,18 @@ public void testIndexWithWriteDelayEnabled() throws Exception {
         for (int j = 0; j < numOfChecks; j++) {
             try {
                 logger.debug("running search");
-                SearchResponse response = prepareSearch("test").get();
-                if (response.getHits().getTotalHits().value != numOfDocs) {
-                    final String message = "Count is "
-                        + response.getHits().getTotalHits().value
-                        + " but "
-                        + numOfDocs
-                        + " was expected. "
-                        + ElasticsearchAssertions.formatShardStatus(response);
-                    logger.error("{}. search response: \n{}", message, response);
-                    fail(message);
-                }
+                assertResponse(prepareSearch("test"), response -> {
+                    if (response.getHits().getTotalHits().value != numOfDocs) {
+                        final String message = "Count is "
+                            + response.getHits().getTotalHits().value
+                            + " but "
+                            + numOfDocs
+                            + " was expected. "
+                            + ElasticsearchAssertions.formatShardStatus(response);
+                        logger.error("{}. search response: \n{}", message, response);
+                        fail(message);
+                    }
+                });
             } catch (Exception e) {
                 logger.error("search failed", e);
                 throw e;