Dry up testChunking (#92530)

Most chunking implementations have a corresponding test regarding the number of chunks emitted, and they're all pretty much copies of each other. Extracting a test utility to capture this pattern is long overdue. Relates #89838
elastic · Jan 3, 2023 · e7fbfa1 · e7fbfa1
1 parent b990254
commit e7fbfa1
Show file tree

Hide file tree

Showing 35 changed files with 243 additions and 469 deletions.
diff --git a/...c/main/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponse.java b/...c/main/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponse.java
@@ -15,7 +15,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.ToXContentObject;
@@ -28,7 +28,7 @@
 import java.util.Objects;
 import java.util.Set;
 
-public class DesiredBalanceResponse extends ActionResponse implements ChunkedToXContent {
+public class DesiredBalanceResponse extends ActionResponse implements ChunkedToXContentObject {
 
     private final DesiredBalanceStats stats;
     private final Map<String, Map<Integer, DesiredShards>> routingTable;

diff --git a/...c/main/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponse.java b/...c/main/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponse.java
@@ -13,14 +13,14 @@
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.List;
 
-public class PendingClusterTasksResponse extends ActionResponse implements ChunkedToXContent {
+public class PendingClusterTasksResponse extends ActionResponse implements ChunkedToXContentObject {
 
     private final List<PendingClusterTask> pendingTasks;
 

diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/get/GetIndexResponse.java b/server/src/main/java/org/elasticsearch/action/admin/indices/get/GetIndexResponse.java
@@ -17,7 +17,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.xcontent.ToXContent;
@@ -35,7 +35,7 @@
 /**
  * A response for a get index action.
  */
-public class GetIndexResponse extends ActionResponse implements ChunkedToXContent {
+public class GetIndexResponse extends ActionResponse implements ChunkedToXContentObject {
 
     private Map<String, MappingMetadata> mappings = Map.of();
     private Map<String, List<AliasMetadata>> aliases = Map.of();

diff --git a/...src/main/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponse.java b/...src/main/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponse.java
@@ -15,7 +15,7 @@
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.xcontent.ParseField;
@@ -28,7 +28,7 @@
 import static org.elasticsearch.rest.BaseRestHandler.DEFAULT_INCLUDE_TYPE_NAME_POLICY;
 import static org.elasticsearch.rest.BaseRestHandler.INCLUDE_TYPE_NAME_PARAMETER;
 
-public class GetMappingsResponse extends ActionResponse implements ChunkedToXContent {
+public class GetMappingsResponse extends ActionResponse implements ChunkedToXContentObject {
 
     private static final ParseField MAPPINGS = new ParseField("mappings");
 

diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponse.java b/server/src/main/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponse.java
@@ -14,7 +14,7 @@
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.indices.recovery.RecoveryState;
 import org.elasticsearch.xcontent.ToXContent;
 
@@ -26,7 +26,7 @@
 /**
  * Information regarding the recovery state of indices and their associated shards.
  */
-public class RecoveryResponse extends BaseBroadcastResponse implements ChunkedToXContent {
+public class RecoveryResponse extends BaseBroadcastResponse implements ChunkedToXContentObject {
 
     private final Map<String, List<RecoveryState>> shardRecoveryStates;
 

diff --git a/...er/src/main/java/org/elasticsearch/action/support/broadcast/ChunkedBroadcastResponse.java b/...er/src/main/java/org/elasticsearch/action/support/broadcast/ChunkedBroadcastResponse.java
@@ -10,15 +10,15 @@
 import org.elasticsearch.action.support.DefaultShardOperationFailedException;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.rest.action.RestActions;
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.List;
 
-public abstract class ChunkedBroadcastResponse extends BaseBroadcastResponse implements ChunkedToXContent {
+public abstract class ChunkedBroadcastResponse extends BaseBroadcastResponse implements ChunkedToXContentObject {
     public ChunkedBroadcastResponse(StreamInput in) throws IOException {
         super(in);
     }

diff --git a/server/src/main/java/org/elasticsearch/health/HealthIndicatorResult.java b/server/src/main/java/org/elasticsearch/health/HealthIndicatorResult.java
@@ -9,7 +9,7 @@
 package org.elasticsearch.health;
 
 import org.elasticsearch.common.collect.Iterators;
-import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.util.Collections;
@@ -23,7 +23,7 @@ public record HealthIndicatorResult(
     HealthIndicatorDetails details,
     List<HealthIndicatorImpact> impacts,
     List<Diagnosis> diagnosisList
-) implements ChunkedToXContent {
+) implements ChunkedToXContentObject {
     @Override
     public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
         final Iterator<? extends ToXContent> diagnosisIterator;

diff --git a/...t/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java b/...t/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java
@@ -12,6 +12,7 @@
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceStats;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentFactory;
@@ -156,14 +157,10 @@ public void testToXContent() throws IOException {
         }
     }
 
-    public void testToChunkedXContent() {
-        DesiredBalanceResponse response = new DesiredBalanceResponse(randomStats(), randomRoutingTable());
-        var toXContentChunked = response.toXContentChunked(ToXContent.EMPTY_PARAMS);
-        int chunks = 0;
-        while (toXContentChunked.hasNext()) {
-            toXContentChunked.next();
-            chunks++;
-        }
-        assertEquals(response.getRoutingTable().size() + 2, chunks);
+    public void testChunking() {
+        AbstractChunkedSerializingTestCase.assertChunkCount(
+            new DesiredBalanceResponse(randomStats(), randomRoutingTable()),
+            response -> response.getRoutingTable().size() + 2
+        );
     }
 }
diff --git a/.../java/org/elasticsearch/action/admin/cluster/snapshots/get/GetSnapshotsResponseTests.java b/.../java/org/elasticsearch/action/admin/cluster/snapshots/get/GetSnapshotsResponseTests.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.snapshots.SnapshotInfo;
 import org.elasticsearch.snapshots.SnapshotInfoTestUtils;
 import org.elasticsearch.snapshots.SnapshotShardFailure;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentParser;
@@ -34,7 +35,6 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -43,7 +43,6 @@
 
 import static org.elasticsearch.snapshots.SnapshotInfo.INDEX_DETAILS_XCONTENT_PARAM;
 import static org.elasticsearch.test.AbstractXContentTestCase.chunkedXContentTester;
-import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
 import static org.hamcrest.CoreMatchers.containsString;
 
 public class GetSnapshotsResponseTests extends ESTestCase {
@@ -178,15 +177,8 @@ public void testFromXContent() throws IOException {
             .test();
     }
 
-    public void testToChunkedXContent() {
-        final GetSnapshotsResponse response = createTestInstance();
-        final Iterator<ToXContent> serialization = response.toXContentChunked(EMPTY_PARAMS);
-        int chunks = 0;
-        while (serialization.hasNext()) {
-            serialization.next();
-            chunks++;
-        }
-        assertEquals(chunks, response.getSnapshots().size() + 2);
+    public void testChunking() {
+        AbstractChunkedSerializingTestCase.assertChunkCount(createTestInstance(), response -> response.getSnapshots().size() + 2);
     }
 
 }
diff --git a/...org/elasticsearch/action/admin/cluster/snapshots/status/SnapshotsStatusResponseTests.java b/...org/elasticsearch/action/admin/cluster/snapshots/status/SnapshotsStatusResponseTests.java
@@ -17,8 +17,6 @@
 import java.util.List;
 import java.util.function.Predicate;
 
-import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
-
 public class SnapshotsStatusResponseTests extends AbstractChunkedSerializingTestCase<SnapshotsStatusResponse> {
 
     @Override
@@ -52,20 +50,10 @@ protected Writeable.Reader<SnapshotsStatusResponse> instanceReader() {
         return SnapshotsStatusResponse::new;
     }
 
-    public void testChunkCount() {
-        final var instance = createTestInstance();
-        // open and close chunk
-        int chunksExpected = 2;
-        for (SnapshotStatus snapshot : instance.getSnapshots()) {
-            // open and close chunk + one chunk per index
-            chunksExpected += 2 + snapshot.getIndices().size();
-        }
-        final var iterator = instance.toXContentChunked(EMPTY_PARAMS);
-        int chunksSeen = 0;
-        while (iterator.hasNext()) {
-            iterator.next();
-            chunksSeen++;
-        }
-        assertEquals(chunksExpected, chunksSeen);
+    public void testChunking() {
+        AbstractChunkedSerializingTestCase.assertChunkCount(
+            createTestInstance(),
+            instance -> 2 + instance.getSnapshots().stream().mapToInt(i -> 2 + i.getIndices().size()).sum()
+        );
     }
 }
diff --git a/...t/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponseTests.java b/...t/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponseTests.java
@@ -11,17 +11,13 @@
 import org.elasticsearch.cluster.service.PendingClusterTask;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.text.Text;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.xcontent.ToXContent;
-import org.elasticsearch.xcontent.XContentBuilder;
 
-import java.io.IOException;
 import java.util.ArrayList;
 
-import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
-
 public class PendingClusterTasksResponseTests extends ESTestCase {
-    public void testPendingClusterTasksResponseChunking() throws IOException {
+    public void testPendingClusterTasksResponseChunking() {
         final var tasks = new ArrayList<PendingClusterTask>();
         for (int i = between(0, 10); i > 0; i--) {
             tasks.add(
@@ -34,16 +30,9 @@ public void testPendingClusterTasksResponseChunking() throws IOException {
                 )
             );
         }
-
-        int chunkCount = 0;
-        try (XContentBuilder builder = jsonBuilder()) {
-            final var iterator = new PendingClusterTasksResponse(tasks).toXContentChunked(ToXContent.EMPTY_PARAMS);
-            while (iterator.hasNext()) {
-                iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS);
-                chunkCount += 1;
-            }
-        } // closing the builder verifies that the XContent is well-formed
-
-        assertEquals(tasks.size() + 2, chunkCount);
+        AbstractChunkedSerializingTestCase.assertChunkCount(
+            new PendingClusterTasksResponse(tasks),
+            response -> response.pendingTasks().size() + 2
+        );
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/get/GetIndexResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/get/GetIndexResponseTests.java
@@ -17,8 +17,8 @@
 import org.elasticsearch.common.settings.IndexScopedSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.RandomCreateIndexGenerator;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
-import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -29,9 +29,6 @@
 import java.util.Locale;
 import java.util.Map;
 
-import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
-import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
-
 public class GetIndexResponseTests extends AbstractWireSerializingTestCase<GetIndexResponse> {
 
     @Override
@@ -80,16 +77,6 @@ protected GetIndexResponse createTestInstance() {
     }
 
     public void testChunking() throws IOException {
-        final var response = createTestInstance();
-
-        try (var builder = jsonBuilder()) {
-            int chunkCount = 0;
-            final var iterator = response.toXContentChunked(EMPTY_PARAMS);
-            while (iterator.hasNext()) {
-                iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS);
-                chunkCount += 1;
-            }
-            assertEquals(response.getIndices().length + 2, chunkCount);
-        } // closing the builder verifies that the XContent is well-formed
+        AbstractChunkedSerializingTestCase.assertChunkCount(createTestInstance(), response -> response.getIndices().length + 2);
     }
 }
diff --git a/...est/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponseTests.java b/...est/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponseTests.java
@@ -11,6 +11,7 @@
 import org.elasticsearch.cluster.metadata.MappingMetadata;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.test.EqualsHashCodeTestUtils;
 
@@ -22,8 +23,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
-import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
-
 public class GetMappingsResponseTests extends AbstractWireSerializingTestCase<GetMappingsResponse> {
 
     public void testCheckEqualsAndHashCode() {
@@ -68,20 +67,15 @@ protected GetMappingsResponse createTestInstance() {
         return resp;
     }
 
-    public void testChunkedXContentUsesChunkPerIndex() {
-        final int indexCount = randomIntBetween(1, 10);
-        final var response = new GetMappingsResponse(
-            IntStream.range(0, indexCount)
-                .mapToObj(i -> "index-" + i)
-                .collect(Collectors.toUnmodifiableMap(Function.identity(), k -> createMappingsForIndex()))
+    public void testChunking() {
+        AbstractChunkedSerializingTestCase.assertChunkCount(
+            new GetMappingsResponse(
+                IntStream.range(0, randomIntBetween(1, 10))
+                    .mapToObj(i -> "index-" + i)
+                    .collect(Collectors.toUnmodifiableMap(Function.identity(), k -> createMappingsForIndex()))
+            ),
+            response -> response.mappings().size() + 2
         );
-        final var chunks = response.toXContentChunked(EMPTY_PARAMS);
-        int chunkCount = 0;
-        while (chunks.hasNext()) {
-            chunks.next();
-            chunkCount++;
-        }
-        assertEquals(2 + indexCount, chunkCount);
     }
 
     // Not meant to be exhaustive