Refractor the Integration tests to abstract the geoshape and geo-poin…

…t index preparation code.

Signed-off-by: Navneet Verma <[email protected]>

modules/geo/src/internalClusterTest/java/org/opensearch/geo/search/aggregations/bucket/AbstractBucketAggregationIntegTest.java

@@ -8,17 +8,30 @@
 
 package org.opensearch.geo.search.aggregations.bucket;
 
+import com.carrotsearch.hppc.ObjectIntHashMap;
+import com.carrotsearch.hppc.ObjectIntMap;
 import org.opensearch.Version;
 import org.opensearch.action.index.IndexRequestBuilder;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.common.geo.GeoPoint;
+import org.opensearch.common.geo.GeoShapeDocValue;
+import org.opensearch.common.settings.Settings;
 import org.opensearch.common.xcontent.XContentBuilder;
 import org.opensearch.geo.GeoModulePluginIntegTestCase;
+import org.opensearch.geo.tests.common.RandomGeoGenerator;
+import org.opensearch.geo.tests.common.RandomGeoGeometryGenerator;
 import org.opensearch.geometry.Geometry;
 import org.opensearch.geometry.Rectangle;
 import org.opensearch.test.VersionUtils;
 
+import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Random;
+import java.util.Set;
 
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
+import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 
 /**
  * This is the base class for all the Bucket Aggregation related integration tests. Use this class to add common
@@ -33,33 +46,169 @@ public abstract class AbstractBucketAggregationIntegTest extends GeoModulePlugin
 
     protected static final String GEO_SHAPE_INDEX_NAME = "geoshape_index";
 
-    // We don't want to generate this BB at random as it may lead to generation of very large or very small BB.
-    // Hence, we are making the parameters of this BB static for simplicity.
-    protected static final Rectangle BOUNDING_RECTANGLE_FOR_GEO_SHAPES_AGG = new Rectangle(-4.1, 20.9, 21.9, -3.1);
+    protected static Rectangle boundingRectangleForGeoShapesAgg;
 
-    protected static final String AGG_NAME = "geohashgrid";
+    protected static ObjectIntMap<String> expectedDocsCountForGeoShapes;
+
+    protected static ObjectIntMap<String> expectedDocCountsForSingleGeoPoint;
+
+    protected static ObjectIntMap<String> multiValuedExpectedDocCountsGeoPoint;
 
     protected static final String GEO_SHAPE_FIELD_NAME = "location_geo_shape";
 
     protected static final String GEO_POINT_FIELD_NAME = "location";
 
     protected static final String KEYWORD_FIELD_NAME = "city";
 
+    protected static String smallestGeoHash = null;
+
     protected final Version version = VersionUtils.randomIndexCompatibleVersion(random());
 
     @Override
     protected boolean forbidPrivateIndexSettings() {
         return false;
     }
 
+    /**
+     * Prepares a GeoShape index for testing the GeoShape bucket aggregations. Different bucket aggregations can use
+     * different techniques for creating buckets. Override the method
+     * {@link AbstractBucketAggregationIntegTest#generateBucketsForGeometry} in the test class for creating the
+     * buckets which will then be used for verifications.
+     *
+     * @param random {@link Random}
+     * @throws Exception thrown during index creation.
+     */
+    protected void prepareGeoShapeIndexForAggregations(final Random random) throws Exception {
+        expectedDocsCountForGeoShapes = new ObjectIntHashMap<>();
+        final Settings settings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, version).build();
+        final List<IndexRequestBuilder> geoshapes = new ArrayList<>();
+        assertAcked(prepareCreate(GEO_SHAPE_INDEX_NAME).setSettings(settings).setMapping(GEO_SHAPE_FIELD_NAME, "type" + "=geo_shape"));
+        boolean isShapeIntersectingBB = false;
+        for (int i = 0; i < NUM_DOCS;) {
+            final Geometry geometry = RandomGeoGeometryGenerator.randomGeometry(random);
+            final GeoShapeDocValue geometryDocValue = GeoShapeDocValue.createGeometryDocValue(geometry);
+            // make sure that there is 1 shape is intersecting with the bounding box
+            if (!isShapeIntersectingBB) {
+                isShapeIntersectingBB = geometryDocValue.isIntersectingRectangle(boundingRectangleForGeoShapesAgg);
+                if (!isShapeIntersectingBB && i == NUM_DOCS - 1) {
+                    continue;
+                }
+            }
+            i++;
+            final Set<String> values = generateBucketsForGeometry(geometry, geometryDocValue);
+            geoshapes.add(indexGeoShape(GEO_SHAPE_INDEX_NAME, geometry));
+            for (final String hash : values) {
+                expectedDocsCountForGeoShapes.put(hash, expectedDocsCountForGeoShapes.getOrDefault(hash, 0) + 1);
+            }
+        }
+        indexRandom(true, geoshapes);
+        ensureGreen(GEO_SHAPE_INDEX_NAME);
+    }
+
+    /**
+     * Returns a set of buckets for the shape at different precision level. Override this method for different bucket
+     * aggregations.
+     *
+     * @param geometry {@link Geometry}
+     * @param geoShapeDocValue {@link GeoShapeDocValue}
+     * @return A {@link Set} of {@link String} which represents the buckets.
+     */
+    protected abstract Set<String> generateBucketsForGeometry(final Geometry geometry, final GeoShapeDocValue geoShapeDocValue);
+
+    /**
+     * Prepares a GeoPoint index for testing the GeoPoint bucket aggregations. Different bucket aggregations can use
+     * different techniques for creating buckets. Override the method
+     * {@link AbstractBucketAggregationIntegTest#generateBucketsForGeoPoint} in the test class for creating the
+     * buckets which will then be used for verifications.
+     *
+     * @param random {@link Random}
+     * @throws Exception thrown during index creation.
+     */
+    protected void prepareSingleValueGeoPointIndex(final Random random) throws Exception {
+        expectedDocCountsForSingleGeoPoint = new ObjectIntHashMap<>();
+        createIndex("idx_unmapped");
+        final Settings settings = Settings.builder()
+            .put(IndexMetadata.SETTING_VERSION_CREATED, version)
+            .put("index.number_of_shards", 4)
+            .put("index.number_of_replicas", 0)
+            .build();
+        assertAcked(
+            prepareCreate("idx").setSettings(settings)
+                .setMapping(GEO_POINT_FIELD_NAME, "type=geo_point", KEYWORD_FIELD_NAME, "type=keyword")
+        );
+        final List<IndexRequestBuilder> cities = new ArrayList<>();
+        for (int i = 0; i < NUM_DOCS; i++) {
+            // generate random point
+            final GeoPoint geoPoint = RandomGeoGenerator.randomPoint(random);
+            cities.add(indexGeoPoint("idx", geoPoint.toString(), geoPoint.getLat() + ", " + geoPoint.getLon()));
+            final Set<String> buckets = generateBucketsForGeoPoint(geoPoint);
+            for (final String bucket : buckets) {
+                expectedDocCountsForSingleGeoPoint.put(bucket, expectedDocCountsForSingleGeoPoint.getOrDefault(bucket, 0) + 1);
+            }
+        }
+        indexRandom(true, cities);
+        ensureGreen("idx_unmapped", "idx");
+    }
+
+    protected void prepareMultiValuedGeoPointIndex(final Random random) throws Exception {
+        multiValuedExpectedDocCountsGeoPoint = new ObjectIntHashMap<>();
+        final Settings settings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, version).build();
+        final List<IndexRequestBuilder> cities = new ArrayList<>();
+        assertAcked(
+            prepareCreate("multi_valued_idx").setSettings(settings)
+                .setMapping(GEO_POINT_FIELD_NAME, "type=geo_point", KEYWORD_FIELD_NAME, "type=keyword")
+        );
+        for (int i = 0; i < NUM_DOCS; i++) {
+            final int numPoints = random.nextInt(4);
+            final List<String> points = new ArrayList<>();
+            final Set<String> buckets = new HashSet<>();
+            for (int j = 0; j < numPoints; ++j) {
+                // generate random point
+                final GeoPoint geoPoint = RandomGeoGenerator.randomPoint(random);
+                points.add(geoPoint.getLat() + "," + geoPoint.getLon());
+                buckets.addAll(generateBucketsForGeoPoint(geoPoint));
+            }
+            cities.add(indexGeoPoints("multi_valued_idx", Integer.toString(i), points));
+            for (final String bucket : buckets) {
+                multiValuedExpectedDocCountsGeoPoint.put(bucket, multiValuedExpectedDocCountsGeoPoint.getOrDefault(bucket, 0) + 1);
+            }
+        }
+        indexRandom(true, cities);
+        ensureGreen("multi_valued_idx");
+    }
+
+    /**
+     * Returns a set of buckets for the GeoPoint at different precision level. Override this method for different bucket
+     * aggregations.
+     *
+     * @param geoPoint {@link GeoPoint}
+     * @return A {@link Set} of {@link String} which represents the buckets.
+     */
+    protected abstract Set<String> generateBucketsForGeoPoint(final GeoPoint geoPoint);
+
+    /**
+     * Indexes a GeoShape in the provided index.
+     * @param index {@link String} index name
+     * @param geometry {@link Geometry} the Geometry to be indexed
+     * @return {@link IndexRequestBuilder}
+     * @throws Exception thrown during creation of {@link IndexRequestBuilder}
+     */
     protected IndexRequestBuilder indexGeoShape(final String index, final Geometry geometry) throws Exception {
         XContentBuilder source = jsonBuilder().startObject();
         source = source.field(GEO_SHAPE_FIELD_NAME, WKT.toWKT(geometry));
         source = source.endObject();
         return client().prepareIndex(index).setSource(source);
     }
 
