refactor unit test

spring-projects · Apr 16, 2024 · bf71a49 · bf71a49
1 parent da00448
commit bf71a49
Showing 1 changed file with 28 additions and 35 deletions.
diff --git a/...ch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java b/...ch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java
@@ -132,13 +132,12 @@ public void addAndSearchTest(String similarityFunction) {
 			vectorStore.add(documents);
 
 			Awaitility.await()
-				.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression")
-					.withTopK(1)
-					.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(1));
+				.until(() -> vectorStore
+					.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThresholdAll()),
+						hasSize(1));
 
-			List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great Depression")
-				.withTopK(1)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL));
+			List<Document> results = vectorStore
+				.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThresholdAll());
 
 			assertThat(results).hasSize(1);
 			Document resultDoc = results.get(0);
@@ -152,9 +151,9 @@ public void addAndSearchTest(String similarityFunction) {
 			vectorStore.delete(documents.stream().map(Document::getId).toList());
 
 			Awaitility.await()
-				.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression")
-					.withTopK(1)
-					.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(0));
+				.until(() -> vectorStore
+					.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThresholdAll()),
+						hasSize(0));
 		});
 	}
 
@@ -177,21 +176,21 @@ public void searchWithFilters(String similarityFunction) {
 			vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
 
 			Awaitility.await()
-				.until(() -> vectorStore.similaritySearch(SearchRequest.query("The World")
-					.withTopK(5)
-					.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(3));
+				.until(() -> vectorStore
+					.similaritySearch(SearchRequest.query("The World").withTopK(5).withSimilarityThresholdAll()),
+						hasSize(3));
 
 			List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("country == 'NL'"));
 
 			assertThat(results).hasSize(1);
 			assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("country == 'BG'"));
 
 			assertThat(results).hasSize(2);
@@ -200,15 +199,15 @@ public void searchWithFilters(String similarityFunction) {
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("country == 'BG' && year == 2020"));
 
 			assertThat(results).hasSize(1);
 			assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("country in ['BG']"));
 
 			assertThat(results).hasSize(2);
@@ -217,22 +216,22 @@ public void searchWithFilters(String similarityFunction) {
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("country in ['BG','NL']"));
 
 			assertThat(results).hasSize(3);
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("country not in ['BG']"));
 
 			assertThat(results).hasSize(1);
 			assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression("NOT(country not in ['BG'])"));
 
 			assertThat(results).hasSize(2);
@@ -241,7 +240,7 @@ public void searchWithFilters(String similarityFunction) {
 
 			results = vectorStore.similaritySearch(SearchRequest.query("The World")
 				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
+				.withSimilarityThresholdAll()
 				.withFilterExpression(
 						"activationDate > " + ZonedDateTime.parse("1970-01-01T00:00:02Z").toInstant().toEpochMilli()));
 
@@ -270,13 +269,12 @@ public void documentUpdateTest(String similarityFunction) {
 			vectorStore.add(List.of(document));
 
 			Awaitility.await()
-				.until(() -> vectorStore.similaritySearch(SearchRequest.query("Spring")
-					.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
-					.withTopK(5)), hasSize(1));
+				.until(() -> vectorStore
+					.similaritySearch(SearchRequest.query("Spring").withSimilarityThresholdAll().withTopK(5)),
+						hasSize(1));
 
-			List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring")
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
-				.withTopK(5));
+			List<Document> results = vectorStore
+				.similaritySearch(SearchRequest.query("Spring").withSimilarityThresholdAll().withTopK(5));
 
 			assertThat(results).hasSize(1);
 			Document resultDoc = results.get(0);
@@ -289,9 +287,7 @@ public void documentUpdateTest(String similarityFunction) {
 					"The World is Big and Salvation Lurks Around the Corner", Map.of("meta2", "meta2"));
 
 			vectorStore.add(List.of(sameIdDocument));
-			SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar")
-				.withTopK(5)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL);
+			SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar").withTopK(5).withSimilarityThresholdAll();
 
 			Awaitility.await()
 				.until(() -> vectorStore.similaritySearch(fooBarSearchRequest).get(0).getContent(),
@@ -324,9 +320,7 @@ public void searchThresholdTest(String similarityFunction) {
 
 			vectorStore.add(documents);
 
-			SearchRequest query = SearchRequest.query("Great Depression")
-				.withTopK(50)
-				.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL);
+			SearchRequest query = SearchRequest.query("Great Depression").withTopK(50).withSimilarityThresholdAll();
 
 			Awaitility.await().until(() -> vectorStore.similaritySearch(query), hasSize(3));
 
@@ -352,9 +346,8 @@ public void searchThresholdTest(String similarityFunction) {
 			vectorStore.delete(documents.stream().map(Document::getId).toList());
 
 			Awaitility.await()
-				.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression")
-					.withTopK(50)
-					.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(0));
+				.until(() -> vectorStore.similaritySearch(
+						SearchRequest.query("Great Depression").withTopK(50).withSimilarityThresholdAll()), hasSize(0));
 		});
 	}