Skip to content

Commit 92cc3b7

Browse files
sobychackosobychacko
committed
Extract common vector store delete tests to base class
This commit extracts shared delete operation tests into a reusable BaseVectorStoreTests class. This reduces code duplication and provides a consistent test suite for delete operations across different vector store implementations. The base class includes tests for: Deleting by ID Deleting by filter expressions Deleting by string filter expressions Most of the vector store implementation now extends this base class and inherits these common tests while maintaining the ability to add vector store specific tests. Signed-off-by: Soby Chacko <[email protected]>
1 parent f650433 commit 92cc3b7

File tree

16 files changed

+341
-1214
lines changed

16 files changed

+341
-1214
lines changed

‎spring-ai-test/src/main/java/org/springframework/ai/test/vectorstore/BaseVectorStoreTests.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*
2+
* Copyright 2023-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.test.vectorstore;
18+
19+
import org.junit.jupiter.api.Test;
20+
21+
import static org.assertj.core.api.Assertions.assertThat;
22+
import static org.awaitility.Awaitility.await;
23+
24+
import java.time.Duration;
25+
import java.util.HashMap;
26+
import java.util.List;
27+
import java.util.Map;
28+
import java.util.concurrent.TimeUnit;
29+
import java.util.function.Consumer;
30+
import java.util.stream.Collectors;
31+
32+
import org.springframework.ai.document.Document;
33+
import org.springframework.ai.vectorstore.SearchRequest;
34+
import org.springframework.ai.vectorstore.VectorStore;
35+
import org.springframework.ai.vectorstore.filter.Filter;
36+
37+
/**
38+
* Base test class for VectorStore implementations. Provides common test scenarios for
39+
* delete operations.
40+
*
41+
* @author Soby Chacko
42+
*/
43+
public abstract class BaseVectorStoreTests {
44+
45+
protected abstract void executeTest(Consumer<VectorStore> testFunction);
46+
47+
protected Document createDocument(String country, Integer year) {
48+
Map<String, Object> metadata = new HashMap<>();
49+
metadata.put("country", country);
50+
if (year != null) {
51+
metadata.put("year", year);
52+
}
53+
return new Document("The World is Big and Salvation Lurks Around the Corner", metadata);
54+
}
55+
56+
protected List<Document> setupTestDocuments(VectorStore vectorStore) {
57+
var doc1 = createDocument("BG", 2020);
58+
var doc2 = createDocument("NL", null);
59+
var doc3 = createDocument("BG", 2023);
60+
61+
List<Document> documents = List.of(doc1, doc2, doc3);
62+
vectorStore.add(documents);
63+
64+
return documents;
65+
}
66+
67+
private String normalizeValue(Object value) {
68+
return value.toString().replaceAll("^\"|\"$", "").trim();
69+
}
70+
71+
private void verifyDocumentsExist(VectorStore vectorStore, List<Document> documents) {
72+
await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> {
73+
List<Document> results = vectorStore.similaritySearch(
74+
SearchRequest.builder().query("The World").topK(documents.size()).similarityThresholdAll().build());
75+
assertThat(results).hasSize(documents.size());
76+
});
77+
}
78+
79+
private void verifyDocumentsDeleted(VectorStore vectorStore, List<String> deletedIds) {
80+
await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> {
81+
List<Document> results = vectorStore
82+
.similaritySearch(SearchRequest.builder().query("The World").topK(10).similarityThresholdAll().build());
83+
84+
List<String> foundIds = results.stream().map(Document::getId).collect(Collectors.toList());
85+
86+
assertThat(foundIds).doesNotContainAnyElementsOf(deletedIds);
87+
});
88+
}
89+
90+
@Test
91+
protected void deleteById() {
92+
executeTest(vectorStore -> {
93+
List<Document> documents = setupTestDocuments(vectorStore);
94+
verifyDocumentsExist(vectorStore, documents);
95+
96+
List<String> idsToDelete = List.of(documents.get(0).getId(), documents.get(1).getId());
97+
vectorStore.delete(idsToDelete);
98+
verifyDocumentsDeleted(vectorStore, idsToDelete);
99+
100+
List<Document> results = vectorStore
101+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
102+
103+
assertThat(results).hasSize(1);
104+
assertThat(results.get(0).getId()).isEqualTo(documents.get(2).getId());
105+
Map<String, Object> metadata = results.get(0).getMetadata();
106+
assertThat(normalizeValue(metadata.get("country"))).isEqualTo("BG");
107+
assertThat(normalizeValue(metadata.get("year"))).isEqualTo("2023");
108+
109+
vectorStore.delete(List.of(documents.get(2).getId()));
110+
});
111+
}
112+
113+
@Test
114+
protected void deleteWithStringFilterExpression() {
115+
executeTest(vectorStore -> {
116+
List<Document> documents = setupTestDocuments(vectorStore);
117+
verifyDocumentsExist(vectorStore, documents);
118+
119+
List<String> bgDocIds = documents.stream()
120+
.filter(d -> "BG".equals(d.getMetadata().get("country")))
121+
.map(Document::getId)
122+
.collect(Collectors.toList());
123+
124+
vectorStore.delete("country == 'BG'");
125+
verifyDocumentsDeleted(vectorStore, bgDocIds);
126+
127+
List<Document> results = vectorStore
128+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
129+
130+
assertThat(results).hasSize(1);
131+
assertThat(normalizeValue(results.get(0).getMetadata().get("country"))).isEqualTo("NL");
132+
133+
vectorStore.delete(List.of(documents.get(1).getId()));
134+
});
135+
}
136+
137+
@Test
138+
protected void deleteByFilter() {
139+
executeTest(vectorStore -> {
140+
List<Document> documents = setupTestDocuments(vectorStore);
141+
verifyDocumentsExist(vectorStore, documents);
142+
143+
List<String> bgDocIds = documents.stream()
144+
.filter(d -> "BG".equals(d.getMetadata().get("country")))
145+
.map(Document::getId)
146+
.collect(Collectors.toList());
147+
148+
Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
149+
new Filter.Key("country"), new Filter.Value("BG"));
150+
151+
vectorStore.delete(filterExpression);
152+
verifyDocumentsDeleted(vectorStore, bgDocIds);
153+
154+
List<Document> results = vectorStore
155+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
156+
157+
assertThat(results).hasSize(1);
158+
assertThat(normalizeValue(results.get(0).getMetadata().get("country"))).isEqualTo("NL");
159+
160+
vectorStore.delete(List.of(documents.get(1).getId()));
161+
});
162+
}
163+
164+
}

