diff --git a/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentHelper.java b/server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentHelper.java index 19ebbe350a53f..29818dccf58b6 100644 --- 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 Iterator wrapWithObject(String private static Iterator map(String name, Map map, Function, 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/server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java index 51106c075bd5b..269e0b0faf3c2 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java +++ b/server/src/test/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/server/src/test/java/org/elasticsearch/action/admin/cluster/snapshots/get/GetSnapshotsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/snapshots/get/GetSnapshotsResponseTests.java index e730665fe5c8c..ed80e44b008ab 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/cluster/snapshots/get/GetSnapshotsResponseTests.java +++ b/server/src/test/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 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/server/src/test/java/org/elasticsearch/action/admin/cluster/snapshots/status/SnapshotsStatusResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/snapshots/status/SnapshotsStatusResponseTests.java index 26baa1a337eff..b30a3bdadf4b0 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/cluster/snapshots/status/SnapshotsStatusResponseTests.java +++ b/server/src/test/java/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 { @Override @@ -52,20 +50,10 @@ protected Writeable.Reader 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/server/src/test/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponseTests.java index 9e5c8bedf756e..179d7840a3a8a 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/cluster/tasks/PendingClusterTasksResponseTests.java +++ b/server/src/test/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(); 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 index 85ab886542a8f..e1392517f3f76 100644 --- 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 { @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/server/src/test/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponseTests.java index 0c58846483f50..ba310fa5b95bc 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/mapping/get/GetMappingsResponseTests.java +++ b/server/src/test/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 { 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/server/src/test/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponseTests.java index 53a2201ec5f5a..2203fd8f27e3c 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/recovery/RecoveryResponseTests.java +++ b/server/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/server/src/test/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponseTests.java index 012b24330ed3b..12ae20dfd3fbe 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponseTests.java +++ b/server/src/test/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 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); } } diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/settings/get/GetSettingsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/settings/get/GetSettingsResponseTests.java index 1fb3590c6a472..c66e7065de774 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/settings/get/GetSettingsResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/settings/get/GetSettingsResponseTests.java @@ -21,8 +21,6 @@ import java.util.Set; import java.util.function.Predicate; -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; - public class GetSettingsResponseTests extends AbstractChunkedSerializingTestCase { @Override @@ -76,15 +74,11 @@ protected Predicate getRandomFieldsExcludeFilter() { return f -> f.equals("") || f.contains(".settings") || f.contains(".defaults"); } - public void testOneChunkPerIndex() { - final var instance = createTestInstance(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - int chunks = 0; - while (iterator.hasNext()) { - chunks++; - iterator.next(); - } - assertEquals(2 + instance.getIndexToSettings().size(), chunks); + public void testChunking() { + AbstractChunkedSerializingTestCase.assertObjectChunkCount( + createTestInstance(), + response -> 2 + response.getIndexToSettings().size() + ); } } diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java index be98f601c5f5b..46e728de01926 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java @@ -14,18 +14,15 @@ import org.elasticsearch.common.util.Maps; import org.elasticsearch.index.search.stats.FieldUsageStats; import org.elasticsearch.index.shard.ShardId; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.xcontent.ToXContent; -import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xcontent.json.JsonXContent; -import java.io.IOException; import java.util.List; import java.util.Map; public class FieldUsageStatsResponseTests extends ESTestCase { - public void testToXContentChunkPerIndex() throws IOException { + public void testToXContentChunkPerIndex() { final int indices = randomIntBetween(0, 100); final Map> perIndex = Maps.newMapWithExpectedSize(indices); for (int i = 0; i < indices; i++) { @@ -46,15 +43,10 @@ public void testToXContentChunkPerIndex() throws IOException { ) ); } - final FieldUsageStatsResponse response = new FieldUsageStatsResponse(indices, indices, 0, List.of(), perIndex); - final XContentBuilder builder = JsonXContent.contentBuilder(); - final var iterator = response.toXContentChunked(ToXContent.EMPTY_PARAMS); - int chunks = 0; - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunks++; - } - assertEquals(indices + 2, chunks); + AbstractChunkedSerializingTestCase.assertObjectChunkCount( + new FieldUsageStatsResponse(indices, indices, 0, List.of(), perIndex), + ignored -> indices + 2 + ); } } diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/IndicesStatsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/IndicesStatsResponseTests.java index 4ee3408a7e1ed..4f812e152080f 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/IndicesStatsResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/IndicesStatsResponseTests.java @@ -13,14 +13,12 @@ import org.elasticsearch.cluster.routing.ShardRoutingState; import org.elasticsearch.cluster.routing.TestShardRouting; import org.elasticsearch.common.UUIDs; -import org.elasticsearch.common.io.Streams; import org.elasticsearch.index.Index; import org.elasticsearch.index.shard.ShardId; import org.elasticsearch.index.shard.ShardPath; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; -import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xcontent.json.JsonXContent; import java.io.IOException; @@ -137,24 +135,16 @@ public void testChunkedEncodingPerIndex() throws IOException { ClusterState.EMPTY_STATE.getMetadata(), ClusterState.EMPTY_STATE.routingTable() ); - final ToXContent.Params paramsClusterLevel = new ToXContent.MapParams(Map.of("level", "cluster")); - final var iteratorClusterLevel = indicesStatsResponse.toXContentChunked(paramsClusterLevel); - int chunksSeenClusterLevel = 0; - final XContentBuilder builder = new XContentBuilder(XContentType.JSON.xContent(), Streams.NULL_OUTPUT_STREAM); - while (iteratorClusterLevel.hasNext()) { - iteratorClusterLevel.next().toXContent(builder, paramsClusterLevel); - chunksSeenClusterLevel++; - } - assertEquals(3, chunksSeenClusterLevel); - - final ToXContent.Params paramsIndexLevel = new ToXContent.MapParams(Map.of("level", "indices")); - final var iteratorIndexLevel = indicesStatsResponse.toXContentChunked(paramsIndexLevel); - int chunksSeenIndexLevel = 0; - while (iteratorIndexLevel.hasNext()) { - iteratorIndexLevel.next().toXContent(builder, paramsIndexLevel); - chunksSeenIndexLevel++; - } - assertEquals(4 + shards, chunksSeenIndexLevel); + AbstractChunkedSerializingTestCase.assertObjectChunkCount( + indicesStatsResponse, + new ToXContent.MapParams(Map.of("level", "cluster")), + ignored -> 3 + ); + AbstractChunkedSerializingTestCase.assertObjectChunkCount( + indicesStatsResponse, + new ToXContent.MapParams(Map.of("level", "indices")), + ignored -> 4 + shards + ); } private ShardRouting createShardRouting(ShardId shardId, boolean isPrimary) { diff --git a/server/src/test/java/org/elasticsearch/action/fieldcaps/MergedFieldCapabilitiesResponseTests.java b/server/src/test/java/org/elasticsearch/action/fieldcaps/MergedFieldCapabilitiesResponseTests.java index 7954543c17940..ab3d4e8061bc4 100644 --- a/server/src/test/java/org/elasticsearch/action/fieldcaps/MergedFieldCapabilitiesResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/fieldcaps/MergedFieldCapabilitiesResponseTests.java @@ -21,8 +21,6 @@ import java.util.Map; import java.util.function.Predicate; -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; - public class MergedFieldCapabilitiesResponseTests extends AbstractChunkedSerializingTestCase { @Override @@ -206,30 +204,12 @@ private static FieldCapabilitiesResponse createSimpleResponse() { return new FieldCapabilitiesResponse(new String[] { "index1", "index2", "index3", "index4" }, responses, failureMap); } - public void testExpectedChunkSizes() { - { - final FieldCapabilitiesResponse instance = FieldCapabilitiesResponseTests.createResponseWithFailures(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - int chunks = 0; - while (iterator.hasNext()) { - iterator.next(); - chunks++; - } - if (instance.getFailures().isEmpty()) { - assertEquals(2, chunks); - } else { - assertEquals(3 + instance.get().size() + instance.getFailures().size(), chunks); - } - } - { - final FieldCapabilitiesResponse instance = createTestInstance(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - int chunks = 0; - while (iterator.hasNext()) { - iterator.next(); - chunks++; - } - assertEquals(2 + instance.get().size(), chunks); - } + public void testChunking() { + AbstractChunkedSerializingTestCase.assertObjectChunkCount( + FieldCapabilitiesResponseTests.createResponseWithFailures(), + instance -> instance.getFailures().isEmpty() ? 2 : (3 + instance.get().size() + instance.getFailures().size()) + ); + + AbstractChunkedSerializingTestCase.assertObjectChunkCount(createTestInstance(), instance -> 2 + instance.get().size()); } } diff --git a/server/src/test/java/org/elasticsearch/cluster/RepositoryCleanupInProgressTests.java b/server/src/test/java/org/elasticsearch/cluster/RepositoryCleanupInProgressTests.java index d001a7c5a2629..dd83cc1f53833 100644 --- a/server/src/test/java/org/elasticsearch/cluster/RepositoryCleanupInProgressTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/RepositoryCleanupInProgressTests.java @@ -8,31 +8,16 @@ package org.elasticsearch.cluster; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.xcontent.ToXContent; - -import java.io.IOException; - -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; -import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; public class RepositoryCleanupInProgressTests extends ESTestCase { - public void testChunking() throws IOException { - final var instance = new RepositoryCleanupInProgress( - randomList(10, () -> new RepositoryCleanupInProgress.Entry(randomAlphaOfLength(10), randomNonNegativeLong())) + public void testChunking() { + AbstractChunkedSerializingTestCase.assertFragmentChunkCount( + new RepositoryCleanupInProgress( + randomList(10, () -> new RepositoryCleanupInProgress.Entry(randomAlphaOfLength(10), randomNonNegativeLong())) + ), + i -> i.entries().size() + 2 ); - - int chunkCount = 0; - try (var builder = jsonBuilder()) { - builder.startObject(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunkCount += 1; - } - builder.endObject(); - } // closing the builder verifies that the XContent is well-formed - - assertEquals(instance.entries().size() + 2, chunkCount); } } diff --git a/server/src/test/java/org/elasticsearch/cluster/RestoreInProgressTests.java b/server/src/test/java/org/elasticsearch/cluster/RestoreInProgressTests.java index a9383fe08af1a..acfe5fddc0d53 100644 --- a/server/src/test/java/org/elasticsearch/cluster/RestoreInProgressTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/RestoreInProgressTests.java @@ -10,16 +10,13 @@ import org.elasticsearch.snapshots.Snapshot; import org.elasticsearch.snapshots.SnapshotId; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.xcontent.ToXContent; import java.io.IOException; import java.util.List; import java.util.Map; -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; -import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; - public class RestoreInProgressTests extends ESTestCase { public void testChunking() throws IOException { final var ripBuilder = new RestoreInProgress.Builder(); @@ -37,19 +34,6 @@ public void testChunking() throws IOException { ); } - final var instance = ripBuilder.build(); - - int chunkCount = 0; - try (var builder = jsonBuilder()) { - builder.startObject(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunkCount += 1; - } - builder.endObject(); - } // closing the builder verifies that the XContent is well-formed - - assertEquals(entryCount + 2, chunkCount); + AbstractChunkedSerializingTestCase.assertFragmentChunkCount(ripBuilder.build(), ignored -> entryCount + 2); } } diff --git a/server/src/test/java/org/elasticsearch/cluster/SnapshotDeletionsInProgressTests.java b/server/src/test/java/org/elasticsearch/cluster/SnapshotDeletionsInProgressTests.java index 8c18ee15afa78..c1edc444b6388 100644 --- a/server/src/test/java/org/elasticsearch/cluster/SnapshotDeletionsInProgressTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/SnapshotDeletionsInProgressTests.java @@ -11,6 +11,7 @@ import org.elasticsearch.common.Strings; import org.elasticsearch.common.xcontent.ChunkedToXContent; import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; @@ -19,7 +20,6 @@ import java.util.Collections; import java.util.List; -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; import static org.hamcrest.Matchers.equalTo; @@ -59,31 +59,21 @@ public void testXContent() throws IOException { } } - public void testChunking() throws IOException { - final var instance = SnapshotDeletionsInProgress.of( - randomList( - 10, - () -> new SnapshotDeletionsInProgress.Entry( - Collections.emptyList(), - randomAlphaOfLength(10), - randomNonNegativeLong(), - randomNonNegativeLong(), - randomFrom(SnapshotDeletionsInProgress.State.values()) + public void testChunking() { + AbstractChunkedSerializingTestCase.assertFragmentChunkCount( + SnapshotDeletionsInProgress.of( + randomList( + 10, + () -> new SnapshotDeletionsInProgress.Entry( + Collections.emptyList(), + randomAlphaOfLength(10), + randomNonNegativeLong(), + randomNonNegativeLong(), + randomFrom(SnapshotDeletionsInProgress.State.values()) + ) ) - ) + ), + instance -> instance.getEntries().size() + 2 ); - - int chunkCount = 0; - try (var builder = jsonBuilder()) { - builder.startObject(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunkCount += 1; - } - builder.endObject(); - } // closing the builder verifies that the XContent is well-formed - - assertEquals(instance.getEntries().size() + 2, chunkCount); } } diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexGraveyardTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexGraveyardTests.java index fa3a779c1a9dc..a40889f79d5d3 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexGraveyardTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexGraveyardTests.java @@ -14,9 +14,11 @@ import org.elasticsearch.common.collect.Iterators; import org.elasticsearch.common.io.stream.BytesStreamOutput; import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.xcontent.ChunkedToXContent; import org.elasticsearch.common.xcontent.ChunkedToXContentHelper; import org.elasticsearch.common.xcontent.XContentElasticsearchExtension; import org.elasticsearch.index.Index; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; @@ -59,13 +61,7 @@ public void testXContent() throws IOException { final IndexGraveyard graveyard = createRandom(); final XContentBuilder builder = JsonXContent.contentBuilder(); builder.startObject(); - final var iterator = graveyard.toXContentChunked(ToXContent.EMPTY_PARAMS); - int chunks = 0; - while (iterator.hasNext()) { - ++chunks; - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - } - assertEquals(2 + graveyard.getTombstones().size(), chunks); + ChunkedToXContent.wrapAsXContentObject(graveyard).toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); if (graveyard.getTombstones().size() > 0) { // check that date properly printed @@ -91,6 +87,10 @@ public void testXContent() throws IOException { assertThat(IndexGraveyard.fromXContent(parser), equalTo(graveyard)); } + public void testChunking() { + AbstractChunkedSerializingTestCase.assertFragmentChunkCount(createRandom(), graveyard -> graveyard.getTombstones().size() + 2); + } + public void testAddTombstones() { final IndexGraveyard graveyard1 = createRandom(); final IndexGraveyard.Builder graveyardBuidler = IndexGraveyard.builder(graveyard1); diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataTests.java index c9062e106aafb..517b8aedf1752 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataTests.java @@ -19,7 +19,6 @@ import org.elasticsearch.common.Strings; import org.elasticsearch.common.UUIDs; import org.elasticsearch.common.bytes.BytesReference; -import org.elasticsearch.common.io.Streams; import org.elasticsearch.common.io.stream.BytesStreamOutput; import org.elasticsearch.common.io.stream.NamedWriteableAwareStreamInput; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; @@ -36,12 +35,12 @@ import org.elasticsearch.index.alias.RandomAliasActionsGenerator; import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.plugins.MapperPlugin; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.VersionUtils; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentParser; -import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xcontent.json.JsonXContent; import java.io.IOException; @@ -2277,41 +2276,29 @@ public void testEmptyDiffReturnsSameInstance() throws IOException { public void testChunkedToXContent() throws IOException { final int datastreams = randomInt(10); - final Metadata instance = randomMetadata(datastreams); - int chunksSeen = 0; - try (XContentBuilder builder = new XContentBuilder(randomFrom(XContentType.values()), Streams.NULL_OUTPUT_STREAM, Set.of())) { - final var iterator = instance.toXContentChunked(ToXContent.EMPTY_PARAMS); - builder.startObject(); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunksSeen++; - } - builder.endObject(); - } - // 2 chunks at the beginning - // 1 chunk for each index + 2 to wrap the indices field - final int indicesChunks = instance.indices().size() + 2; - // 2 chunks for wrapping reserved state + 1 chunk for each item - final int reservedStateChunks = instance.reservedStateMetadata().size() + 2; - // 2 chunks wrapping templates and one chunk per template - final int templatesChunks = instance.templates().size() + 2; - // 2 chunks to wrap each custom - final int customChunks = 2 * instance.customs().size(); - // 1 chunk per datastream, 4 chunks to wrap ds and ds-aliases, or 0 if there are no datastreams - final int dsChunks = datastreams == 0 ? 0 : (datastreams + 4); - // 2 chunks to wrap index graveyard and one per tombstone - final int graveYardChunks = instance.indexGraveyard().getTombstones().size() + 2; - // 2 chunks to wrap component templates and one per component template - final int componentTemplateChunks = instance.componentTemplates().size() + 2; - // 2 chunks to wrap v2 templates and one per v2 template - final int v2TemplateChunks = instance.templatesV2().size() + 2; - // 1 chunk to close metadata - - assertEquals( - 2 + indicesChunks + reservedStateChunks + templatesChunks + customChunks + dsChunks + graveYardChunks + componentTemplateChunks - + v2TemplateChunks + 1, - chunksSeen - ); + AbstractChunkedSerializingTestCase.assertFragmentChunkCount(randomMetadata(datastreams), instance -> { + // 2 chunks at the beginning + // 1 chunk for each index + 2 to wrap the indices field + final int indicesChunks = instance.indices().size() + 2; + // 2 chunks for wrapping reserved state + 1 chunk for each item + final int reservedStateChunks = instance.reservedStateMetadata().size() + 2; + // 2 chunks wrapping templates and one chunk per template + final int templatesChunks = instance.templates().size() + 2; + // 2 chunks to wrap each custom + final int customChunks = 2 * instance.customs().size(); + // 1 chunk per datastream, 4 chunks to wrap ds and ds-aliases, or 0 if there are no datastreams + final int dsChunks = datastreams == 0 ? 0 : (datastreams + 4); + // 2 chunks to wrap index graveyard and one per tombstone + final int graveYardChunks = instance.indexGraveyard().getTombstones().size() + 2; + // 2 chunks to wrap component templates and one per component template + final int componentTemplateChunks = instance.componentTemplates().size() + 2; + // 2 chunks to wrap v2 templates and one per v2 template + final int v2TemplateChunks = instance.templatesV2().size() + 2; + // 1 chunk to close metadata + + return 2 + indicesChunks + reservedStateChunks + templatesChunks + customChunks + dsChunks + graveYardChunks + + componentTemplateChunks + v2TemplateChunks + 1; + }); } /** diff --git a/server/src/test/java/org/elasticsearch/health/HealthIndicatorResultTests.java b/server/src/test/java/org/elasticsearch/health/HealthIndicatorResultTests.java index 8825e7d6f1134..960ae3ae24886 100644 --- a/server/src/test/java/org/elasticsearch/health/HealthIndicatorResultTests.java +++ b/server/src/test/java/org/elasticsearch/health/HealthIndicatorResultTests.java @@ -10,6 +10,7 @@ import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; @@ -198,13 +199,6 @@ public void testChunkCount() { // -> each Diagnosis yields 5 chunks => 10 chunks from both diagnosis // -> HealthIndicatorResult surrounds the diagnosis list by 2 chunks - int chunksExpected = 12; - var iterator = result.toXContentChunked(ToXContent.EMPTY_PARAMS); - int chunksSeen = 0; - while (iterator.hasNext()) { - iterator.next(); - chunksSeen++; - } - assertEquals(chunksExpected, chunksSeen); + AbstractChunkedSerializingTestCase.assertObjectChunkCount(result, ignored -> 12); } } diff --git a/server/src/test/java/org/elasticsearch/health/metadata/HealthMetadataSerializationTests.java b/server/src/test/java/org/elasticsearch/health/metadata/HealthMetadataSerializationTests.java index 433a4a784efc9..975c83cac6519 100644 --- a/server/src/test/java/org/elasticsearch/health/metadata/HealthMetadataSerializationTests.java +++ b/server/src/test/java/org/elasticsearch/health/metadata/HealthMetadataSerializationTests.java @@ -15,15 +15,11 @@ import org.elasticsearch.common.unit.ByteSizeValue; import org.elasticsearch.common.unit.RatioValue; import org.elasticsearch.common.unit.RelativeByteSizeValue; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.SimpleDiffableWireSerializationTestCase; -import org.elasticsearch.xcontent.ToXContent; -import java.io.IOException; import java.util.List; -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; -import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; - public class HealthMetadataSerializationTests extends SimpleDiffableWireSerializationTestCase { @Override @@ -108,20 +104,7 @@ private HealthMetadata mutate(HealthMetadata base) { return new HealthMetadata(mutateDiskMetadata(base.getDiskMetadata())); } - public void testToXContentChunking() throws IOException { - final var instance = createTestInstance(); - - int chunkCount = 0; - try (var builder = jsonBuilder()) { - builder.startObject(); - final var iterator = instance.toXContentChunked(EMPTY_PARAMS); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunkCount += 1; - } - builder.endObject(); - } // closing the builder verifies that the XContent is well-formed - - assertEquals(1, chunkCount); + public void testChunking() { + AbstractChunkedSerializingTestCase.assertFragmentChunkCount(createTestInstance(), ignored -> 1); } } diff --git a/server/src/test/java/org/elasticsearch/ingest/IngestMetadataTests.java b/server/src/test/java/org/elasticsearch/ingest/IngestMetadataTests.java index 2f5690380eb8b..d97c0de6de2aa 100644 --- a/server/src/test/java/org/elasticsearch/ingest/IngestMetadataTests.java +++ b/server/src/test/java/org/elasticsearch/ingest/IngestMetadataTests.java @@ -12,9 +12,9 @@ import org.elasticsearch.cluster.metadata.Metadata; import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; -import org.elasticsearch.common.io.Streams; import org.elasticsearch.common.util.Maps; import org.elasticsearch.common.xcontent.ChunkedToXContent; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; @@ -25,7 +25,6 @@ import java.io.IOException; import java.util.HashMap; import java.util.Map; -import java.util.Set; import static org.hamcrest.Matchers.aMapWithSize; import static org.hamcrest.Matchers.allOf; @@ -125,7 +124,7 @@ public void testDiff() throws Exception { ); } - public void testChunkedToXContent() throws IOException { + public void testChunkedToXContent() { final BytesReference pipelineConfig = new BytesArray("{}"); final int pipelines = randomInt(10); final Map pipelineConfigurations = new HashMap<>(); @@ -133,16 +132,9 @@ public void testChunkedToXContent() throws IOException { final String id = Integer.toString(i); pipelineConfigurations.put(id, new PipelineConfiguration(id, pipelineConfig, XContentType.JSON)); } - int chunksSeen = 0; - try (XContentBuilder builder = new XContentBuilder(randomFrom(XContentType.values()), Streams.NULL_OUTPUT_STREAM, Set.of())) { - final var iterator = new IngestMetadata(pipelineConfigurations).toXContentChunked(ToXContent.EMPTY_PARAMS); - builder.startObject(); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunksSeen++; - } - builder.endObject(); - } - assertEquals(2 + pipelines, chunksSeen); + AbstractChunkedSerializingTestCase.assertFragmentChunkCount( + new IngestMetadata(pipelineConfigurations), + response -> 2 + response.getPipelines().size() + ); } } diff --git a/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java b/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java index 2842c593c6a7e..169d6633c89c5 100644 --- a/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java +++ b/server/src/test/java/org/elasticsearch/snapshots/SnapshotsInProgressSerializationTests.java @@ -20,17 +20,17 @@ import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.unit.ByteSizeValue; +import org.elasticsearch.common.xcontent.ChunkedToXContentHelper; import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.index.Index; import org.elasticsearch.index.shard.ShardId; import org.elasticsearch.repositories.IndexId; import org.elasticsearch.repositories.ShardGeneration; import org.elasticsearch.repositories.ShardSnapshotResult; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.SimpleDiffableWireSerializationTestCase; import org.elasticsearch.test.VersionUtils; -import org.elasticsearch.xcontent.ToXContent; -import org.elasticsearch.xcontent.XContentBuilder; import java.io.IOException; import java.util.ArrayList; @@ -42,8 +42,6 @@ import java.util.stream.Collectors; import java.util.stream.Stream; -import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; -import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; import static org.hamcrest.Matchers.anyOf; import static org.hamcrest.Matchers.equalTo; @@ -419,21 +417,8 @@ public void testXContent() throws IOException { ) ); - String json; - long chunkCount = 0; - try (XContentBuilder builder = jsonBuilder()) { - builder.humanReadable(true); - builder.startObject(); - final var iterator = sip.toXContentChunked(EMPTY_PARAMS); - while (iterator.hasNext()) { - iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); - chunkCount += 1; - } - builder.endObject(); - json = Strings.toString(builder); - } - - assertEquals(2 + sip.asStream().count(), chunkCount); + AbstractChunkedSerializingTestCase.assertFragmentChunkCount(sip, instance -> Math.toIntExact(instance.asStream().count() + 2)); + final var json = Strings.toString(ChunkedToXContentHelper.objectFromFragment(sip), false, true); assertThat( json, anyOf( diff --git a/server/src/test/java/org/elasticsearch/tasks/ListTasksResponseTests.java b/server/src/test/java/org/elasticsearch/tasks/ListTasksResponseTests.java index 990f6d1a62025..cec9009e12aa5 100644 --- a/server/src/test/java/org/elasticsearch/tasks/ListTasksResponseTests.java +++ b/server/src/test/java/org/elasticsearch/tasks/ListTasksResponseTests.java @@ -15,8 +15,8 @@ import org.elasticsearch.cluster.node.DiscoveryNodes; import org.elasticsearch.common.Strings; import org.elasticsearch.common.xcontent.ChunkedToXContent; +import org.elasticsearch.test.AbstractChunkedSerializingTestCase; import org.elasticsearch.test.AbstractXContentTestCase; -import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.ToXContentObject; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentParser; @@ -32,7 +32,6 @@ import static java.util.Collections.emptyList; import static java.util.Collections.singletonList; import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; -import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.instanceOf; @@ -185,38 +184,14 @@ public void testFromXContentWithFailures() throws IOException { ); } - public void testChunkedEncoding() throws IOException { + public void testChunkedEncoding() { final var response = createTestInstanceWithFailures().in(); - - try (var builder = jsonBuilder()) { - int chunkCount = 0; - final var iterator = response.groupedByNone().toXContentChunked(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(response.getTasks().size() + 2, chunkCount); - } - - try (var builder = jsonBuilder()) { - int chunkCount = 0; - final var iterator = response.groupedByParent().toXContentChunked(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(response.getTaskGroups().size() + 2, chunkCount); - } - - try (var builder = jsonBuilder()) { - int chunkCount = 0; - final var iterator = response.groupedByNode(() -> DiscoveryNodes.EMPTY_NODES).toXContentChunked(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(2 + response.getPerNodeTasks().values().stream().mapToInt(entry -> 2 + entry.size()).sum(), chunkCount); - } + AbstractChunkedSerializingTestCase.assertObjectChunkCount(response.groupedByNone(), o -> response.getTasks().size() + 2); + AbstractChunkedSerializingTestCase.assertObjectChunkCount(response.groupedByParent(), o -> response.getTaskGroups().size() + 2); + AbstractChunkedSerializingTestCase.assertObjectChunkCount( + response.groupedByNode(() -> DiscoveryNodes.EMPTY_NODES), + o -> 2 + response.getPerNodeTasks().values().stream().mapToInt(entry -> 2 + entry.size()).sum() + ); } private static ListTasksResponseWrapper createTestInstanceWithFailures() { diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractChunkedSerializingTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractChunkedSerializingTestCase.java index 39f8ba0eab12c..4cd2633b4ff3b 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/AbstractChunkedSerializingTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/test/AbstractChunkedSerializingTestCase.java @@ -16,8 +16,11 @@ import org.elasticsearch.xcontent.XContentType; import java.io.IOException; +import java.util.function.ToIntFunction; import static org.elasticsearch.test.AbstractXContentTestCase.chunkedXContentTester; +import static org.elasticsearch.xcontent.ToXContent.EMPTY_PARAMS; +import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; public abstract class AbstractChunkedSerializingTestCase extends AbstractSerializationTestCase { @@ -57,4 +60,45 @@ protected T createXContextTestInstance(XContentType xContentType) { protected boolean isFragment() { return false; } + + public static void assertObjectChunkCount(T instance, ToIntFunction expectedChunkCount) { + assertObjectChunkCount(instance, EMPTY_PARAMS, expectedChunkCount); + } + + public static void assertObjectChunkCount( + T instance, + ToXContent.Params params, + ToIntFunction expectedChunkCount + ) { + assertChunkCount(instance, false, params, expectedChunkCount); + } + + public static void assertFragmentChunkCount(T instance, ToIntFunction expectedChunkCount) { + assertChunkCount(instance, true, EMPTY_PARAMS, expectedChunkCount); + } + + private static void assertChunkCount( + T instance, + boolean fragment, + ToXContent.Params params, + ToIntFunction expectedChunkCount + ) { + int chunkCount = 0; + try (var builder = jsonBuilder()) { + if (fragment) { + builder.startObject(); + } + final var iterator = instance.toXContentChunked(params); + while (iterator.hasNext()) { + iterator.next().toXContent(builder, params); + chunkCount += 1; + } + if (fragment) { + builder.endObject(); + } + } catch (IOException e) { + throw new AssertionError("unexpected", e); + } // closing the builder verifies that the XContent is well-formed + assertEquals(expectedChunkCount.applyAsInt(instance), chunkCount); + } }