[opensearch-project/geospatial#212] Fixing the IT for GeoTilesAggrega…

…tion. Signed-off-by: Navneet Verma <[email protected]>
navneet1v · Feb 1, 2023 · 57178fc · 57178fc
1 parent b1cf2d1
commit 57178fc
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 16 deletions.
diff --git a/.../org/opensearch/geo/search/aggregations/bucket/AbstractGeoBucketAggregationIntegTest.java b/.../org/opensearch/geo/search/aggregations/bucket/AbstractGeoBucketAggregationIntegTest.java
@@ -42,6 +42,8 @@ public abstract class AbstractGeoBucketAggregationIntegTest extends GeoModulePlu
 
     protected static final int MAX_PRECISION_FOR_GEO_SHAPES_AGG_TESTING = 4;
 
+    protected static final int MIN_PRECISION_WITHOUT_BB_AGGS = 2;
+
     protected static final int NUM_DOCS = 100;
 
     protected static final String GEO_SHAPE_INDEX_NAME = "geoshape_index";
@@ -88,14 +90,13 @@ protected void prepareGeoShapeIndexForAggregations(final Random random) throws E
             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;
-                }
+            isShapeIntersectingBB = geometryDocValue.isIntersectingRectangle(boundingRectangleForGeoShapesAgg);
+            if (!isShapeIntersectingBB && i == NUM_DOCS - 1) {
+                continue;
             }
+
             i++;
-            final Set<String> values = generateBucketsForGeometry(geometry, geometryDocValue);
+            final Set<String> values = generateBucketsForGeometry(geometry, geometryDocValue, isShapeIntersectingBB);
             geoshapes.add(indexGeoShape(GEO_SHAPE_INDEX_NAME, geometry));
             for (final String hash : values) {
                 expectedDocsCountForGeoShapes.put(hash, expectedDocsCountForGeoShapes.getOrDefault(hash, 0) + 1);
@@ -109,11 +110,16 @@ protected void prepareGeoShapeIndexForAggregations(final Random random) throws E
      * 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}
+     * @param geometry           {@link Geometry}
+     * @param geoShapeDocValue   {@link GeoShapeDocValue}
+     * @param intersectingWithBB boolean
      * @return A {@link Set} of {@link String} which represents the buckets.
      */
-    protected abstract Set<String> generateBucketsForGeometry(final Geometry geometry, final GeoShapeDocValue geoShapeDocValue);
+    protected abstract Set<String> generateBucketsForGeometry(
+        final Geometry geometry,
+        final GeoShapeDocValue geoShapeDocValue,
+        final boolean intersectingWithBB
+    );
 
     /**
      * Prepares a GeoPoint index for testing the GeoPoint bucket aggregations. Different bucket aggregations can use

diff --git a/...internalClusterTest/java/org/opensearch/geo/search/aggregations/bucket/GeoHashGridIT.java b/...internalClusterTest/java/org/opensearch/geo/search/aggregations/bucket/GeoHashGridIT.java
@@ -112,7 +112,7 @@ public void testGeoShapes() {
             GeoGridAggregationBuilder builder = AggregationBuilders.geohashGrid(AGG_NAME).field(GEO_SHAPE_FIELD_NAME).precision(precision);
             // This makes sure that for only higher precision we are providing the GeoBounding Box. This also ensures
             // that we are able to test both bounded and unbounded aggregations
-            if (precision > 2) {
+            if (precision > MIN_PRECISION_WITHOUT_BB_AGGS) {
                 builder.setGeoBoundingBox(boundingBox);
             }
             final SearchResponse response = client().prepareSearch(GEO_SHAPE_INDEX_NAME).addAggregation(builder).get();
@@ -268,14 +268,18 @@ public void testShardSizeIsZero() {
     }
 
     @Override
-    protected Set<String> generateBucketsForGeometry(final Geometry geometry, final GeoShapeDocValue geometryDocValue) {
+    protected Set<String> generateBucketsForGeometry(
+        final Geometry geometry,
+        final GeoShapeDocValue geometryDocValue,
+        boolean intersectingWithBB
+    ) {
         final GeoPoint topLeft = new GeoPoint();
         final GeoPoint bottomRight = new GeoPoint();
         assert geometry != null;
         GeoBoundsHelper.updateBoundsForGeometry(geometry, topLeft, bottomRight);
         final Set<String> geoHashes = new HashSet<>();
         for (int precision = MAX_PRECISION_FOR_GEO_SHAPES_AGG_TESTING; precision > 0; precision--) {
-            if (precision > 2 && !geometryDocValue.isIntersectingRectangle(boundingRectangleForGeoShapesAgg)) {
+            if (precision > MIN_PRECISION_WITHOUT_BB_AGGS && !intersectingWithBB) {
                 continue;
             }
             final GeoPoint topRight = new GeoPoint(topLeft.getLat(), bottomRight.getLon());

diff --git a/...internalClusterTest/java/org/opensearch/geo/search/aggregations/bucket/GeoTileGridIT.java b/...internalClusterTest/java/org/opensearch/geo/search/aggregations/bucket/GeoTileGridIT.java
@@ -60,7 +60,7 @@ public void testGeoShapes() {
                 .precision(precision);
             // This makes sure that for only higher precision we are providing the GeoBounding Box. This also ensures
             // that we are able to test both bounded and unbounded aggregations
-            if (precision > 2) {
+            if (precision > MIN_PRECISION_WITHOUT_BB_AGGS) {
                 builder.setGeoBoundingBox(boundingBox);
             }
             final SearchResponse response = client().prepareSearch(GEO_SHAPE_INDEX_NAME).addAggregation(builder).get();
@@ -134,18 +134,22 @@ public void testMultivaluedGeoPointsAggregation() throws Exception {
      * 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}
+     * @param geometry           {@link Geometry}
+     * @param geoShapeDocValue   {@link GeoShapeDocValue}
+     * @param intersectingWithBB
      * @return A {@link Set} of {@link String} which represents the buckets.
      */
     @Override
-    protected Set<String> generateBucketsForGeometry(Geometry geometry, GeoShapeDocValue geoShapeDocValue) {
+    protected Set<String> generateBucketsForGeometry(Geometry geometry, GeoShapeDocValue geoShapeDocValue, boolean intersectingWithBB) {
         final GeoPoint topLeft = new GeoPoint();
         final GeoPoint bottomRight = new GeoPoint();
         assert geometry != null;
         GeoBoundsHelper.updateBoundsForGeometry(geometry, topLeft, bottomRight);
         final Set<String> geoTiles = new HashSet<>();
         for (int precision = MAX_PRECISION_FOR_GEO_SHAPES_AGG_TESTING; precision > 0; precision--) {
+            if (precision > MIN_PRECISION_WITHOUT_BB_AGGS && !intersectingWithBB) {
+                continue;
+            }
             geoTiles.addAll(
                 GeoTileUtils.encodeShape(geoShapeDocValue, precision).stream().map(GeoTileUtils::stringEncode).collect(Collectors.toSet())
             );