Skip to content

Commit

Permalink
Fix flaky tests (#1355)
Browse files Browse the repository at this point in the history 
* Fix flaky bwc tests

Signed-off-by: Ryan Bogan <[email protected]>
  • Loading branch information
@ryanbogan
ryanbogan authored Jan 3, 2024
1 parent 8d60054 commit 391a2ef
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 17 deletions.
2 changes: 2 additions & 0 deletions qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/ModelIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ public void testKNNModel() throws Exception {
createKnnIndex(testIndex, modelIndexMapping(TEST_FIELD, TEST_MODEL_ID));
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
} else {
Thread.sleep(1000);
DOC_ID = NUM_DOCS;
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
QUERY_COUNT = 2 * NUM_DOCS;
Expand Down Expand Up @@ -109,6 +110,7 @@ public void testKNNModelDefault() throws Exception {
createKnnIndex(testIndex, modelIndexMapping(TEST_FIELD, TEST_MODEL_ID_DEFAULT));
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
} else {
Thread.sleep(1000);
DOC_ID = NUM_DOCS;
addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS);
QUERY_COUNT = 2 * NUM_DOCS;
Expand Down
51 changes: 34 additions & 17 deletions src/test/java/org/opensearch/knn/plugin/transport/GetModelResponseTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,13 @@

package org.opensearch.knn.plugin.transport;

import org.mockito.MockedStatic;
import org.opensearch.Version;
import org.opensearch.common.io.stream.BytesStreamOutput;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.common.xcontent.XContentFactory;
import org.opensearch.knn.KNNTestCase;
import org.opensearch.knn.index.KNNClusterUtil;
import org.opensearch.knn.index.SpaceType;
import org.opensearch.knn.index.util.KNNEngine;
import org.opensearch.knn.indices.Model;
Expand All @@ -23,6 +26,10 @@

import java.io.IOException;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;

public class GetModelResponseTests extends KNNTestCase {

private ModelMetadata getModelMetadata(ModelState state) {
Expand All @@ -41,25 +48,35 @@ public void testStreams() throws IOException {
}

public void testXContent() throws IOException {
String modelId = "test-model";
byte[] testModelBlob = "hello".getBytes();
Model model = new Model(getModelMetadata(ModelState.CREATED), testModelBlob, modelId);
GetModelResponse getModelResponse = new GetModelResponse(model);
String expectedResponseString =
"{\"model_id\":\"test-model\",\"model_blob\":\"aGVsbG8=\",\"state\":\"created\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}";
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder();
getModelResponse.toXContent(xContentBuilder, null);
assertEquals(expectedResponseString, xContentBuilder.toString());
try (MockedStatic<KNNClusterUtil> knnClusterUtilMockedStatic = mockStatic(KNNClusterUtil.class)) {
final KNNClusterUtil knnClusterUtil = mock(KNNClusterUtil.class);
when(knnClusterUtil.getClusterMinVersion()).thenReturn(Version.CURRENT);
knnClusterUtilMockedStatic.when(KNNClusterUtil::instance).thenReturn(knnClusterUtil);
String modelId = "test-model";
byte[] testModelBlob = "hello".getBytes();
Model model = new Model(getModelMetadata(ModelState.CREATED), testModelBlob, modelId);
GetModelResponse getModelResponse = new GetModelResponse(model);
String expectedResponseString =
"{\"model_id\":\"test-model\",\"model_blob\":\"aGVsbG8=\",\"state\":\"created\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}";
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder();
getModelResponse.toXContent(xContentBuilder, null);
assertEquals(expectedResponseString, xContentBuilder.toString());
}
}

public void testXContentWithNoModelBlob() throws IOException {
String modelId = "test-model";
Model model = new Model(getModelMetadata(ModelState.FAILED), null, modelId);
GetModelResponse getModelResponse = new GetModelResponse(model);
String expectedResponseString =
"{\"model_id\":\"test-model\",\"model_blob\":\"\",\"state\":\"failed\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}";
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder();
getModelResponse.toXContent(xContentBuilder, null);
assertEquals(expectedResponseString, xContentBuilder.toString());
try (MockedStatic<KNNClusterUtil> knnClusterUtilMockedStatic = mockStatic(KNNClusterUtil.class)) {
final KNNClusterUtil knnClusterUtil = mock(KNNClusterUtil.class);
when(knnClusterUtil.getClusterMinVersion()).thenReturn(Version.CURRENT);
knnClusterUtilMockedStatic.when(KNNClusterUtil::instance).thenReturn(knnClusterUtil);
String modelId = "test-model";
Model model = new Model(getModelMetadata(ModelState.FAILED), null, modelId);
GetModelResponse getModelResponse = new GetModelResponse(model);
String expectedResponseString =
"{\"model_id\":\"test-model\",\"model_blob\":\"\",\"state\":\"failed\",\"timestamp\":\"2021-03-27 10:15:30 AM +05:30\",\"description\":\"test model\",\"error\":\"\",\"space_type\":\"l2\",\"dimension\":4,\"engine\":\"nmslib\",\"training_node_assignment\":\"\"}";
XContentBuilder xContentBuilder = XContentFactory.jsonBuilder();
getModelResponse.toXContent(xContentBuilder, null);
assertEquals(expectedResponseString, xContentBuilder.toString());
}
}
}

0 comments on commit 391a2ef

Please sign in to comment.