Skip to content

Commit 7106d20

Browse files
martin-gaievskimartin-gaievski
authored
Optimizing integ tests for less model upload calls (#683)
* Optimizing tests for less model upload calls Signed-off-by: Martin Gaievski <gaievski@amazon.com>
1 parent ff95604 commit 7106d20

File tree

8 files changed

+139
-289
lines changed

8 files changed

+139
-289
lines changed

‎CONTRIBUTING.md

+7-4
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,13 @@ To send us a pull request, please:
3131

3232
1. Fork the repository.
3333
2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change.
34-
3. Ensure local tests pass.
35-
4. Commit to your fork using clear commit messages.
36-
5. Send us a pull request, answering any default questions in the pull request interface.
37-
6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation.
34+
3. Include tests that check your new feature or bug fix. Ideally, we're looking for unit, integration, and BWC tests, but that depends on how big and critical your change is.
35+
If you're adding an integration test and it is using local ML models, please make sure that the number of model deployments is limited, and you're using the smallest possible model.
36+
Each model deployment consumes resources, and having too many models may cause unexpected test failures.
37+
4. Ensure local tests pass.
38+
5. Commit to your fork using clear commit messages.
39+
6. Send us a pull request, answering any default questions in the pull request interface.
40+
7. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation.
3841

3942
GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and
4043
[creating a pull request](https://help.github.com/articles/creating-a-pull-request/).

‎src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java

+17-35
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSu
164164
}
165165

166166
@SneakyThrows
167-
public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful() {
167+
public void testQueryMatches_whenMultipleShards_thenSuccessful() {
168168
String modelId = null;
169169
try {
170170
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME);
@@ -223,55 +223,37 @@ public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful
223223

224224
// verify that all ids are unique
225225
assertEquals(Set.copyOf(ids).size(), ids.size());
226-
} finally {
227-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE);
228-
}
229-
}
230-
231-
@SneakyThrows
232-
public void testResultProcessor_whenMultipleShardsAndNoMatches_thenSuccessful() {
233-
try {
234-
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME);
235-
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
236226

237-
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
238-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT6));
239-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
227+
// verify case when there are partial match
228+
HybridQueryBuilder hybridQueryBuilderPartialMatch = new HybridQueryBuilder();
229+
hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
230+
hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4));
231+
hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
240232

241-
Map<String, Object> searchResponseAsMap = search(
233+
Map<String, Object> searchResponseAsMapPartialMatch = search(
242234
TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME,
243-
hybridQueryBuilder,
235+
hybridQueryBuilderPartialMatch,
244236
null,
245237
5,
246238
Map.of("search_pipeline", SEARCH_PIPELINE)
247239
);
248-
assertQueryResults(searchResponseAsMap, 0, true);
249-
} finally {
250-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE);
251-
}
252-
}
253-
254-
@SneakyThrows
255-
public void testResultProcessor_whenMultipleShardsAndPartialMatches_thenSuccessful() {
256-
try {
257-
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME);
258-
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
240+
assertQueryResults(searchResponseAsMapPartialMatch, 4, true, Range.between(0.33f, 1.0f));
259241

260-
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
261-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
262-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4));
263-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
242+
// verify case when query doesn't have a match
243+
HybridQueryBuilder hybridQueryBuilderNoMatches = new HybridQueryBuilder();
244+
hybridQueryBuilderNoMatches.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT6));
245+
hybridQueryBuilderNoMatches.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
264246

265-
Map<String, Object> searchResponseAsMap = search(
247+
Map<String, Object> searchResponseAsMapNoMatches = search(
266248
TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME,
267-
hybridQueryBuilder,
249+
hybridQueryBuilderNoMatches,
268250
null,
269251
5,
270252
Map.of("search_pipeline", SEARCH_PIPELINE)
271253
);
272-
assertQueryResults(searchResponseAsMap, 4, true, Range.between(0.33f, 1.0f));
254+
assertQueryResults(searchResponseAsMapNoMatches, 0, true);
273255
} finally {
274-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE);
256+
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE);
275257
}
276258
}
277259

‎src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java

+4-14
Original file line numberDiff line numberDiff line change
@@ -34,29 +34,19 @@ public void setUp() throws Exception {
3434
updateClusterSettings();
3535
}
3636

37-
public void testEmbeddingProcessor_whenIngestingDocumentWithSourceMatchingTextMapping_thenSuccessful() throws Exception {
37+
public void testEmbeddingProcessor_whenIngestingDocumentWithOrWithoutSourceMatchingMapping_thenSuccessful() throws Exception {
3838
String modelId = null;
3939
try {
4040
modelId = uploadModel();
4141
loadModel(modelId);
4242
createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING);
4343
createTextImageEmbeddingIndex();
44+
// verify doc with mapping
4445
ingestDocumentWithTextMappedToEmbeddingField();
4546
assertEquals(1, getDocCount(INDEX_NAME));
46-
} finally {
47-
wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null);
48-
}
49-
}
50-
51-
public void testEmbeddingProcessor_whenIngestingDocumentWithSourceWithoutMatchingInMapping_thenSuccessful() throws Exception {
52-
String modelId = null;
53-
try {
54-
modelId = uploadModel();
55-
loadModel(modelId);
56-
createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING);
57-
createTextImageEmbeddingIndex();
47+
// verify doc without mapping
5848
ingestDocumentWithoutMappedFields();
59-
assertEquals(1, getDocCount(INDEX_NAME));
49+
assertEquals(2, getDocCount(INDEX_NAME));
6050
} finally {
6151
wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null);
6252
}

‎src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java

