Skip to content

Commit bf71a49

Browse files
l-trottal-trotta
committed
refactor unit test
1 parent da00448 commit bf71a49

File tree

1 file changed

+28
-35
lines changed

1 file changed

+28
-35
lines changed

‎vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java

Lines changed: 28 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -132,13 +132,12 @@ public void addAndSearchTest(String similarityFunction) {
132132
vectorStore.add(documents);
133133

134134
Awaitility.await()
135-
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression")
136-
.withTopK(1)
137-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(1));
135+
.until(() -> vectorStore
136+
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThresholdAll()),
137+
hasSize(1));
138138

139-
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great Depression")
140-
.withTopK(1)
141-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL));
139+
List<Document> results = vectorStore
140+
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThresholdAll());
142141

143142
assertThat(results).hasSize(1);
144143
Document resultDoc = results.get(0);
@@ -152,9 +151,9 @@ public void addAndSearchTest(String similarityFunction) {
152151
vectorStore.delete(documents.stream().map(Document::getId).toList());
153152

154153
Awaitility.await()
155-
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression")
156-
.withTopK(1)
157-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(0));
154+
.until(() -> vectorStore
155+
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1).withSimilarityThresholdAll()),
156+
hasSize(0));
158157
});
159158
}
160159

@@ -177,21 +176,21 @@ public void searchWithFilters(String similarityFunction) {
177176
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
178177

179178
Awaitility.await()
180-
.until(() -> vectorStore.similaritySearch(SearchRequest.query("The World")
181-
.withTopK(5)
182-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(3));
179+
.until(() -> vectorStore
180+
.similaritySearch(SearchRequest.query("The World").withTopK(5).withSimilarityThresholdAll()),
181+
hasSize(3));
183182

184183
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("The World")
185184
.withTopK(5)
186-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
185+
.withSimilarityThresholdAll()
187186
.withFilterExpression("country == 'NL'"));
188187

189188
assertThat(results).hasSize(1);
190189
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
191190

192191
results = vectorStore.similaritySearch(SearchRequest.query("The World")
193192
.withTopK(5)
194-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
193+
.withSimilarityThresholdAll()
195194
.withFilterExpression("country == 'BG'"));
196195

197196
assertThat(results).hasSize(2);
@@ -200,15 +199,15 @@ public void searchWithFilters(String similarityFunction) {
200199

201200
results = vectorStore.similaritySearch(SearchRequest.query("The World")
202201
.withTopK(5)
203-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
202+
.withSimilarityThresholdAll()
204203
.withFilterExpression("country == 'BG' && year == 2020"));
205204

206205
assertThat(results).hasSize(1);
207206
assertThat(results.get(0).getId()).isEqualTo(bgDocument.getId());
208207

209208
results = vectorStore.similaritySearch(SearchRequest.query("The World")
210209
.withTopK(5)
211-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
210+
.withSimilarityThresholdAll()
212211
.withFilterExpression("country in ['BG']"));
213212

214213
assertThat(results).hasSize(2);
@@ -217,22 +216,22 @@ public void searchWithFilters(String similarityFunction) {
217216

218217
results = vectorStore.similaritySearch(SearchRequest.query("The World")
219218
.withTopK(5)
220-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
219+
.withSimilarityThresholdAll()
221220
.withFilterExpression("country in ['BG','NL']"));
222221

223222
assertThat(results).hasSize(3);
224223

225224
results = vectorStore.similaritySearch(SearchRequest.query("The World")
226225
.withTopK(5)
227-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
226+
.withSimilarityThresholdAll()
228227
.withFilterExpression("country not in ['BG']"));
229228

230229
assertThat(results).hasSize(1);
231230
assertThat(results.get(0).getId()).isEqualTo(nlDocument.getId());
232231

233232
results = vectorStore.similaritySearch(SearchRequest.query("The World")
234233
.withTopK(5)
235-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
234+
.withSimilarityThresholdAll()
236235
.withFilterExpression("NOT(country not in ['BG'])"));
237236

238237
assertThat(results).hasSize(2);
@@ -241,7 +240,7 @@ public void searchWithFilters(String similarityFunction) {
241240

242241
results = vectorStore.similaritySearch(SearchRequest.query("The World")
243242
.withTopK(5)
244-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
243+
.withSimilarityThresholdAll()
245244
.withFilterExpression(
246245
"activationDate > " + ZonedDateTime.parse("1970-01-01T00:00:02Z").toInstant().toEpochMilli()));
247246

@@ -270,13 +269,12 @@ public void documentUpdateTest(String similarityFunction) {
270269
vectorStore.add(List.of(document));
271270

272271
Awaitility.await()
273-
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Spring")
274-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
275-
.withTopK(5)), hasSize(1));
272+
.until(() -> vectorStore
273+
.similaritySearch(SearchRequest.query("Spring").withSimilarityThresholdAll().withTopK(5)),
274+
hasSize(1));
276275

277-
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring")
278-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)
279-
.withTopK(5));
276+
List<Document> results = vectorStore
277+
.similaritySearch(SearchRequest.query("Spring").withSimilarityThresholdAll().withTopK(5));
280278

281279
assertThat(results).hasSize(1);
282280
Document resultDoc = results.get(0);
@@ -289,9 +287,7 @@ public void documentUpdateTest(String similarityFunction) {
289287
"The World is Big and Salvation Lurks Around the Corner", Map.of("meta2", "meta2"));
290288

291289
vectorStore.add(List.of(sameIdDocument));
292-
SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar")
293-
.withTopK(5)
294-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL);
290+
SearchRequest fooBarSearchRequest = SearchRequest.query("FooBar").withTopK(5).withSimilarityThresholdAll();
295291

296292
Awaitility.await()
297293
.until(() -> vectorStore.similaritySearch(fooBarSearchRequest).get(0).getContent(),
@@ -324,9 +320,7 @@ public void searchThresholdTest(String similarityFunction) {
324320

325321
vectorStore.add(documents);
326322

327-
SearchRequest query = SearchRequest.query("Great Depression")
328-
.withTopK(50)
329-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL);
323+
SearchRequest query = SearchRequest.query("Great Depression").withTopK(50).withSimilarityThresholdAll();
330324

331325
Awaitility.await().until(() -> vectorStore.similaritySearch(query), hasSize(3));
332326

@@ -352,9 +346,8 @@ public void searchThresholdTest(String similarityFunction) {
352346
vectorStore.delete(documents.stream().map(Document::getId).toList());
353347

354348
Awaitility.await()
355-
.until(() -> vectorStore.similaritySearch(SearchRequest.query("Great Depression")
356-
.withTopK(50)
357-
.withSimilarityThreshold(SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL)), hasSize(0));
349+
.until(() -> vectorStore.similaritySearch(
350+
SearchRequest.query("Great Depression").withTopK(50).withSimilarityThresholdAll()), hasSize(0));
358351
});
359352
}
360353

0 commit comments

Comments
 (0)