Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add integration test for geohex_grid #176

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
import static org.opensearch.geospatial.action.upload.geojson.UploadGeoJSONRequestContent.FIELD_DATA;
import static org.opensearch.geospatial.shared.URLBuilder.getPluginURLPrefix;
import static org.opensearch.index.query.AbstractGeometryQueryBuilder.DEFAULT_SHAPE_FIELD_NAME;
import static org.opensearch.rest.action.search.RestSearchAction.TYPED_KEYS_PARAM;
import static org.opensearch.search.aggregations.Aggregations.AGGREGATIONS_FIELD;

import java.io.IOException;
import java.util.Collections;
Expand Down Expand Up @@ -69,6 +71,7 @@ public abstract class GeospatialRestTestCase extends OpenSearchSecureRestTestCas
public static final String SHAPE_INDEX_FIELD = "index";
public static final String SHAPE_ID_FIELD = "id";
public static final String SHAPE_INDEX_PATH_FIELD = "path";
public static final String QUERY_PARAM_TOKEN = "?";

private static String buildPipelinePath(String name) {
return String.join(URL_DELIMITER, "_ingest", "pipeline", name);
Expand Down Expand Up @@ -251,7 +254,16 @@ public String buildContentAsString(CheckedConsumer<XContentBuilder, IOException>
return Strings.toString(builder);
}

public String buildSearchBodyAsString(
public String buildSearchAggregationsBodyAsString(CheckedConsumer<XContentBuilder, IOException> aggregationsBuilder)
throws IOException {
return buildContentAsString(builder -> {
builder.startObject(AGGREGATIONS_FIELD);
aggregationsBuilder.accept(builder);
builder.endObject();
});
}

public String buildSearchQueryBodyAsString(
CheckedConsumer<XContentBuilder, IOException> searchQueryBuilder,
String queryType,
String fieldName
Expand All @@ -264,9 +276,12 @@ public String buildSearchBodyAsString(
});
}

public SearchResponse searchIndex(String indexName, String entity) throws Exception {
String path = String.join(URL_DELIMITER, indexName, SEARCH);
final Request request = new Request("GET", path);
public SearchResponse searchIndex(String indexName, String entity, boolean includeType) throws Exception {
var urlPathBuilder = new StringBuilder().append(indexName).append(URL_DELIMITER).append(SEARCH);
if (includeType) {
urlPathBuilder.append(QUERY_PARAM_TOKEN).append(TYPED_KEYS_PARAM);
}
final Request request = new Request("GET", urlPathBuilder.toString());
request.setJsonEntity(entity);
final Response response = client().performRequest(request);
return SearchResponse.fromXContent(createParser(XContentType.JSON.xContent(), EntityUtils.toString(response.getEntity())));
Expand Down Expand Up @@ -299,13 +314,13 @@ public String indexDocumentUsingGeoJSON(String indexName, String fieldName, Geom

public SearchResponse searchUsingShapeRelation(String indexName, String fieldName, Geometry geometry, ShapeRelation shapeRelation)
throws Exception {
String searchEntity = buildSearchBodyAsString(builder -> {
String searchEntity = buildSearchQueryBodyAsString(builder -> {
builder.field(DEFAULT_SHAPE_FIELD_NAME);
GeoJson.toXContent(geometry, builder, EMPTY_PARAMS);
builder.field(SHAPE_RELATION, shapeRelation.getRelationName());
}, XYShapeQueryBuilder.NAME, fieldName);

return searchIndex(indexName, searchEntity);
return searchIndex(indexName, searchEntity, false);
}

public void createIndexedShapeIndex() throws IOException {
Expand All @@ -321,15 +336,15 @@ public SearchResponse searchUsingIndexedShapeIndex(
String docId,
String fieldName
) throws Exception {
String searchEntity = buildSearchBodyAsString(builder -> {
String searchEntity = buildSearchQueryBodyAsString(builder -> {
builder.startObject(INDEXED_SHAPE_FIELD);
builder.field(SHAPE_INDEX_FIELD, indexedShapeIndex);
builder.field(SHAPE_ID_FIELD, docId);
builder.field(SHAPE_INDEX_PATH_FIELD, indexedShapePath);
builder.endObject();
}, XYShapeQueryBuilder.NAME, fieldName);

return searchIndex(indexName, searchEntity);
return searchIndex(indexName, searchEntity, false);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -56,12 +56,12 @@ public void testIndexPointsFilterRectangleWithDefaultRelation() throws Exception
indexDocumentUsingGeoJSON(indexName, xyShapeFieldName, new Point(-45, -50));

Rectangle rectangle = new Rectangle(-45, 45, 45, -45);
String searchEntity = buildSearchBodyAsString(builder -> {
String searchEntity = buildSearchQueryBodyAsString(builder -> {
builder.field(DEFAULT_SHAPE_FIELD_NAME);
GeoJson.toXContent(rectangle, builder, EMPTY_PARAMS);
}, XYShapeQueryBuilder.NAME, xyShapeFieldName);

final SearchResponse searchResponse = searchIndex(indexName, searchEntity);
final SearchResponse searchResponse = searchIndex(indexName, searchEntity, false);
assertSearchResponse(searchResponse);
assertHitCount(searchResponse, 1);
MatcherAssert.assertThat(searchResponse.getHits().getAt(0).getId(), equalTo(firstDocumentID));
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,172 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.geospatial.search.aggregations.bucket.geogrid;

import static java.util.function.Function.identity;
import static java.util.stream.Collectors.toMap;
import static org.hamcrest.Matchers.containsString;
import static org.opensearch.geo.GeometryTestUtils.randomPoint;
import static org.opensearch.geospatial.GeospatialTestHelper.randomHexGridPrecision;
import static org.opensearch.geospatial.GeospatialTestHelper.randomLowerCaseString;
import static org.opensearch.geospatial.h3.H3.geoToH3Address;
import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertSearchResponse;

import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import org.hamcrest.MatcherAssert;
import org.opensearch.client.ResponseException;
import org.opensearch.cluster.ClusterModule;
import org.opensearch.common.ParseField;
import org.opensearch.common.settings.Settings;
import org.opensearch.common.xcontent.ContextParser;
import org.opensearch.common.xcontent.NamedXContentRegistry;
import org.opensearch.geometry.Point;
import org.opensearch.geospatial.GeospatialRestTestCase;
import org.opensearch.geospatial.h3.H3;
import org.opensearch.index.mapper.GeoPointFieldMapper;
import org.opensearch.search.aggregations.Aggregation;
import org.opensearch.search.aggregations.bucket.MultiBucketsAggregation;

public class GeoHexAggregationIT extends GeospatialRestTestCase {

private static final String FIELD = "field";
private static final String FIELD_PRECISION = "precision";
private static final String FIELD_SIZE = "size";
private static int MAX_DOCUMENTS = 15;
private static int MIN_DOCUMENTS = 2;
private String indexName;
private String geospatialFieldName;

@Override
public void setUp() throws Exception {
super.setUp();
indexName = randomLowerCaseString();
geospatialFieldName = randomLowerCaseString();
}

public void testGeoHexGridBucket() throws Exception {
// Step 1: Create an index
createIndex(indexName, Settings.EMPTY, Map.of(geospatialFieldName, GeoPointFieldMapper.CONTENT_TYPE));

// Generate metadata for Test data
final var randomDocumentsForTesting = randomIntBetween(MIN_DOCUMENTS, MAX_DOCUMENTS);
final var randomPrecision = randomHexGridPrecision();

// Generate Test data
final Map<Point, String> pointStringMap = generateRandomPointH3CellMap(randomDocumentsForTesting, randomPrecision);
for (var point : pointStringMap.keySet()) {
indexDocumentUsingWKT(indexName, geospatialFieldName, point.toString());
}

// do in-memory aggregation for comparison
final Map<String, Long> expectedAggregationMap = pointStringMap.values()
.stream()
.collect(Collectors.groupingBy(identity(), Collectors.counting()));

// build test aggregation search query
var context = randomLowerCaseString();
var content = buildSearchAggregationsBodyAsString(builder -> {
builder.startObject(context)
.startObject(GeoHexGridAggregationBuilder.NAME)
.field(FIELD, geospatialFieldName)
.field(FIELD_PRECISION, randomPrecision)
.field(FIELD_SIZE, expectedAggregationMap.size())
.endObject()
.endObject();
});

// execute Aggregation
final var searchResponse = searchIndex(indexName, content, true);
// Assert Search succeeded
assertSearchResponse(searchResponse);
// Fetch Aggregation
final var aggregation = searchResponse.getAggregations().asMap().get(context);
assertNotNull(aggregation);

// Assert Results
assertTrue(aggregation instanceof MultiBucketsAggregation);
final var multiBucketsAggregation = (MultiBucketsAggregation) aggregation;

// Assert size before checking contents
assertEquals(expectedAggregationMap.size(), multiBucketsAggregation.getBuckets().size());
final Map<String, Long> actualAggregationMap = multiBucketsAggregation.getBuckets()
.stream()
.collect(toMap(MultiBucketsAggregation.Bucket::getKeyAsString, MultiBucketsAggregation.Bucket::getDocCount));

// compare in-memory aggregation with cluster aggregation
assertEquals(expectedAggregationMap, actualAggregationMap);

}

public void testSizeIsZero() throws Exception {

// build test aggregation search query
var context = randomLowerCaseString();
var content = buildSearchAggregationsBodyAsString(builder -> {
builder.startObject(context)
.startObject(GeoHexGridAggregationBuilder.NAME)
.field(FIELD, geospatialFieldName)
.field(FIELD_PRECISION, randomHexGridPrecision())
.field(FIELD_SIZE, 0)
.endObject()
.endObject();
});

// execute Aggregation
ResponseException exception = expectThrows(ResponseException.class, () -> searchIndex(indexName, content, true));
MatcherAssert.assertThat(exception.getMessage(), containsString("[size] must be greater than 0."));
}

public void testInvalidPrecision() throws Exception {

// build test aggregation search query
var invalidPrecision = H3.MAX_H3_RES + 1;
var content = buildSearchAggregationsBodyAsString(builder -> {
builder.startObject(randomLowerCaseString())
.startObject(GeoHexGridAggregationBuilder.NAME)
.field(FIELD, geospatialFieldName)
.field(FIELD_PRECISION, invalidPrecision)
.endObject()
.endObject();
});

// execute Aggregation
ResponseException exception = expectThrows(ResponseException.class, () -> searchIndex(indexName, content, true));
MatcherAssert.assertThat(
exception.getMessage(),
containsString(
String.format(
Locale.getDefault(),
"Invalid precision of %d . Must be between %d and %d",
invalidPrecision,
H3.MIN_H3_RES,
H3.MAX_H3_RES
)
)
);
}

@Override
protected NamedXContentRegistry xContentRegistry() {
final List<NamedXContentRegistry.Entry> namedXContents = new ArrayList<>(ClusterModule.getNamedXWriteables());
final ContextParser<Object, Aggregation> hexGridParser = (p, c) -> ParsedGeoHexGrid.fromXContent(p, (String) c);
namedXContents.add(
new NamedXContentRegistry.Entry(Aggregation.class, new ParseField(GeoHexGridAggregationBuilder.NAME), hexGridParser)
);
return new NamedXContentRegistry(namedXContents);
}

private Map<Point, String> generateRandomPointH3CellMap(int size, int randomPrecision) {
return IntStream.range(0, size)
.mapToObj(unUsed -> randomPoint())
.collect(toMap(identity(), point -> geoToH3Address(point.getLat(), point.getLon(), randomPrecision)));
}
}