Skip to content

Commit

Permalink
Clean up test resource dynamically after each test case
Browse files Browse the repository at this point in the history 
Signed-off-by: Weijia Zhao <[email protected]>
  • Loading branch information
@weijia-aws
weijia-aws committed Feb 4, 2025
1 parent e83ad5d commit b9dd581
Show file tree
Hide file tree
Showing 44 changed files with 3,846 additions and 4,383 deletions.
10 changes: 3 additions & 7 deletions ...start-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,9 @@ public void testBatchIngestionWithNeuralSparseProcessor_E2EFlow() throws Excepti
String modelId = null;
modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
try {
List<Map<String, String>> docs = prepareDataForBulkIngestion(5, 5);
bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs);
validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class);
} finally {
wipeOfTestResources(indexName, PIPELINE_NAME, modelId, null);
}
List<Map<String, String>> docs = prepareDataForBulkIngestion(5, 5);
bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs);
validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class);
}
}

Expand Down
19 changes: 7 additions & 12 deletions qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,18 +67,13 @@ private void validateNormalizationProcessor(final String fileName, final String
addDocuments(getIndexNameForTest(), true);
createSearchPipeline(searchPipelineName);
} else {
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocuments(getIndexNameForTest(), false);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null);
validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder);
hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder);
} finally {
wipeOfTestResources(getIndexNameForTest(), pipelineName, modelId, searchPipelineName);
}
String modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocuments(getIndexNameForTest(), false);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null);
validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder);
hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder);
}
}

Expand Down
21 changes: 8 additions & 13 deletions ...rade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,19 +59,14 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr
Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f }))
);
} else {
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_UPGRADED, null, null);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null);
QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f);
validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer);
hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
}
String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_UPGRADED, null, null);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null);
QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f);
validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer);
hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer);
}
}

Expand Down
13 changes: 4 additions & 9 deletions ...tart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,10 @@ public void testKnnRadialSearch_E2EFlow() throws Exception {
);
addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT);
} else {
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
validateIndexQuery(modelId);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
}
String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
validateIndexQuery(modelId);
}
}

Expand Down
13 changes: 4 additions & 9 deletions ...art-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,15 +38,10 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
);
addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT);
} else {
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
validateTestIndex(modelId);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
}
String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
validateTestIndex(modelId);
}
}

Expand Down
38 changes: 14 additions & 24 deletions ...src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,18 +59,13 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws
search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits")
);
} else {
String modelId = null;
try {
modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
sparseEncodingQueryBuilderWithModelId.modelId(modelId);
assertEquals(
search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"),
search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits")
);
} finally {
wipeOfTestResources(getIndexNameForTest(), SPARSE_INGEST_PIPELINE_NAME, modelId, SPARSE_SEARCH_PIPELINE_NAME);
}
String modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
sparseEncodingQueryBuilderWithModelId.modelId(modelId);
assertEquals(
search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"),
search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits")
);
}
}

