Harden milvus IT test fixtures

spring-projects · Sep 28, 2023 · a45e34e · a45e34e
1 parent 5ff2c3b
commit a45e34e
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 17 deletions.
diff --git a/...g-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java b/...g-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java
@@ -48,8 +48,8 @@
 import io.milvus.param.index.CreateIndexParam;
 import io.milvus.param.index.DescribeIndexParam;
 import io.milvus.param.index.DropIndexParam;
-import io.milvus.response.SearchResultsWrapper;
 import io.milvus.response.QueryResultsWrapper.RowRecord;
+import io.milvus.response.SearchResultsWrapper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -359,7 +359,7 @@ public List<Document> similaritySearch(String query, int topK, double similarity
 				metadata.put(DISTANCE_FIELD_NAME, 1 - getResultSimilarity(rowRecord));
 				return new Document(docId, content, metadata.getInnerMap());
 			})
-			.collect(Collectors.toList());
+			.toList();
 	}
 
 	private float getResultSimilarity(RowRecord rowRecord) {
@@ -368,7 +368,7 @@ private float getResultSimilarity(RowRecord rowRecord) {
 	}
 
 	private List<Float> toFloatList(List<Double> embeddingDouble) {
-		return embeddingDouble.stream().map(Number::floatValue).collect(Collectors.toList());
+		return embeddingDouble.stream().map(Number::floatValue).toList();
 	}
 
 	// ---------------------------------------------------------------------------------

diff --git a/...ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java b/...ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java
@@ -21,18 +21,17 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.UUID;
-import java.util.stream.Collectors;
 
 import io.milvus.client.MilvusServiceClient;
 import io.milvus.param.ConnectParam;
 import io.milvus.param.IndexType;
 import io.milvus.param.MetricType;
 import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.testcontainers.containers.DockerComposeContainer;
 import org.testcontainers.containers.wait.strategy.Wait;
-import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
 
 import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
@@ -56,14 +55,9 @@
 @Testcontainers
 public class MilvusVectorStoreIT {
 
-	@Container
-	public static DockerComposeContainer milvusContainer = new DockerComposeContainer(
-			new File("src/test/resources/docker-compose.yml"))
-		.withExposedService("standalone", 19530)
-		.withExposedService("standalone", 9091,
-				Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
-		.waitingFor("standalone",
-				Wait.forLogMessage(".*Proxy successfully started.*\\s", 1).withStartupTimeout(Duration.ofSeconds(100)));
+	private static DockerComposeContainer milvusContainer;
+
+	private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString());
 
 	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
 		.withUserConfiguration(TestApplication.class)
@@ -77,9 +71,25 @@ public class MilvusVectorStoreIT {
 					"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
 					Collections.singletonMap("meta2", "meta2")));
 
+	@BeforeAll
+	public static void beforeAll() {
+		FileSystemUtils.deleteRecursively(TEMP_FOLDER);
+		TEMP_FOLDER.mkdirs();
+
+		milvusContainer = new DockerComposeContainer(new File("src/test/resources/docker-compose.yml"))
+			.withEnv("DOCKER_VOLUME_DIRECTORY", TEMP_FOLDER.getAbsolutePath())
+			.withExposedService("standalone", 19530)
+			.withExposedService("standalone", 9091,
+					Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
+			.waitingFor("standalone", Wait.forLogMessage(".*Proxy successfully started.*\\s", 1)
+				.withStartupTimeout(Duration.ofSeconds(100)));
+		milvusContainer.start();
+	}
+
 	@AfterAll
 	public static void afterAll() {
-		FileSystemUtils.deleteRecursively(new File("src/test/resources/volumes"));
+		milvusContainer.stop();
+		FileSystemUtils.deleteRecursively(TEMP_FOLDER);
 	}
 
 	private void resetCollection(VectorStore vectorStore) {
@@ -113,7 +123,7 @@ public void addAndSearchTest(String metricType) {
 				assertThat(resultDoc.getMetadata()).containsKey("distance");
 
 				// Remove all documents from the store
-				vectorStore.delete(documents.stream().map(doc -> doc.getId()).collect(Collectors.toList()));
+				vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
 
 				List<Document> results2 = vectorStore.similaritySearch("Hello", 1);
 				assertThat(results2).hasSize(0);
@@ -184,11 +194,11 @@ public void searchThresholdTest(String metricType) {
 
 				List<Float> distances = fullResult.stream()
 					.map(doc -> (Float) doc.getMetadata().get("distance"))
-					.collect(Collectors.toList());
+					.toList();
 
 				assertThat(distances).hasSize(3);
 
-				List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.01)));
+				List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.001)));
 
 				assertThat(results).hasSize(1);
 				Document resultDoc = results.get(0);