Dry up testChunking

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 elastic#89838
DaveCTurner · Dec 22, 2022 · 6015782 · 6015782
1 parent 08515ea
commit 6015782
Show file tree

Hide file tree

Showing 24 changed files with 225 additions and 409 deletions.
diff --git a/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentHelper.java b/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentHelper.java
@@ -73,4 +73,8 @@ public static <T extends ToXContent> Iterator<ToXContent> wrapWithObject(String
     private static <T> Iterator<ToXContent> map(String name, Map<String, T> map, Function<Map.Entry<String, T>, ToXContent> toXContent) {
         return wrapWithObject(name, map.entrySet().stream().map(toXContent).iterator());
     }
+
+    public static ChunkedToXContent objectFromFragment(ChunkedToXContent fragment) {
+        return params -> Iterators.concat(startObject(), fragment.toXContentChunked(params), endObject());
+    }
 }
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.assertObjectChunkCount(
+            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 {
@@ -177,15 +176,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.assertObjectChunkCount(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.assertObjectChunkCount(
+            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.assertObjectChunkCount(
+            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.assertObjectChunkCount(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.assertObjectChunkCount(
+            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

diff --git a/.../src/test/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponseTests.java b/.../src/test/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponseTests.java
@@ -15,6 +15,7 @@
 import org.elasticsearch.cluster.routing.UnassignedInfo;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.recovery.RecoveryState;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.ESTestCase;
 
 import java.util.List;
@@ -23,7 +24,6 @@
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.emptySet;
-import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS;
 
 public class RecoveryResponseTests extends ESTestCase {
 
@@ -33,37 +33,33 @@ public void testChunkedToXContent() {
         DiscoveryNode sourceNode = new DiscoveryNode("foo", buildNewFakeTransportAddress(), emptyMap(), emptySet(), Version.CURRENT);
         DiscoveryNode targetNode = new DiscoveryNode("bar", buildNewFakeTransportAddress(), emptyMap(), emptySet(), Version.CURRENT);
         final int shards = randomInt(50);
-        final RecoveryResponse recoveryResponse = new RecoveryResponse(
-            successfulShards + failedShards,
-            successfulShards,
-            failedShards,
-            IntStream.range(0, shards)
-                .boxed()
-                .collect(
-                    Collectors.toUnmodifiableMap(
-                        i -> "index-" + i,
-                        i -> List.of(
-                            new RecoveryState(
-                                ShardRouting.newUnassigned(
-                                    new ShardId("index-" + i, "index-uuid-" + i, 0),
-                                    randomBoolean(),
-                                    RecoverySource.PeerRecoverySource.INSTANCE,
-                                    new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, null)
-                                ).initialize(sourceNode.getId(), null, randomNonNegativeLong()),
-                                sourceNode,
-                                targetNode
+        AbstractChunkedSerializingTestCase.assertObjectChunkCount(
+            new RecoveryResponse(
+                successfulShards + failedShards,
+                successfulShards,
+                failedShards,
+                IntStream.range(0, shards)
+                    .boxed()
+                    .collect(
+                        Collectors.toUnmodifiableMap(
+                            i -> "index-" + i,
+                            i -> List.of(
+                                new RecoveryState(
+                                    ShardRouting.newUnassigned(
+                                        new ShardId("index-" + i, "index-uuid-" + i, 0),
+                                        randomBoolean(),
+                                        RecoverySource.PeerRecoverySource.INSTANCE,
+                                        new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, null)
+                                    ).initialize(sourceNode.getId(), null, randomNonNegativeLong()),
+                                    sourceNode,
+                                    targetNode
+                                )
                             )
                         )
-                    )
-                ),
-            List.of()
+                    ),
+                List.of()
+            ),
+            ignored -> shards + 2
         );
-        final var iterator = recoveryResponse.toXContentChunked(EMPTY_PARAMS);
-        int chunks = 0;
-        while (iterator.hasNext()) {
-            iterator.next();
-            chunks++;
-        }
-        assertEquals(shards + 2, chunks);
     }
 }
diff --git a/...est/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponseTests.java b/...est/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponseTests.java
@@ -13,9 +13,10 @@
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.TestShardRouting;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.index.engine.Segment;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.util.ArrayList;
@@ -43,15 +44,12 @@ public void testToXContentSerialiationWithSortedFields() throws Exception {
             0,
             Collections.emptyList()
         );
-        var serialization = response.toXContentChunked(EMPTY_PARAMS);
         try (XContentBuilder builder = jsonBuilder()) {
-            while (serialization.hasNext()) {
-                serialization.next().toXContent(builder, ToXContent.EMPTY_PARAMS);
-            }
+            ChunkedToXContent.wrapAsXContentObject(response).toXContent(builder, EMPTY_PARAMS);
         }
     }
 
-    public void testSerializesOneChunkPerIndex() {
+    public void testChunking() {
         final int indices = randomIntBetween(1, 10);
         final List<ShardRouting> routings = new ArrayList<>(indices);
         for (int i = 0; i < indices; i++) {
@@ -61,19 +59,15 @@ public void testSerializesOneChunkPerIndex() {
         SortField sortField = new SortField("foo", SortField.Type.STRING);
         sortField.setMissingValue(SortField.STRING_LAST);
         segment.segmentSort = new Sort(sortField);
-        IndicesSegmentResponse response = new IndicesSegmentResponse(
-            routings.stream().map(routing -> new ShardSegments(routing, List.of(segment))).toArray(ShardSegments[]::new),
-            indices,
-            indices,
-            0,
-            Collections.emptyList()
+        AbstractChunkedSerializingTestCase.assertObjectChunkCount(
+            new IndicesSegmentResponse(
+                routings.stream().map(routing -> new ShardSegments(routing, List.of(segment))).toArray(ShardSegments[]::new),
+                indices,
+                indices,
+                0,
+                Collections.emptyList()
+            ),
+            response -> response.getIndices().size() + 4
         );
-        int chunks = 0;
-        final var iterator = response.toXContentChunked(EMPTY_PARAMS);
-        while (iterator.hasNext()) {
-            iterator.next();
-            chunks++;
-        }
-        assertEquals(indices + 4, chunks);
     }
 }