Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 072c50c

Browse files
martin-gaievskimartin-gaievski
committedApr 10, 2024·
Optimizing tests for less model upload calls
Signed-off-by: Martin Gaievski <gaievski@amazon.com>
1 parent cb83e20 commit 072c50c

File tree

7 files changed

+136
-286
lines changed

7 files changed

+136
-286
lines changed
 

‎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+
// very 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)
Please sign in to comment.