-    protected IndexRequestBuilder indexCity(final String index, final String name, final List<String> latLon) throws Exception {
+    /**
+     * Indexes a {@link List} of {@link GeoPoint}s in the provided Index name.
+     * @param index {@link String} index name
+     * @param name {@link String} value for the string field in index
+     * @param latLon {@link List} of {@link String} representing the String representation of GeoPoint
+     * @return {@link IndexRequestBuilder}
+     * @throws Exception thrown during indexing.
+     */
+    protected IndexRequestBuilder indexGeoPoints(final String index, final String name, final List<String> latLon) throws Exception {
         XContentBuilder source = jsonBuilder().startObject().field(KEYWORD_FIELD_NAME, name);
         if (latLon != null) {
             source = source.field(GEO_POINT_FIELD_NAME, latLon);
@@ -68,8 +217,38 @@ protected IndexRequestBuilder indexCity(final String index, final String name, f
         return client().prepareIndex(index).setSource(source);
     }
 
-    protected IndexRequestBuilder indexCity(final String index, final String name, final String latLon) throws Exception {
-        return indexCity(index, name, List.of(latLon));
+    /**
+     * Indexes a {@link GeoPoint} in the provided Index name.
+     * @param index {@link String} index name
+     * @param name {@link String} value for the string field in index
+     * @param latLon {@link String} representing the String representation of GeoPoint
+     * @return {@link IndexRequestBuilder}
+     * @throws Exception thrown during indexing.
+     */
+    protected IndexRequestBuilder indexGeoPoint(final String index, final String name, final String latLon) throws Exception {
+        return indexGeoPoints(index, name, List.of(latLon));
+    }
+
+    /**
+     * Generates a Bounding Box of a fixed radius that can be used for shapes aggregations to reduce the size of
+     * aggregation results.
+     * @param random {@link Random}
+     * @return {@link Rectangle}
+     */
+    protected Rectangle getGridAggregationBoundingBox(final Random random) {
+        final double radius = getRadiusOfBoundingBox();
+        assertTrue("The radius of Bounding Box is less than or equal to 0", radius > 0);
+        return RandomGeoGeometryGenerator.randomRectangle(random, radius);
+    }
+
+    /**
+     * Returns a radius for the Bounding box. Test classes can override this method to change the radius of BBox for
+     * the test cases. If we increase this value, it will lead to creation of a lot of buckets that can lead of
+     * IndexOutOfBoundsExceptions.
+     * @return double
+     */
+    protected double getRadiusOfBoundingBox() {
+        return 5.0;
     }
 
 }