Skip to content

Commit a45e34e

Browse files
committedSep 28, 2023
Harden milvus IT test fixtures
1 parent 5ff2c3b commit a45e34e

File tree

2 files changed

+27
-17
lines changed

2 files changed

+27
-17
lines changed
 

‎vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,8 @@
4848
import io.milvus.param.index.CreateIndexParam;
4949
import io.milvus.param.index.DescribeIndexParam;
5050
import io.milvus.param.index.DropIndexParam;
51-
import io.milvus.response.SearchResultsWrapper;
5251
import io.milvus.response.QueryResultsWrapper.RowRecord;
52+
import io.milvus.response.SearchResultsWrapper;
5353
import org.slf4j.Logger;
5454
import org.slf4j.LoggerFactory;
5555

@@ -359,7 +359,7 @@ public List<Document> similaritySearch(String query, int topK, double similarity
359359
metadata.put(DISTANCE_FIELD_NAME, 1 - getResultSimilarity(rowRecord));
360360
return new Document(docId, content, metadata.getInnerMap());
361361
})
362-
.collect(Collectors.toList());
362+
.toList();
363363
}
364364

365365
private float getResultSimilarity(RowRecord rowRecord) {
@@ -368,7 +368,7 @@ private float getResultSimilarity(RowRecord rowRecord) {
368368
}
369369

370370
private List<Float> toFloatList(List<Double> embeddingDouble) {
371-
return embeddingDouble.stream().map(Number::floatValue).collect(Collectors.toList());
371+
return embeddingDouble.stream().map(Number::floatValue).toList();
372372
}
373373

374374
// ---------------------------------------------------------------------------------

‎vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,17 @@
2121
import java.util.Collections;
2222
import java.util.List;
2323
import java.util.UUID;
24-
import java.util.stream.Collectors;
2524

2625
import io.milvus.client.MilvusServiceClient;
2726
import io.milvus.param.ConnectParam;
2827
import io.milvus.param.IndexType;
2928
import io.milvus.param.MetricType;
3029
import org.junit.jupiter.api.AfterAll;
30+
import org.junit.jupiter.api.BeforeAll;
3131
import org.junit.jupiter.params.ParameterizedTest;
3232
import org.junit.jupiter.params.provider.ValueSource;
3333
import org.testcontainers.containers.DockerComposeContainer;
3434
import org.testcontainers.containers.wait.strategy.Wait;
35-
import org.testcontainers.junit.jupiter.Container;
3635
import org.testcontainers.junit.jupiter.Testcontainers;
3736

3837
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
@@ -56,14 +55,9 @@
5655
@Testcontainers
5756
public class MilvusVectorStoreIT {
5857

59-
@Container
60-
public static DockerComposeContainer milvusContainer = new DockerComposeContainer(
61-
new File("src/test/resources/docker-compose.yml"))
62-
.withExposedService("standalone", 19530)
63-
.withExposedService("standalone", 9091,
64-
Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
65-
.waitingFor("standalone",
66-
Wait.forLogMessage(".*Proxy successfully started.*\\s", 1).withStartupTimeout(Duration.ofSeconds(100)));
58+
private static DockerComposeContainer milvusContainer;
59+
60+
private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString());
6761

6862
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
6963
.withUserConfiguration(TestApplication.class)
@@ -77,9 +71,25 @@ public class MilvusVectorStoreIT {
7771
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
7872
Collections.singletonMap("meta2", "meta2")));
7973

74+
@BeforeAll
75+
public static void beforeAll() {
76+
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
77+
TEMP_FOLDER.mkdirs();
78+
79+
milvusContainer = new DockerComposeContainer(new File("src/test/resources/docker-compose.yml"))
80+
.withEnv("DOCKER_VOLUME_DIRECTORY", TEMP_FOLDER.getAbsolutePath())
81+
.withExposedService("standalone", 19530)
82+
.withExposedService("standalone", 9091,
83+
Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
84+
.waitingFor("standalone", Wait.forLogMessage(".*Proxy successfully started.*\\s", 1)
85+
.withStartupTimeout(Duration.ofSeconds(100)));
86+
milvusContainer.start();
87+
}
88+
8089
@AfterAll
8190
public static void afterAll() {
82-
FileSystemUtils.deleteRecursively(new File("src/test/resources/volumes"));
91+
milvusContainer.stop();
92+
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
8393
}
8494

8595
private void resetCollection(VectorStore vectorStore) {
@@ -113,7 +123,7 @@ public void addAndSearchTest(String metricType) {
113123
assertThat(resultDoc.getMetadata()).containsKey("distance");
114124

115125
// Remove all documents from the store
116-
vectorStore.delete(documents.stream().map(doc -> doc.getId()).collect(Collectors.toList()));
126+
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
117127

118128
List<Document> results2 = vectorStore.similaritySearch("Hello", 1);
119129
assertThat(results2).hasSize(0);
@@ -184,11 +194,11 @@ public void searchThresholdTest(String metricType) {
184194

185195
List<Float> distances = fullResult.stream()
186196
.map(doc -> (Float) doc.getMetadata().get("distance"))
187-
.collect(Collectors.toList());
197+
.toList();
188198

189199
assertThat(distances).hasSize(3);
190200

191-
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.01)));
201+
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.001)));
192202

193203
assertThat(results).hasSize(1);
194204
Document resultDoc = results.get(0);

0 commit comments

Comments
 (0)
Please sign in to comment.