add and delete unit test

l-trotta · Sep 10, 2024 · 692ed87 · 692ed87
1 parent b22a577
commit 692ed87
Showing 1 changed file with 25 additions and 0 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
@@ -27,6 +27,7 @@
 
 import co.elastic.clients.elasticsearch.ElasticsearchClient;
 import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord;
+import co.elastic.clients.elasticsearch.indices.stats.IndicesStats;
 import co.elastic.clients.json.jackson.JacksonJsonpMapper;
 import co.elastic.clients.transport.rest_client.RestClientTransport;
 import com.fasterxml.jackson.databind.DeserializationFeature;
@@ -36,6 +37,7 @@
 import org.elasticsearch.client.RestClient;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -105,6 +107,29 @@ void cleanDatabase() {
 		});
 	}
 
+	@Test
+	public void addAndDeleteDocumentsTest(){
+		getContextRunner().run(context -> {
+			ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
+					ElasticsearchVectorStore.class);
+			ElasticsearchClient elasticsearchClient = context.getBean(ElasticsearchClient.class);
+
+			IndicesStats stats = elasticsearchClient.indices().stats(s -> s.index("spring-ai-document-index")).indices().get("spring-ai-document-index");
+
+			assertThat(stats.total().docs().count()).isEqualTo(0L);
+
+			vectorStore.add(documents);
+			elasticsearchClient.indices().refresh();
+			stats = elasticsearchClient.indices().stats(s -> s.index("spring-ai-document-index")).indices().get("spring-ai-document-index");
+			assertThat(stats.total().docs().count()).isEqualTo(3L);
+
+			vectorStore.doDelete(List.of("1", "2", "3"));
+			elasticsearchClient.indices().refresh();
+			stats = elasticsearchClient.indices().stats(s -> s.index("spring-ai-document-index")).indices().get("spring-ai-document-index");
+			assertThat(stats.total().docs().count()).isEqualTo(0L);
+		});
+	}
+
 	@ParameterizedTest(name = "{0} : {displayName} ")
 	@ValueSource(strings = { "cosine", "l2_norm", "dot_product" })
 	public void addAndSearchTest(String similarityFunction) {