+10-26
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,6 @@
1212
import org.apache.hc.core5.http.HttpHeaders;
1313
import org.apache.hc.core5.http.io.entity.EntityUtils;
1414
import org.apache.hc.core5.http.message.BasicHeader;
15-
import org.junit.After;
16-
import org.junit.Before;
1715
import org.opensearch.client.Request;
1816
import org.opensearch.client.Response;
1917
import org.opensearch.common.xcontent.XContentHelper;
@@ -36,33 +34,19 @@ public class MLOpenSearchRerankProcessorIT extends BaseNeuralSearchIT {
3634
private final static String TEXT_REP_1 = "Jacques loves fish. Fish make Jacques happy";
3735
private final static String TEXT_REP_2 = "Fish like to eat plankton";
3836
private final static String INDEX_CONFIG = "{\"mappings\": {\"properties\": {\"text_representation\": {\"type\": \"text\"}}}}";
39-
private String modelId;
40-
41-
@After
42-
@SneakyThrows
43-
public void tearDown() {
44-
super.tearDown();
45-
/* this is required to minimize chance of model not being deployed due to open memory CB,
46-
* this happens in case we leave model from previous test case. We use new model for every test, and old model
47-
* can be undeployed and deleted to free resources after each test case execution.
48-
*/
49-
deleteModel(modelId);
50-
deleteSearchPipeline(PIPELINE_NAME);
51-
deleteIndex(INDEX_NAME);
52-
}
53-
54-
@Before
55-
@SneakyThrows
56-
public void setup() {
57-
modelId = uploadTextSimilarityModel();
58-
loadModel(modelId);
59-
}
6037

6138
@SneakyThrows
6239
public void testCrossEncoderRerankProcessor() {
63-
createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json");
64-
setupIndex();
65-
runQueries();
40+
String modelId = null;
41+
try {
42+
modelId = uploadTextSimilarityModel();
43+
loadModel(modelId);
44+
createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json");
45+
setupIndex();
46+
runQueries();
47+
} finally {
48+
wipeOfTestResources(INDEX_NAME, null, modelId, PIPELINE_NAME);
49+
}
6650
}
6751

6852
private String uploadTextSimilarityModel() throws Exception {

‎src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java

+13-27
Original file line numberDiff line numberDiff line change
@@ -566,28 +566,17 @@ public void testRequestCache_whenMultipleShardsQueryReturnResults_thenSuccessful
566566

567567
@SneakyThrows
568568
public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNestedFields_thenSuccess() {
569-
String modelId = null;
570569
String alias = "alias_with_filter";
571570
try {
572571
initializeIndexIfNotExist(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME);
573-
modelId = prepareModel();
574572
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
575573
// create alias for index
576574
QueryBuilder aliasFilter = QueryBuilders.boolQuery()
577575
.mustNot(QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
578576
createIndexAlias(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, alias, aliasFilter);
579577

580-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
581-
TEST_KNN_VECTOR_FIELD_NAME_1,
582-
TEST_QUERY_TEXT,
583-
"",
584-
modelId,
585-
5,
586-
null,
587-
null
588-
);
589578
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
590-
hybridQueryBuilder.add(neuralQueryBuilder);
579+
hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1));
591580

592581
Map<String, Object> searchResponseAsMap = search(
593582
alias,
@@ -608,34 +597,23 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNested
608597
assertEquals(RELATION_EQUAL_TO, total.get("relation"));
609598
} finally {
610599
deleteIndexAlias(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, alias);
611-
wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, modelId, SEARCH_PIPELINE);
600+
wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, null, SEARCH_PIPELINE);
612601
}
613602
}
614603

615604
@SneakyThrows
616605
public void testWrappedQueryWithFilter_whenIndexAliasHasFilters_thenSuccess() {
617-
String modelId = null;
618606
String alias = "alias_with_filter";
619607
try {
620608
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME);
621-
modelId = prepareModel();
622609
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
623610
// create alias for index
624611
QueryBuilder aliasFilter = QueryBuilders.boolQuery()
625612
.mustNot(QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
626613
createIndexAlias(TEST_MULTI_DOC_INDEX_NAME, alias, aliasFilter);
627614

628-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
629-
TEST_KNN_VECTOR_FIELD_NAME_1,
630-
TEST_QUERY_TEXT,
631-
"",
632-
modelId,
633-
5,
634-
null,
635-
null
636-
);
637615
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
638-
hybridQueryBuilder.add(neuralQueryBuilder);
616+
hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1));
639617

640618
Map<String, Object> searchResponseAsMap = search(
641619
alias,
@@ -656,7 +634,7 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilters_thenSuccess() {
656634
assertEquals(RELATION_EQUAL_TO, total.get("relation"));
657635
} finally {
658636
deleteIndexAlias(TEST_MULTI_DOC_INDEX_NAME, alias);
659-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, modelId, SEARCH_PIPELINE);
637+
wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, SEARCH_PIPELINE);
660638
}
661639
}
662640

@@ -892,7 +870,15 @@ private void addDocsToIndex(final String testMultiDocIndexName) {
892870
Collections.singletonList(TEST_TEXT_FIELD_NAME_1),
893871
Collections.singletonList(TEST_DOC_TEXT2)
894872
);
895-
assertEquals(3, getDocCount(testMultiDocIndexName));
873+
addKnnDoc(
874+
testMultiDocIndexName,
875+
"4",
876+
Collections.emptyList(),
877+
Collections.emptyList(),
878+
Collections.singletonList(TEST_TEXT_FIELD_NAME_1),
879+
Collections.singletonList(TEST_DOC_TEXT3)
880+
);
881+
assertEquals(4, getDocCount(testMultiDocIndexName));
896882
}
897883

898884
private List<Map<String, Object>> getNestedHits(Map<String, Object> searchResponseAsMap) {

0 commit comments

Comments
 (0)