Expand Down Expand Up @@ -105,19 +100,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep
search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits")
);
} else {
String modelId = null;
try {
modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
neuralQueryBuilderWithModelId.modelId(modelId);
String modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
neuralQueryBuilderWithModelId.modelId(modelId);

assertEquals(
search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"),
search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits")
);
} finally {
wipeOfTestResources(getIndexNameForTest(), DENSE_INGEST_PIPELINE_NAME, modelId, DENSE_SEARCH_PIPELINE_NAME);
}
assertEquals(
search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"),
search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits")
);
}
}
}
27 changes: 11 additions & 16 deletions ...t-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,22 +51,17 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception {
List.of(TEXT_1)
);
} else {
String modelId = null;
try {
modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
addSparseEncodingDoc(
getIndexNameForTest(),
"1",
List.of(TEST_SPARSE_ENCODING_FIELD),
List.of(testRankFeaturesDoc2),
List.of(TEST_TEXT_FIELD),
List.of(TEXT_2)
);
validateTestIndex(modelId);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
}
String modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
addSparseEncodingDoc(
getIndexNameForTest(),
"1",
List.of(TEST_SPARSE_ENCODING_FIELD),
List.of(testRankFeaturesDoc2),
List.of(TEST_TEXT_FIELD),
List.of(TEXT_2)
);
validateTestIndex(modelId);
}

}
Expand Down
20 changes: 5 additions & 15 deletions ...rc/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,21 +45,11 @@ public void testNeuralSparseQueryTwoPhaseProcessor_NeuralSearch_E2EFlow() throws
Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits");
assertNotNull(resultWith2PhasePipeline);
} else {
String modelId = null;
try {
modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
neuralSparseQueryBuilder.modelId(modelId);
Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits");
assertNotNull(resultWith2PhasePipeline);
} finally {
wipeOfTestResources(
getIndexNameForTest(),
NEURAL_SPARSE_INGEST_PIPELINE_NAME,
modelId,
NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME
);
}
String modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR);
loadModel(modelId);
neuralSparseQueryBuilder.modelId(modelId);
Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits");
assertNotNull(resultWith2PhasePipeline);
}
}
}
13 changes: 4 additions & 9 deletions ...start-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,15 +36,10 @@ public void testTextEmbeddingProcessor_E2EFlow() throws Exception {
);
addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null);
} else {
String modelId = null;
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, null, null);
validateTestIndex(modelId);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
}
String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
loadModel(modelId);
addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, null, null);
validateTestIndex(modelId);
}
}

Expand Down
8 changes: 2 additions & 6 deletions ...pgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,8 @@ public void testTextChunkingProcessor_E2EFlow() throws Exception {
addDocument(indexName, "0", INPUT_FIELD, TEST_INGEST_TEXT, null, null);
validateTestIndex(indexName, OUTPUT_FIELD, 1, expectedPassages);
} else {
try {
addDocument(indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null);
validateTestIndex(indexName, OUTPUT_FIELD, 2, expectedPassages);
} finally {
wipeOfTestResources(indexName, PIPELINE_NAME, null, null);
}
addDocument(indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null);
validateTestIndex(indexName, OUTPUT_FIELD, 2, expectedPassages);
}
}

Expand Down
14 changes: 5 additions & 9 deletions ...lling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,15 +46,11 @@ public void testBatchIngestion_SparseEncodingProcessor_E2EFlow() throws Exceptio
validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class);
break;
case UPGRADED:
try {
sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR);
loadModel(sparseModelId);
List<Map<String, String>> docsForUpgraded = prepareDataForBulkIngestion(10, 5);
bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded);
validateDocCountAndInfo(indexName, 15, () -> getDocById(indexName, "14"), EMBEDDING_FIELD_NAME, Map.class);
} finally {
wipeOfTestResources(indexName, SPARSE_PIPELINE, sparseModelId, null);
}
sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR);
loadModel(sparseModelId);
List<Map<String, String>> docsForUpgraded = prepareDataForBulkIngestion(10, 5);
bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded);
validateDocCountAndInfo(indexName, 15, () -> getDocById(indexName, "14"), EMBEDDING_FIELD_NAME, Map.class);
break;
default:
throw new IllegalStateException("Unexpected value: " + getClusterType());
Expand Down
20 changes: 8 additions & 12 deletions qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,18 +76,14 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr
}
break;
case UPGRADED:
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND;
loadModel(modelId);
addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null);
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null);
hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME);
}
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND;
loadModel(modelId);
addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null);
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null);
hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null);
break;
default:
throw new IllegalStateException("Unexpected value: " + getClusterType());
Expand Down
22 changes: 9 additions & 13 deletions ...rade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,19 +78,15 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr
}
break;
case UPGRADED:
try {
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND;
loadModel(modelId);
addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null);
QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f);
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer);
hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer);
} finally {
wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME);
}
modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR);
int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND;
loadModel(modelId);
addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null);
HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null);
QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f);
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer);
hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault());
validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer);
break;
default:
throw new IllegalStateException("Unexpected value: " + getClusterType());
Expand Down
Loading

0 comments on commit b9dd581

Please sign in to comment.