‎vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java‎

Lines changed: 25 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@
1919
import java.io.IOException;
2020
import java.nio.charset.StandardCharsets;
2121
import java.util.Collections;
22+
import java.util.HashMap;
2223
import java.util.List;
2324
import java.util.Map;
2425
import java.util.Optional;
2526
import java.util.UUID;
27+
import java.util.function.Consumer;
2628
import java.util.stream.Collectors;
2729

2830
import com.datastax.oss.driver.api.core.CqlSession;
@@ -40,8 +42,10 @@
4042
import org.springframework.ai.document.Document;
4143
import org.springframework.ai.document.DocumentMetadata;
4244
import org.springframework.ai.embedding.EmbeddingModel;
45+
import org.springframework.ai.test.vectorstore.BaseVectorStoreTests;
4346
import org.springframework.ai.transformers.TransformersEmbeddingModel;
4447
import org.springframework.ai.vectorstore.SearchRequest;
48+
import org.springframework.ai.vectorstore.VectorStore;
4549
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn;
4650
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumnTags;
4751
import org.springframework.ai.vectorstore.filter.Filter;
@@ -64,7 +68,7 @@
6468
* @since 1.0.0
6569
*/
6670
@Testcontainers
67-
class CassandraVectorStoreIT {
71+
class CassandraVectorStoreIT extends BaseVectorStoreTests {
6872

6973
@Container
7074
static CassandraContainer<?> cassandraContainer = new CassandraContainer<>(CassandraImage.DEFAULT_IMAGE);
@@ -110,6 +114,24 @@ private static CassandraVectorStore createTestStore(ApplicationContext context,
110114
return store;
111115
}
112116

117+
@Override
118+
protected void executeTest(Consumer<VectorStore> testFunction) {
119+
contextRunner.run(context -> {
120+
VectorStore vectorStore = context.getBean(VectorStore.class);
121+
testFunction.accept(vectorStore);
122+
});
123+
}
124+
125+
@Override
126+
protected Document createDocument(String country, Integer year) {
127+
Map<String, Object> metadata = new HashMap<>();
128+
metadata.put("country", country);
129+
if (year != null) {
130+
metadata.put("year", year.shortValue());
131+
}
132+
return new Document("The World is Big and Salvation Lurks Around the Corner", metadata);
133+
}
134+
113135
@Test
114136
void ensureBeanGetsCreated() {
115137
this.contextRunner.run(context -> {
@@ -422,39 +444,7 @@ void searchWithThreshold() {
422444
}
423445

424446
@Test
425-
void deleteById() {
426-
this.contextRunner.run(context -> {
427-
try (CassandraVectorStore store = createTestStore(context,
428-
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
429-
new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) {
430-
431-
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
432-
Map.of("country", "BG", "year", (short) 2020));
433-
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
434-
Map.of("country", "NL"));
435-
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
436-
Map.of("country", "BG", "year", (short) 2023));
437-
438-
store.add(List.of(bgDocument, nlDocument, bgDocument2));
439-
440-
// Verify initial state
441-
List<Document> results = store
442-
.similaritySearch(SearchRequest.builder().query("The World").topK(5).build());
443-
assertThat(results).hasSize(3);
444-
445-
store.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
446-
447-
results = store.similaritySearch(
448-
SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
449-
450-
assertThat(results).hasSize(1);
451-
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
452-
}
453-
});
454-
}
455-
456-
@Test
457-
void deleteByFilter() {
447+
protected void deleteByFilter() {
458448
this.contextRunner.run(context -> {
459449
try (CassandraVectorStore store = createTestStore(context,
460450
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
@@ -490,7 +480,7 @@ void deleteByFilter() {
490480
}
491481

492482
@Test
493-
void deleteWithStringFilterExpression() {
483+
protected void deleteWithStringFilterExpression() {
494484
this.contextRunner.run(context -> {
495485
try (CassandraVectorStore store = createTestStore(context,
496486
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),

0 commit comments

Comments
 (0)