Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Added Elasticsearch Vector Store Unit Test (#1339)

Browse Source 
* add and delete unit test
* format
This commit is contained in:
Laura Trotta
2024-09-10 17:39:46 +02:00
committed by GitHub
parent c658c7b52e
commit 67c55cbcbe
1 changed files with 34 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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) {