Skip to content

Commit

Permalink
Remove test cases that directly update model index
Browse files Browse the repository at this point in the history
Signed-off-by: John Mazanec <[email protected]>
  • Loading branch information
jmazanec15 committed Dec 29, 2023
1 parent 1a54cf8 commit db931fe
Show file tree
Hide file tree
Showing 2 changed files with 1 addition and 61 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -10,21 +10,16 @@
import org.opensearch.action.search.SearchResponse;
import org.opensearch.client.Request;
import org.opensearch.client.Response;
import org.opensearch.client.ResponseException;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.core.xcontent.MediaTypeRegistry;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.index.util.KNNEngine;
import org.opensearch.knn.indices.ModelMetadata;
import org.opensearch.knn.indices.ModelState;
import org.opensearch.knn.plugin.KNNPlugin;
import org.opensearch.core.rest.RestStatus;
import org.opensearch.search.SearchHit;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Map;

Expand Down Expand Up @@ -55,7 +50,7 @@ public class ModelIT extends AbstractRestartUpgradeTestCase {
private static int DOC_ID_TEST_MODEL_INDEX = 0;
private static int DOC_ID_TEST_MODEL_INDEX_DEFAULT = 0;
private static final int DELAY_MILLI_SEC = 1000;
private static final int EXP_NUM_OF_MODELS = 3;
private static final int EXP_NUM_OF_MODELS = 2;
private static final int K = 5;
private static final int NUM_DOCS = 10;
private static final int NUM_DOCS_TEST_MODEL_INDEX = 100;
Expand All @@ -66,7 +61,6 @@ public class ModelIT extends AbstractRestartUpgradeTestCase {
private static int QUERY_COUNT_TEST_MODEL_INDEX_DEFAULT = 0;
private static final String TEST_MODEL_ID = "test-model-id";
private static final String TEST_MODEL_ID_DEFAULT = "test-model-id-default";
private static final String TEST_MODEL_ID_TRAINING = "test-model-id-training";
private static final String MODEL_DESCRIPTION = "Description for train model test";

// KNN model test
Expand Down Expand Up @@ -139,22 +133,6 @@ public void testKNNModelDefault() throws Exception {
}
}

// KNN Delete Model test for model in Training State
public void testDeleteTrainingModel() throws Exception {
byte[] testModelBlob = "hello".getBytes(StandardCharsets.UTF_8);
ModelMetadata testModelMetadata = getModelMetadata();
testModelMetadata.setState(ModelState.TRAINING);
if (isRunningAgainstOldCluster()) {
addModelToSystemIndex(TEST_MODEL_ID_TRAINING, testModelMetadata, testModelBlob);
} else {
String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, TEST_MODEL_ID_TRAINING);
Request request = new Request("DELETE", restURI);

ResponseException ex = expectThrows(ResponseException.class, () -> client().performRequest(request));
assertEquals(RestStatus.CONFLICT.getStatus(), ex.getResponse().getStatusLine().getStatusCode());
}
}

// Delete Models and ".opensearch-knn-models" index to clear cluster metadata
@AfterClass
public static void wipeAllModels() throws IOException {
Expand Down Expand Up @@ -255,8 +233,4 @@ public String modelIndexMapping(String fieldName, String modelId) throws IOExcep
.endObject()
.toString();
}

private ModelMetadata getModelMetadata() {
return new ModelMetadata(KNNEngine.DEFAULT, SpaceType.DEFAULT, 4, ModelState.CREATED, "2021-03-27", "test model", "", "");
}
}
34 changes: 0 additions & 34 deletions src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import org.opensearch.knn.index.KNNSettings;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.indices.ModelDao;
import org.opensearch.knn.indices.ModelMetadata;
import org.opensearch.knn.indices.ModelState;
import org.opensearch.knn.plugin.KNNPlugin;
import org.opensearch.knn.plugin.script.KNNScoringScriptEngine;
Expand Down Expand Up @@ -56,7 +55,6 @@
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
Expand All @@ -76,15 +74,10 @@
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_EF_SEARCH;
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_NLIST;
import static org.opensearch.knn.common.KNNConstants.METHOD_PARAMETER_SPACE_TYPE;
import static org.opensearch.knn.common.KNNConstants.MODEL_BLOB_PARAMETER;
import static org.opensearch.knn.common.KNNConstants.MODEL_DESCRIPTION;
import static org.opensearch.knn.common.KNNConstants.MODEL_ERROR;
import static org.opensearch.knn.common.KNNConstants.MODEL_NODE_ASSIGNMENT;
import static org.opensearch.knn.common.KNNConstants.MODEL_ID;
import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_MAPPING_PATH;
import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME;
import static org.opensearch.knn.common.KNNConstants.MODEL_STATE;
import static org.opensearch.knn.common.KNNConstants.MODEL_TIMESTAMP;
import static org.opensearch.knn.common.KNNConstants.TRAIN_FIELD_PARAMETER;
import static org.opensearch.knn.common.KNNConstants.TRAIN_INDEX_PARAMETER;
import static org.opensearch.knn.common.KNNConstants.NAME;
Expand Down Expand Up @@ -764,33 +757,6 @@ protected void createModelSystemIndex() throws IOException {
}
}

protected void addModelToSystemIndex(String modelId, ModelMetadata modelMetadata, byte[] model) throws IOException {
assertFalse(org.opensearch.core.common.Strings.isNullOrEmpty(modelId));
String modelBase64 = Base64.getEncoder().encodeToString(model);

Request request = new Request("POST", "/" + MODEL_INDEX_NAME + "/_doc/" + modelId + "?refresh=true");

XContentBuilder builder = XContentFactory.jsonBuilder()
.startObject()
.field(MODEL_ID, modelId)
.field(MODEL_STATE, modelMetadata.getState().getName())
.field(KNN_ENGINE, modelMetadata.getKnnEngine().getName())
.field(METHOD_PARAMETER_SPACE_TYPE, modelMetadata.getSpaceType().getValue())
.field(DIMENSION, modelMetadata.getDimension())
.field(MODEL_BLOB_PARAMETER, modelBase64)
.field(MODEL_TIMESTAMP, modelMetadata.getTimestamp())
.field(MODEL_DESCRIPTION, modelMetadata.getDescription())
.field(MODEL_ERROR, modelMetadata.getError())
.field(MODEL_NODE_ASSIGNMENT, modelMetadata.getNodeAssignment())
.endObject();

request.setJsonEntity(builder.toString());

Response response = client().performRequest(request);

assertEquals(request.getEndpoint() + ": failed", RestStatus.CREATED, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
}

/**
* Clear cache
* <p>
Expand Down

0 comments on commit db931fe

Please sign in to comment.