Skip to content

Commit 596f752

Browse files
committed
Refactor integ tests that access model index (#1423)
Refactors integration tests that directly access the model system index. End users should not be directly accessing the model system index. It is supposed to be an implementation detail. We have written restful integration tests that directly access the model system index in order to initialize the cluster state. However, we should not do this because users should not be able to interact with it through restful APIs That being said, some of this implementation detail leaks out into the interface. For instance, in k-NN stats we have a stat that is the model system index status. So, in order to test this, we do need direct access to the system index. Similarly, for search, we execute the search against the system index and directly return the results. This is probably a bug - but we still need to test it. Signed-off-by: John Mazanec <jmazane@amazon.com> (cherry picked from commit 2b963b4)
1 parent 85fb70f commit 596f752

File tree

6 files changed

+45
-107
lines changed

6 files changed

+45
-107
lines changed

CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
2929
### Infrastructure
3030
* Upgrade gradle to 8.4 [1289](https://github.com/opensearch-project/k-NN/pull/1289)
3131
* Refactor security testing to install from individual components [#1307](https://github.com/opensearch-project/k-NN/pull/1307)
32+
* Refactor integ tests that access model index [#1423](https://github.com/opensearch-project/k-NN/pull/1423)
3233
### Documentation
3334
### Maintenance
3435
* Update developer guide to include M1 Setup [#1222](https://github.com/opensearch-project/k-NN/pull/1222)

src/test/java/org/opensearch/knn/plugin/action/RestDeleteModelHandlerIT.java

+6-42
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
package org.opensearch.knn.plugin.action;
1313

14+
import lombok.SneakyThrows;
1415
import org.apache.http.util.EntityUtils;
1516
import org.opensearch.client.Request;
1617
import org.opensearch.client.Response;
@@ -30,7 +31,6 @@
3031
import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE;
3132
import static org.opensearch.knn.common.KNNConstants.MODEL_ID;
3233
import static org.opensearch.knn.common.KNNConstants.MODELS;
33-
import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME;
3434
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_SPACE_TYPE;
3535
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NLIST;
3636
import static org.opensearch.knn.common.KNNConstants.METHOD_ENCODER_PARAMETER;
@@ -43,9 +43,8 @@
4343

4444
public class RestDeleteModelHandlerIT extends KNNRestTestCase {
4545

46-
public void testDeleteModelExists() throws Exception {
47-
createModelSystemIndex();
48-
46+
@SneakyThrows
47+
public void testDelete_whenModelExists_thenDeletionSucceeds() {
4948
String modelId = "test-model-id";
5049
String trainingIndexName = "train-index";
5150
String trainingFieldName = "train-field";
@@ -80,42 +79,8 @@ public void testDeleteModelExists() throws Exception {
8079
assertTrue(ex.getMessage().contains(modelId));
8180
}
8281

83-
public void testDeleteTrainingModel() throws Exception {
84-
createModelSystemIndex();
85-
86-
String modelId = "test-model-id";
87-
String trainingIndexName = "train-index";
88-
String trainingFieldName = "train-field";
89-
int dimension = 8;
90-
String modelDescription = "dummy description";
91-
92-
createBasicKnnIndex(trainingIndexName, trainingFieldName, dimension);
93-
// we do not wait for training to be completed
94-
ingestDataAndTrainModel(modelId, trainingIndexName, trainingFieldName, dimension, modelDescription);
95-
96-
Response getModelResponse = getModel(modelId, List.of());
97-
assertEquals(RestStatus.OK, RestStatus.fromCode(getModelResponse.getStatusLine().getStatusCode()));
98-
99-
String responseBody = EntityUtils.toString(getModelResponse.getEntity());
100-
assertNotNull(responseBody);
101-
102-
Map<String, Object> responseMap = createParser(XContentType.JSON.xContent(), responseBody).map();
103-
104-
assertEquals(modelId, responseMap.get(MODEL_ID));
105-
106-
String deleteModelRestURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, modelId);
107-
Request deleteModelRequest = new Request("DELETE", deleteModelRestURI);
108-
109-
ResponseException ex = expectThrows(ResponseException.class, () -> client().performRequest(deleteModelRequest));
110-
assertEquals(RestStatus.CONFLICT.getStatus(), ex.getResponse().getStatusLine().getStatusCode());
111-
112-
// need to wait for training operation as it's required for after test cleanup
113-
assertTrainingSucceeds(modelId, NUM_OF_ATTEMPTS, DELAY_MILLI_SEC);
114-
}
115-
116-
public void testDeleteModelFailsInvalid() throws Exception {
82+
public void testDelete_whenModelIDIsInvalid_thenFail() {
11783
String modelId = "invalid-model-id";
118-
createModelSystemIndex();
11984
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, modelId);
12085
Request request = new Request("DELETE", restURI);
12186

@@ -124,7 +89,8 @@ public void testDeleteModelFailsInvalid() throws Exception {
12489
}
12590

12691
// Test Train Model -> Delete Model -> Train Model with same modelId
127-
public void testTrainingDeletedModel() throws Exception {
92+
@SneakyThrows
93+
public void testTraining_whenModelHasBeenDeleted_thenSucceedTrainingModelWithSameID() {
12894
String modelId = "test-model-id1";
12995
String trainingIndexName1 = "train-index-1";
13096
String trainingIndexName2 = "train-index-2";
@@ -141,8 +107,6 @@ public void testTrainingDeletedModel() throws Exception {
141107
Response response = client().performRequest(request);
142108
assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
143109

144-
assertEquals(0, getDocCount(MODEL_INDEX_NAME));
145-
146110
// Train Model again with same ModelId
147111
trainModel(modelId, trainingIndexName2, trainingFieldName, dimension);
148112
}

src/test/java/org/opensearch/knn/plugin/action/RestGetModelHandlerIT.java

+7-10
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
package org.opensearch.knn.plugin.action;
1313

1414
import joptsimple.internal.Strings;
15+
import lombok.SneakyThrows;
1516
import org.apache.http.util.EntityUtils;
1617
import org.opensearch.client.Request;
1718
import org.opensearch.client.Response;
@@ -21,7 +22,6 @@
2122
import org.opensearch.knn.plugin.KNNPlugin;
2223
import org.opensearch.core.rest.RestStatus;
2324

24-
import java.io.IOException;
2525
import java.util.Arrays;
2626
import java.util.List;
2727
import java.util.Map;
@@ -44,9 +44,8 @@
4444

4545
public class RestGetModelHandlerIT extends KNNRestTestCase {
4646

47-
public void testGetModelExists() throws Exception {
48-
createModelSystemIndex();
49-
47+
@SneakyThrows
48+
public void testGetModel_whenModelIdExists_thenSucceed() {
5049
String modelId = "test-model-id";
5150
String trainingIndexName = "train-index";
5251
String trainingFieldName = "train-field";
@@ -81,8 +80,8 @@ public void testGetModelExists() throws Exception {
8180
assertEquals(L2.getValue(), responseMap.get(METHOD_PARAMETER_SPACE_TYPE));
8281
}
8382

84-
public void testGetModelExistsWithFilter() throws Exception {
85-
createModelSystemIndex();
83+
@SneakyThrows
84+
public void testGetModel_whenFilterApplied_thenReturnExpectedFields() {
8685
String modelId = "test-model-id";
8786
String trainingIndexName = "train-index";
8887
String trainingFieldName = "train-field";
@@ -118,17 +117,15 @@ public void testGetModelExistsWithFilter() throws Exception {
118117
assertFalse(responseMap.containsKey(MODEL_STATE));
119118
}
120119

121-
public void testGetModelFailsInvalid() throws IOException {
122-
createModelSystemIndex();
120+
public void testGetModel_whenModelIDIsInValid_thenFail() {
123121
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, "invalid-model-id");
124122
Request request = new Request("GET", restURI);
125123

126124
ResponseException ex = expectThrows(ResponseException.class, () -> client().performRequest(request));
127125
assertTrue(ex.getMessage().contains("\"invalid-model-id\""));
128126
}
129127

130-
public void testGetModelFailsBlank() throws IOException {
131-
createModelSystemIndex();
128+
public void testGetModel_whenIDIsBlank_thenFail() {
132129
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, " ");
133130
Request request = new Request("GET", restURI);
134131

src/test/java/org/opensearch/knn/plugin/action/RestKNNStatsHandlerIT.java

+6-3
Original file line numberDiff line numberDiff line change
@@ -338,8 +338,7 @@ public void testScriptStats_multipleShards() throws Exception {
338338
assertEquals(initialScriptQueryErrors + 2, (int) (nodeStats.get(0).get(StatNames.SCRIPT_QUERY_ERRORS.getName())));
339339
}
340340

341-
public void testModelIndexHealthMetricsStats() throws IOException {
342-
// Create request that filters only model index
341+
public void testModelIndexHealthMetricsStats() throws Exception {
343342
String modelIndexStatusName = StatNames.MODEL_INDEX_STATUS.getName();
344343
// index can be created in one of previous tests, and as we do not delete it each test the check below became optional
345344
if (!systemIndexExists(MODEL_INDEX_NAME)) {
@@ -351,7 +350,11 @@ public void testModelIndexHealthMetricsStats() throws IOException {
351350
// Check that model health status is null since model index is not created to system yet
352351
assertNull(statsMap.get(StatNames.MODEL_INDEX_STATUS.getName()));
353352

354-
createModelSystemIndex();
353+
// Train a model so that the system index will get created
354+
createBasicKnnIndex(TRAINING_INDEX, TRAINING_FIELD, DIMENSION);
355+
bulkIngestRandomVectors(TRAINING_INDEX, TRAINING_FIELD, NUM_DOCS, DIMENSION);
356+
trainKnnModel(TEST_MODEL_ID, TRAINING_INDEX, TRAINING_FIELD, DIMENSION, MODEL_DESCRIPTION);
357+
validateModelCreated(TEST_MODEL_ID);
355358
}
356359

357360
Response response = getKnnStats(Collections.emptyList(), Arrays.asList(modelIndexStatusName));

src/test/java/org/opensearch/knn/plugin/action/RestSearchModelHandlerIT.java

+25-32
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
package org.opensearch.knn.plugin.action;
1313

14+
import lombok.SneakyThrows;
1415
import org.apache.http.util.EntityUtils;
1516
import org.opensearch.action.search.SearchResponse;
1617
import org.opensearch.client.Request;
@@ -19,22 +20,18 @@
1920
import org.opensearch.core.xcontent.XContentParser;
2021
import org.opensearch.common.xcontent.XContentType;
2122
import org.opensearch.knn.KNNRestTestCase;
22-
import org.opensearch.knn.index.SpaceType;
23-
import org.opensearch.knn.index.util.KNNEngine;
2423
import org.opensearch.knn.indices.Model;
25-
import org.opensearch.knn.indices.ModelMetadata;
26-
import org.opensearch.knn.indices.ModelState;
2724
import org.opensearch.knn.plugin.KNNPlugin;
2825
import org.opensearch.core.rest.RestStatus;
2926
import org.opensearch.search.SearchHit;
3027

31-
import java.io.IOException;
3228
import java.util.Arrays;
3329
import java.util.HashMap;
3430
import java.util.List;
3531
import java.util.Map;
3632

3733
import static org.opensearch.knn.common.KNNConstants.MODELS;
34+
import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME;
3835
import static org.opensearch.knn.common.KNNConstants.PARAM_SIZE;
3936
import static org.opensearch.knn.common.KNNConstants.SEARCH_MODEL_MAX_SIZE;
4037
import static org.opensearch.knn.common.KNNConstants.SEARCH_MODEL_MIN_SIZE;
@@ -47,12 +44,7 @@
4744

4845
public class RestSearchModelHandlerIT extends KNNRestTestCase {
4946

50-
private ModelMetadata getModelMetadata() {
51-
return new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, 4, ModelState.CREATED, "2021-03-27", "test model", "", "");
52-
}
53-
54-
public void testNotSupportedParams() throws IOException {
55-
createModelSystemIndex();
47+
public void testSearch_whenUnSupportedParamsPassed_thenFail() {
5648
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, "_search");
5749
Map<String, String> invalidParams = new HashMap<>();
5850
invalidParams.put("index", "index-name");
@@ -61,27 +53,31 @@ public void testNotSupportedParams() throws IOException {
6153
expectThrows(ResponseException.class, () -> client().performRequest(request));
6254
}
6355

64-
public void testNoModelExists() throws IOException {
65-
createModelSystemIndex();
56+
@SneakyThrows
57+
public void testSearch_whenNoModelExists_thenReturnEmptyResults() {
58+
// Currently, if the model index exists, we will return empty hits. If it does not exist, we will
59+
// throw an exception. This is somewhat of a bug considering that the model index is supposed to be
60+
// an implementation detail abstracted away from the user. However, in order to test, we need to handle
61+
// the 2 different scenarios
6662
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, "_search");
6763
Request request = new Request("GET", restURI);
6864
request.setJsonEntity("{\n" + " \"query\": {\n" + " \"match_all\": {}\n" + " }\n" + "}");
69-
70-
Response response = client().performRequest(request);
71-
assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
72-
73-
String responseBody = EntityUtils.toString(response.getEntity());
74-
assertNotNull(responseBody);
75-
76-
XContentParser parser = createParser(XContentType.JSON.xContent(), responseBody);
77-
SearchResponse searchResponse = SearchResponse.fromXContent(parser);
78-
assertNotNull(searchResponse);
79-
assertEquals(searchResponse.getHits().getHits().length, 0);
80-
65+
if (!systemIndexExists(MODEL_INDEX_NAME)) {
66+
ResponseException ex = expectThrows(ResponseException.class, () -> client().performRequest(request));
67+
assertEquals(RestStatus.NOT_FOUND.getStatus(), ex.getResponse().getStatusLine().getStatusCode());
68+
} else {
69+
Response response = client().performRequest(request);
70+
assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
71+
String responseBody = EntityUtils.toString(response.getEntity());
72+
assertNotNull(responseBody);
73+
XContentParser parser = createParser(XContentType.JSON.xContent(), responseBody);
74+
SearchResponse searchResponse = SearchResponse.fromXContent(parser);
75+
assertNotNull(searchResponse);
76+
assertEquals(searchResponse.getHits().getHits().length, 0);
77+
}
8178
}
8279

83-
public void testSizeValidationFailsInvalidSize() throws IOException {
84-
createModelSystemIndex();
80+
public void testSearch_whenInvalidSizePassed_thenFail() {
8581
for (Integer invalidSize : Arrays.asList(SEARCH_MODEL_MIN_SIZE - 1, SEARCH_MODEL_MAX_SIZE + 1)) {
8682
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, "_search?" + PARAM_SIZE + "=" + invalidSize);
8783
Request request = new Request("GET", restURI);
@@ -101,8 +97,8 @@ public void testSizeValidationFailsInvalidSize() throws IOException {
10197

10298
}
10399

104-
public void testSearchModelExists() throws Exception {
105-
createModelSystemIndex();
100+
@SneakyThrows
101+
public void testSearch_whenModelExists_thenSuccess() {
106102
String trainingIndex = "irrelevant-index";
107103
String trainingFieldName = "train-field";
108104
int dimension = 8;
@@ -151,7 +147,6 @@ public void testSearchModelExists() throws Exception {
151147
}
152148

153149
public void testSearchModelWithoutSource() throws Exception {
154-
createModelSystemIndex();
155150
String trainingIndex = "irrelevant-index";
156151
String trainingFieldName = "train-field";
157152
int dimension = 8;
@@ -192,7 +187,6 @@ public void testSearchModelWithoutSource() throws Exception {
192187
}
193188

194189
public void testSearchModelWithSourceFilteringIncludes() throws Exception {
195-
createModelSystemIndex();
196190
String trainingIndex = "irrelevant-index";
197191
String trainingFieldName = "train-field";
198192
int dimension = 8;
@@ -244,7 +238,6 @@ public void testSearchModelWithSourceFilteringIncludes() throws Exception {
244238
}
245239

246240
public void testSearchModelWithSourceFilteringExcludes() throws Exception {
247-
createModelSystemIndex();
248241
String trainingIndex = "irrelevant-index";
249242
String trainingFieldName = "train-field";
250243
int dimension = 8;

src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java

-20
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55

66
package org.opensearch.knn;
77

8-
import com.google.common.base.Charsets;
9-
import com.google.common.io.Resources;
108
import com.google.common.primitives.Floats;
119
import lombok.extern.log4j.Log4j2;
1210
import org.apache.commons.lang.StringUtils;
@@ -23,7 +21,6 @@
2321
import org.opensearch.knn.index.query.KNNQueryBuilder;
2422
import org.opensearch.knn.index.KNNSettings;
2523
import org.opensearch.knn.index.SpaceType;
26-
import org.opensearch.knn.indices.ModelDao;
2724
import org.opensearch.knn.indices.ModelState;
2825
import org.opensearch.knn.plugin.KNNPlugin;
2926
import org.opensearch.knn.plugin.script.KNNScoringScriptEngine;
@@ -52,7 +49,6 @@
5249

5350
import java.io.IOException;
5451
import java.io.InputStream;
55-
import java.net.URL;
5652
import java.nio.file.Files;
5753
import java.nio.file.Path;
5854
import java.util.ArrayList;
@@ -78,8 +74,6 @@
7874
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NLIST;
7975
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_SPACE_TYPE;
8076
import static org.opensearch.knn.common.KNNConstants.MODEL_DESCRIPTION;
81-
import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_MAPPING_PATH;
82-
import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME;
8377
import static org.opensearch.knn.common.KNNConstants.MODEL_STATE;
8478
import static org.opensearch.knn.common.KNNConstants.TRAIN_FIELD_PARAMETER;
8579
import static org.opensearch.knn.common.KNNConstants.TRAIN_INDEX_PARAMETER;
@@ -744,20 +738,6 @@ protected String getIndexSettingByName(String indexName, String settingName, boo
744738
}
745739
}
746740

747-
protected void createModelSystemIndex() throws IOException {
748-
URL url = ModelDao.class.getClassLoader().getResource(MODEL_INDEX_MAPPING_PATH);
749-
if (url == null) {
750-
throw new IllegalStateException("Unable to retrieve mapping for \"" + MODEL_INDEX_NAME + "\"");
751-
}
752-
753-
String mapping = Resources.toString(url, Charsets.UTF_8);
754-
mapping = mapping.substring(1, mapping.length() - 1);
755-
756-
if (!systemIndexExists(MODEL_INDEX_NAME)) {
757-
createIndex(MODEL_INDEX_NAME, Settings.builder().put("number_of_shards", 1).put("number_of_replicas", 0).build(), mapping);
758-
}
759-
}
760-
761741
/**
762742
* Clear cache
763743
* <p>

0 commit comments

Comments
 (0)