Skip to content

Commit bcb07b2

Browse files
craigtavernercraigtaverner
committed
Add license checks and integration tests to test license checks
1 parent 5268880 commit bcb07b2

File tree

13 files changed

+447
-45
lines changed

13 files changed

+447
-45
lines changed

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialExtentAggregationNoLicenseIT.java‎

Lines changed: 1 addition & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,7 @@
77

88
package org.elasticsearch.xpack.esql.spatial;
99

10-
import org.elasticsearch.license.License;
11-
import org.elasticsearch.license.XPackLicenseState;
12-
import org.elasticsearch.license.internal.XPackLicenseStatus;
1310
import org.elasticsearch.plugins.Plugin;
14-
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
15-
import org.elasticsearch.xpack.spatial.SpatialPlugin;
1611

1712
import java.util.Collection;
1813
import java.util.List;
@@ -21,45 +16,11 @@ public class SpatialExtentAggregationNoLicenseIT extends SpatialExtentAggregatio
2116

2217
@Override
2318
protected Collection<Class<? extends Plugin>> nodePlugins() {
24-
return List.of(TestSpatialPlugin.class, TestEsqlPlugin.class);
19+
return List.of(SpatialNoLicenseTestCase.TestSpatialPlugin.class, SpatialNoLicenseTestCase.TestEsqlPlugin.class);
2520
}
2621

