Skip to content

Commit 4d692a5

Browse files
sobychackoilayaperumalg
sobychacko
authored and
ilayaperumalg
committed
Add missing integration tests for delete by ID API in vector store implementations.
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. Adding javadoc Signed-off-by: Soby Chacko <[email protected]>
1 parent a8e305d commit 4d692a5

File tree

17 files changed

+379
-867
lines changed

17 files changed

+379
-867
lines changed

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

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
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+
/**
46+
* Execute a test function with a configured VectorStore instance. This method is
47+
* responsible for providing a properly initialized VectorStore within the appropriate
48+
* Spring application context for testing.
49+
* @param testFunction the consumer that executes test operations on the VectorStore
50+
*/
51+
protected abstract void executeTest(Consumer<VectorStore> testFunction);
52+
53+
protected Document createDocument(String country, Integer year) {
54+
Map<String, Object> metadata = new HashMap<>();
55+
metadata.put("country", country);
56+
if (year != null) {
57+
metadata.put("year", year);
58+
}
59+
return new Document("The World is Big and Salvation Lurks Around the Corner", metadata);
60+
}
61+
62+
protected List<Document> setupTestDocuments(VectorStore vectorStore) {
63+
var doc1 = createDocument("BG", 2020);
64+
var doc2 = createDocument("NL", null);
65+
var doc3 = createDocument("BG", 2023);
66+
67+
List<Document> documents = List.of(doc1, doc2, doc3);
68+
vectorStore.add(documents);
69+
70+
return documents;
71+
}
72+
73+
private String normalizeValue(Object value) {
74+
return value.toString().replaceAll("^\"|\"$", "").trim();
75+
}
76+
77+
private void verifyDocumentsExist(VectorStore vectorStore, List<Document> documents) {
78+
await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> {
79+
List<Document> results = vectorStore.similaritySearch(
80+
SearchRequest.builder().query("The World").topK(documents.size()).similarityThresholdAll().build());
81+
assertThat(results).hasSize(documents.size());
82+
});
83+
}
84+
85+
private void verifyDocumentsDeleted(VectorStore vectorStore, List<String> deletedIds) {
86+
await().atMost(5, TimeUnit.SECONDS).pollInterval(Duration.ofMillis(500)).untilAsserted(() -> {
87+
List<Document> results = vectorStore
88+
.similaritySearch(SearchRequest.builder().query("The World").topK(10).similarityThresholdAll().build());
89+
90+
List<String> foundIds = results.stream().map(Document::getId).collect(Collectors.toList());
91+
92+
assertThat(foundIds).doesNotContainAnyElementsOf(deletedIds);
93+
});
94+
}
95+
96+
@Test
97+
protected void deleteById() {
98+
executeTest(vectorStore -> {
99+
List<Document> documents = setupTestDocuments(vectorStore);
100+
verifyDocumentsExist(vectorStore, documents);
101+
102+
List<String> idsToDelete = List.of(documents.get(0).getId(), documents.get(1).getId());
103+
vectorStore.delete(idsToDelete);
104+
verifyDocumentsDeleted(vectorStore, idsToDelete);
105+
106+
List<Document> results = vectorStore
107+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
108+
109+
assertThat(results).hasSize(1);
110+
assertThat(results.get(0).getId()).isEqualTo(documents.get(2).getId());
111+
Map<String, Object> metadata = results.get(0).getMetadata();
112+
assertThat(normalizeValue(metadata.get("country"))).isEqualTo("BG");
113+
assertThat(normalizeValue(metadata.get("year"))).isEqualTo("2023");
114+
115+
vectorStore.delete(List.of(documents.get(2).getId()));
116+
});
117+
}
118+
119+
@Test
120+
protected void deleteWithStringFilterExpression() {
121+
executeTest(vectorStore -> {
122+
List<Document> documents = setupTestDocuments(vectorStore);
123+
verifyDocumentsExist(vectorStore, documents);
124+
125+
List<String> bgDocIds = documents.stream()
126+
.filter(d -> "BG".equals(d.getMetadata().get("country")))
127+
.map(Document::getId)
128+
.collect(Collectors.toList());
129+
130+
vectorStore.delete("country == 'BG'");
131+
verifyDocumentsDeleted(vectorStore, bgDocIds);
132+
133+
List<Document> results = vectorStore
134+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
135+
136+
assertThat(results).hasSize(1);
137+
assertThat(normalizeValue(results.get(0).getMetadata().get("country"))).isEqualTo("NL");
138+
139+
vectorStore.delete(List.of(documents.get(1).getId()));
140+
});
141+
}
142+
143+
@Test
144+
protected void deleteByFilter() {
145+
executeTest(vectorStore -> {
146+
List<Document> documents = setupTestDocuments(vectorStore);
147+
verifyDocumentsExist(vectorStore, documents);
148+
149+
List<String> bgDocIds = documents.stream()
150+
.filter(d -> "BG".equals(d.getMetadata().get("country")))
151+
.map(Document::getId)
152+
.collect(Collectors.toList());
153+
154+
Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
155+
new Filter.Key("country"), new Filter.Value("BG"));
156+
157+
vectorStore.delete(filterExpression);
158+
verifyDocumentsDeleted(vectorStore, bgDocIds);
159+
160+
List<Document> results = vectorStore
161+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
162+
163+
assertThat(results).hasSize(1);
164+
assertThat(normalizeValue(results.get(0).getMetadata().get("country"))).isEqualTo("NL");
165+
166+
vectorStore.delete(List.of(documents.get(1).getId()));
167+
});
168+
}
169+
170+
}

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

