From 7f97b60e45fa1a4e27d9623f40c6389d1c4c6ba9 Mon Sep 17 00:00:00 2001
From: "Daniel (dB.) Doubrovkine" <dblock@amazon.com>
Date: Tue, 7 May 2024 13:20:51 +0200
Subject: [PATCH] Fix: test. (#972)

Signed-off-by: dblock <dblock@amazon.com>
Signed-off-by: Vacha Shah <vachshah@amazon.com>
---
 .../integTest/AbstractSearchRequestIT.java      | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/java-client/src/test/java11/org/opensearch/client/opensearch/integTest/AbstractSearchRequestIT.java b/java-client/src/test/java11/org/opensearch/client/opensearch/integTest/AbstractSearchRequestIT.java
index efeb7dd7a2..495f568546 100644
--- a/java-client/src/test/java11/org/opensearch/client/opensearch/integTest/AbstractSearchRequestIT.java
+++ b/java-client/src/test/java11/org/opensearch/client/opensearch/integTest/AbstractSearchRequestIT.java
@@ -19,7 +19,9 @@
 import org.opensearch.Version;
 import org.opensearch.client.opensearch._types.FieldValue;
 import org.opensearch.client.opensearch._types.SortOrder;
+import org.opensearch.client.opensearch._types.aggregations.Aggregate;
 import org.opensearch.client.opensearch._types.aggregations.Aggregation;
+import org.opensearch.client.opensearch._types.aggregations.CompositeAggregate;
 import org.opensearch.client.opensearch._types.aggregations.CompositeAggregation;
 import org.opensearch.client.opensearch._types.aggregations.CompositeAggregationSource;
 import org.opensearch.client.opensearch._types.aggregations.CompositeBucket;
@@ -94,13 +96,14 @@ public void shouldReturnSearchResultsWithCompositeAgg() throws Exception {
 
         final SearchResponse<ShopItem> response = javaClient().search(request, ShopItem.class);
         assertEquals(response.hits().hits().size(), 2);
-        for (Map.Entry<String, Aggregation> entry : response.aggregations().entrySet()) {
-            CompositeAggregation compositeAggregation = entry.getValue().composite();
-            assertEquals(2, compositeAggregation.buckets().size());
-            assertEquals(1, Integer.parseInt(compositeAggregation.buckets().get(0).key().get("quantity").toString()));
-            assertEquals(1, compositeAggregation.buckets().get(0).docCount());
-            assertEquals(2, Integer.parseInt(compositeAggregation.buckets().get(1).key().get("quantity").toString()));
-            assertEquals(1, compositeAggregation.buckets().get(1).docCount());
+        for (Map.Entry<String, Aggregate> entry : response.aggregations().entrySet()) {
+            CompositeAggregate compositeAggregation = entry.getValue().composite();
+            List<CompositeBucket> buckets = compositeAggregation.buckets().array();
+            assertEquals(2, buckets.size());
+            assertEquals(1, Integer.parseInt(buckets.get(0).key().get("quantity").toString()));
+            assertEquals(1, buckets.get(0).docCount());
+            assertEquals(2, Integer.parseInt(buckets.get(1).key().get("quantity").toString()));
+            assertEquals(1, buckets.get(1).docCount());
         }
         List<CompositeBucket> buckets = response.aggregations()
             .entrySet()