2722
@Override
2823
public void testStExtentAggregationWithShapes() {
2924
assertStExtentFailsWith("index_geo_shape");
3025
}
31-
32-
private static XPackLicenseState getLicenseState() {
33-
License.OperationMode operationMode;
34-
boolean active;
35-
if (randomBoolean()) {
36-
operationMode = randomFrom(
37-
License.OperationMode.GOLD,
38-
License.OperationMode.BASIC,
39-
License.OperationMode.MISSING,
40-
License.OperationMode.STANDARD
41-
);
42-
active = true;
43-
} else {
44-
operationMode = randomFrom(License.OperationMode.PLATINUM, License.OperationMode.ENTERPRISE, License.OperationMode.TRIAL);
45-
active = false; // expired
46-
}
47-
48-
return new XPackLicenseState(
49-
() -> System.currentTimeMillis(),
50-
new XPackLicenseStatus(operationMode, active, "Test license expired")
51-
);
52-
}
53-
54-
public static class TestEsqlPlugin extends EsqlPlugin {
55-
protected XPackLicenseState getLicenseState() {
56-
return SpatialExtentAggregationNoLicenseIT.getLicenseState();
57-
}
58-
}
59-
60-
public static class TestSpatialPlugin extends SpatialPlugin {
61-
protected XPackLicenseState getLicenseState() {
62-
return SpatialExtentAggregationNoLicenseIT.getLicenseState();
63-
}
64-
}
6526
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialGridLicenseTestCase.java‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.spatial;
9+
10+
import org.elasticsearch.ElasticsearchException;
11+
import org.elasticsearch.action.bulk.BulkRequestBuilder;
12+
import org.elasticsearch.action.index.IndexRequest;
13+
import org.elasticsearch.action.support.WriteRequest;
14+
import org.elasticsearch.geometry.Point;
15+
import org.elasticsearch.xpack.esql.VerificationException;
16+
import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase;
17+
import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
18+
import org.junit.Before;
19+
20+
import java.util.ArrayList;
21+
import java.util.Iterator;
22+
import java.util.LinkedHashMap;
23+
import java.util.List;
24+
import java.util.Locale;
25+
import java.util.Map;
26+
27+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
28+
import static org.hamcrest.Matchers.containsString;
29+
import static org.hamcrest.Matchers.equalTo;
30+
31+
public abstract class SpatialGridLicenseTestCase extends AbstractEsqlIntegTestCase {
32+
33+
@Before
34+
public void setupIndex() throws Exception {
35+
assumeTrue("requires SPATIAL_GRID capability", EsqlCapabilities.Cap.SPATIAL_GRID.isEnabled());
36+
createAndPopulateIndexes(10);
37+
}
38+
39+
protected List<Point> testData = new ArrayList<>();
40+
41+
protected abstract String gridFunction();
42+
43+
protected abstract Map<Long, Long> expectedValues();
44+
45+
protected int precision() {
46+
return 1; // Default precision for grid function tests, can be overridden in subclasses
47+
}
48+
49+
/**
50+
* This test will fail without a platinum license
51+
*/
52+
public abstract void testGeoGridWithShapes();
53+
54+
/**
55+
* This test should pass with and without platinum licenses.
56+
* This should only be overridden for ST_GEOHEX which also licenses geo_point use cases.
57+
*/
58+
public void testGeoGridWithPoints() {
59+
assertGeoGridFromIndex("index_geo_point");
60+
}
61+
62+
protected void assertGeoGridFromIndex(String index) {
63+
assumeTrue("geo_shape capability not yet implemented", index.equals("index_geo_point"));
64+
var query = String.format(Locale.ROOT, """
65+
FROM %s
66+
| EVAL gridId = %s(location, %s)
67+
| STATS count=COUNT() BY gridId
68+
| KEEP gridId, count
69+
""", index, gridFunction(), precision());
70+
try (var resp = run(query)) {
71+
assertColumnNames(resp.columns(), List.of("gridId", "count"));
72+
assertColumnTypes(resp.columns(), List.of("long", "long"));
73+
Map<Long, Long> values = getValuesMap(resp.values());
74+
Map<Long, Long> expected = expectedValues();
75+
assertThat(values.size(), equalTo(expected.size()));
76+
for (Long gridId : expected.keySet()) {
77+
assertThat("Missing grid-id: " + gridId, values.containsKey(gridId), equalTo(true));
78+
assertThat("Unexpected count for grid-id: " + gridId, values.get(gridId), equalTo(expected.get(gridId)));
79+
}
80+
}
81+
}
82+
83+
protected void assertGeoGridFailsWith(String index) {
84+
assumeTrue("geo_shape capability not yet implemented", index.equals("index_geo_point"));
85+
var query = String.format(Locale.ROOT, """
86+
FROM %s
87+
| EVAL gridId = %s(location, %d)
88+
| STATS count=COUNT() BY gridId
89+
""", index, gridFunction(), precision());
90+
var expectedError = String.format(
91+
Locale.ROOT,
92+
"current license is non-compliant for [%s(location, %d)]",
93+
gridFunction(),
94+
precision()
95+
);
96+
ElasticsearchException e = expectThrows(VerificationException.class, () -> run(query));
97+
assertThat(e.getMessage(), containsString(expectedError));
98+
}
99+
100+
public static Map<Long, Long> getValuesMap(Iterator<Iterator<Object>> values) {
101+
var valuesMap = new LinkedHashMap<Long, Long>();
102+
values.forEachRemaining(row -> { valuesMap.put(((Number) row.next()).longValue(), ((Number) row.next()).longValue()); });
103+
return valuesMap;
104+
}
105+
106+
private void createAndPopulateIndexes(int count) throws Exception {
107+
initIndex("index_", "geo_point");
108+
initIndex("index_", "geo_shape");
109+
BulkRequestBuilder points = client().prepareBulk();
110+
BulkRequestBuilder shapes = client().prepareBulk();
111+
StringBuilder coords = new StringBuilder();
112+
for (int i = 0; i < count; i++) {
113+
double x = randomDoubleBetween(-10.0, 10.0, true);
114+
double y = randomDoubleBetween(-10.0, 10.0, true);
115+
Point point = new Point(x, y);
116+
testData.add(point);
117+
points.add(new IndexRequest("index_geo_point").id(x + ":" + y).source("location", point.toString()));
118+
if (coords.length() > 0) {
119+
coords.append(", ");
120+
}
121+
coords.append(x).append(" ").append(y);
122+
}
123+
if (coords.length() > 0) {
124+
String lineString = "LINESTRING (" + coords + ")";
125+
shapes.add(new IndexRequest("index_geo_shape").id("polygon").source("location", lineString));
126+
}
127+
points.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
128+
shapes.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
129+
ensureYellow("index_geo_point");
130+
ensureYellow("index_geo_shape");
131+
}
132+
133+
protected void initIndex(String prefix, String fieldType) {
134+
assertAcked(prepareCreate(prefix + fieldType).setMapping(String.format(Locale.ROOT, """
135+
{
136+
"properties" : {
137+
"location": { "type" : "%s" }
138+
}
139+
}
140+
""", fieldType)));
141+
}
142+
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialNoLicenseTestCase.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.spatial;
9+
10+
import org.elasticsearch.license.License;
11+
import org.elasticsearch.license.XPackLicenseState;
12+
import org.elasticsearch.license.internal.XPackLicenseStatus;
13+
import org.elasticsearch.test.ESIntegTestCase;
14+
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
15+
import org.elasticsearch.xpack.spatial.SpatialPlugin;
16+
17+
/**
18+
* Utility class to provide a license state that does not allow spatial features.
19+
* This is used in tests to ensure that spatial functions behave correctly when no valid license is present.
20+
*/
21+
public abstract class SpatialNoLicenseTestCase extends ESIntegTestCase {
22+
23+
private static XPackLicenseState getLicenseState() {
24+
License.OperationMode operationMode;
25+
boolean active;
26+
if (randomBoolean()) {
27+
// Randomly chosen licenses that are not Platinum, Enterprise, or Trial
28+
operationMode = randomFrom(
29+
License.OperationMode.GOLD,
30+
License.OperationMode.BASIC,
31+
License.OperationMode.MISSING,
32+
License.OperationMode.STANDARD
33+
);
34+
active = true;
35+
} else {
36+
// Randomly chosen licenses that are Platinum, Enterprise, or Trial but marked as expired
37+
operationMode = randomFrom(License.OperationMode.PLATINUM, License.OperationMode.ENTERPRISE, License.OperationMode.TRIAL);
38+
active = false; // expired
39+
}
40+
41+
return new XPackLicenseState(
42+
() -> System.currentTimeMillis(),
43+
new XPackLicenseStatus(operationMode, active, "Test license expired")
44+
);
45+
}
46+
47+
/**
48+
* Test plugin that provides a license state that does not allow spatial features.
49+
* This is used to test the behavior of spatial functions when no valid license is present.
50+
*/
51+
public static class TestEsqlPlugin extends EsqlPlugin {
52+
protected XPackLicenseState getLicenseState() {
53+
return SpatialNoLicenseTestCase.getLicenseState();
54+
}
55+
}
56+
57+
/**
58+
* Test plugin that provides a license state that does not allow spatial features.
59+
* This is used to test the behavior of spatial functions when no valid license is present.
60+
*/
61+
public static class TestSpatialPlugin extends SpatialPlugin {
62+
protected XPackLicenseState getLicenseState() {
63+
return SpatialNoLicenseTestCase.getLicenseState();
64+
}
65+
}
66+
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/StGeohashLicenseIT.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.spatial;
9+
10+
import org.elasticsearch.geometry.Point;
11+
import org.elasticsearch.plugins.Plugin;
12+
import org.elasticsearch.xpack.esql.action.EsqlPluginWithEnterpriseOrTrialLicense;
13+
import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StGeohash;
14+
import org.elasticsearch.xpack.spatial.SpatialPlugin;
15+
16+
import java.util.Collection;
17+
import java.util.HashMap;
18+
import java.util.List;
19+
import java.util.Map;
20+
21+
public class StGeohashLicenseIT extends SpatialGridLicenseTestCase {
22+
@Override
23+
protected Collection<Class<? extends Plugin>> nodePlugins() {
24+
return List.of(SpatialPlugin.class, EsqlPluginWithEnterpriseOrTrialLicense.class);
25+
}
26+
27+
@Override
28+
protected String gridFunction() {
29+
return "ST_GEOHASH";
30+
}
31+
32+
@Override
33+
protected Map<Long, Long> expectedValues() {
34+
Map<Long, Long> expected = new HashMap<>();
35+
for (Point point : testData) {
36+
long gridId = StGeohash.unboundedGrid.calculateGridId(point, precision());
37+
expected.compute(gridId, (k, v) -> v == null ? 1 : v + 1);
38+
}
39+
return expected;
40+
}
41+
42+
public void testGeoGridWithShapes() {
43+
assertGeoGridFromIndex("index_geo_shape");
44+
}
45+
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/StGeohashNoLicenseIT.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.spatial;
9+
10+
import org.elasticsearch.plugins.Plugin;
11+
12+
import java.util.Collection;
13+
import java.util.List;
14+
15+
public class StGeohashNoLicenseIT extends StGeohashLicenseIT {
16+
@Override
17+
protected Collection<Class<? extends Plugin>> nodePlugins() {
18+
return List.of(SpatialNoLicenseTestCase.TestSpatialPlugin.class, SpatialNoLicenseTestCase.TestEsqlPlugin.class);
19+
}
20+
21+
public void testGeoGridWithShapes() {
22+
assertGeoGridFailsWith("index_geo_shape");
23+
}
24+
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/StGeohexLicenseIT.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.spatial;
9+
10+
import org.elasticsearch.geometry.Point;
11+
import org.elasticsearch.plugins.Plugin;
12+
import org.elasticsearch.xpack.esql.action.EsqlPluginWithEnterpriseOrTrialLicense;
13+
import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StGeohex;
14+
import org.elasticsearch.xpack.spatial.SpatialPlugin;
15+
16+
import java.util.Collection;
17+
import java.util.HashMap;
18+
import java.util.List;
19+
import java.util.Map;
20+
21+
public class StGeohexLicenseIT extends SpatialGridLicenseTestCase {
22+
@Override
23+
protected Collection<Class<? extends Plugin>> nodePlugins() {
24+
return List.of(SpatialPlugin.class, EsqlPluginWithEnterpriseOrTrialLicense.class);
25+
}
26+
27+
@Override
28+
protected String gridFunction() {
29+
return "ST_GEOHEX";
30+
}
31+
32+
@Override
33+
protected Map<Long, Long> expectedValues() {
34+
Map<Long, Long> expected = new HashMap<>();
35+
for (Point point : testData) {
36+
long gridId = StGeohex.unboundedGrid.calculateGridId(point, precision());
37+
expected.compute(gridId, (k, v) -> v == null ? 1 : v + 1);
38+
}
39+
return expected;
40+
}
41+
42+
public void testGeoGridWithShapes() {
43+
assertGeoGridFromIndex("index_geo_shape");
44+
}
45+
}

0 commit comments

Comments
 (0)