Lines changed: 25 additions & 3 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,7 +444,7 @@ void searchWithThreshold() {
422444
}
423445

424446
@Test
425-
void deleteByFilter() {
447+
protected void deleteByFilter() {
426448
this.contextRunner.run(context -> {
427449
try (CassandraVectorStore store = createTestStore(context,
428450
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
@@ -458,7 +480,7 @@ void deleteByFilter() {
458480
}
459481

460482
@Test
461-
void deleteWithStringFilterExpression() {
483+
protected void deleteWithStringFilterExpression() {
462484
this.contextRunner.run(context -> {
463485
try (CassandraVectorStore store = createTestStore(context,
464486
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),

‎vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java‎

Lines changed: 11 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.List;
2121
import java.util.Map;
2222
import java.util.UUID;
23+
import java.util.function.Consumer;
2324

2425
import org.junit.jupiter.api.Test;
2526
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -33,6 +34,7 @@
3334
import org.springframework.ai.embedding.EmbeddingModel;
3435
import org.springframework.ai.openai.OpenAiEmbeddingModel;
3536
import org.springframework.ai.openai.api.OpenAiApi;
37+
import org.springframework.ai.test.vectorstore.BaseVectorStoreTests;
3638
import org.springframework.ai.vectorstore.SearchRequest;
3739
import org.springframework.ai.vectorstore.VectorStore;
3840
import org.springframework.ai.vectorstore.filter.Filter;
@@ -51,7 +53,7 @@
5153
*/
5254
@Testcontainers
5355
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
54-
public class ChromaVectorStoreIT {
56+
public class ChromaVectorStoreIT extends BaseVectorStoreTests {
5557

5658
@Container
5759
static ChromaDBContainer chromaContainer = new ChromaDBContainer(ChromaImage.DEFAULT_IMAGE);
@@ -68,6 +70,14 @@ public class ChromaVectorStoreIT {
6870
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
6971
Collections.singletonMap("meta2", "meta2")));
7072

73+
@Override
74+
protected void executeTest(Consumer<VectorStore> testFunction) {
75+
contextRunner.run(context -> {
76+
VectorStore vectorStore = context.getBean(VectorStore.class);
77+
testFunction.accept(vectorStore);
78+
});
79+
}
80+
7181
@Test
7282
public void addAndSearch() {
7383
this.contextRunner.run(context -> {
@@ -168,69 +178,6 @@ public void addAndSearchWithFilters() {
168178
});
169179
}
170180

171-
@Test
172-
public void deleteWithFilterExpression() {
173-
this.contextRunner.run(context -> {
174-
VectorStore vectorStore = context.getBean(VectorStore.class);
175-
176-
// Create test documents with different metadata
177-
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
178-
Map.of("country", "Bulgaria"));
179-
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
180-
Map.of("country", "Netherlands"));
181-
182-
// Add documents to the store
183-
vectorStore.add(List.of(bgDocument, nlDocument));
184-
185-
// Verify initial state
186-
var request = SearchRequest.builder().query("The World").topK(5).build();
187-
List<Document> results = vectorStore.similaritySearch(request);
188-
assertThat(results).hasSize(2);
189-
190-
// Delete document with country = Bulgaria
191-
Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
192-
new Filter.Key("country"), new Filter.Value("Bulgaria"));
193-
194-
vectorStore.delete(filterExpression);
195-
196-
// Verify Bulgaria document was deleted
197-
results = vectorStore
198-
.similaritySearch(SearchRequest.from(request).filterExpression("country == 'Bulgaria'").build());
199-
assertThat(results).isEmpty();
200-
201-
// Verify Netherlands document still exists
202-
results = vectorStore
203-
.similaritySearch(SearchRequest.from(request).filterExpression("country == 'Netherlands'").build());
204-
assertThat(results).hasSize(1);
205-
assertThat(results.get(0).getMetadata().get("country")).isEqualTo("Netherlands");
206-
207-
// Clean up
208-
vectorStore.delete(List.of(nlDocument.getId()));
209-
});
210-
}
211-
212-
@Test
213-
public void deleteWithStringFilterExpression() {
214-
this.contextRunner.run(context -> {
215-
VectorStore vectorStore = context.getBean(VectorStore.class);
216-
217-
var bgDocument = new Document("The World is Big", Map.of("country", "Bulgaria"));
218-
var nlDocument = new Document("The World is Big", Map.of("country", "Netherlands"));
219-
vectorStore.add(List.of(bgDocument, nlDocument));
220-
221-
var request = SearchRequest.builder().query("World").topK(5).build();
222-
assertThat(vectorStore.similaritySearch(request)).hasSize(2);
223-
224-
vectorStore.delete("country == 'Bulgaria'");
225-
226-
var results = vectorStore.similaritySearch(request);
227-
assertThat(results).hasSize(1);
228-
assertThat(results.get(0).getMetadata().get("country")).isEqualTo("Netherlands");
229-
230-
vectorStore.delete(List.of(nlDocument.getId()));
231-
});
232-
}
233-
234181
@Test
235182
public void documentUpdateTest() {
236183

0 commit comments

Comments
 (0)