From 67c55cbcbedc4bccf1bcd92505334b4f4c01fb9a Mon Sep 17 00:00:00 2001
From: Laura Trotta <153528055+l-trotta@users.noreply.github.com>
Date: Tue, 10 Sep 2024 17:39:46 +0200
Subject: [PATCH] Added Elasticsearch Vector Store Unit Test (#1339)

* add and delete unit test
* format
---
 .../ElasticsearchVectorStoreIT.java           | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java
index 3832bb2f9..43cc435b9 100644
--- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java
+++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/ElasticsearchVectorStoreIT.java
@@ -27,6 +27,7 @@ import java.util.concurrent.TimeUnit;
 
 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.awaitility.Awaitility;
 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,38 @@ class ElasticsearchVectorStoreIT {
 		});
 	}
 
+	@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) {