Skip to content

Commit 8b5607c

Browse files
christophstroblchristophstrobl
committed
Make use if EnableIfVectorSearchAvailable
1 parent fbbb099 commit 8b5607c

File tree

4 files changed

+85
-34
lines changed

4 files changed

+85
-34
lines changed

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/index/VectorIndexIntegrationTests.java

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import org.jspecify.annotations.Nullable;
2626
import org.junit.jupiter.api.AfterEach;
2727
import org.junit.jupiter.api.BeforeEach;
28-
import org.junit.jupiter.api.Tag;
2928
import org.junit.jupiter.api.Test;
3029
import org.junit.jupiter.api.extension.ExtendWith;
3130
import org.junit.jupiter.params.ParameterizedTest;
@@ -34,6 +33,7 @@
3433
import org.springframework.data.mongodb.core.index.VectorIndex.SimilarityFunction;
3534
import org.springframework.data.mongodb.core.mapping.Field;
3635
import org.springframework.data.mongodb.test.util.AtlasContainer;
36+
import org.springframework.data.mongodb.test.util.EnableIfVectorSearchAvailable;
3737
import org.springframework.data.mongodb.test.util.MongoServerCondition;
3838
import org.springframework.data.mongodb.test.util.MongoTestTemplate;
3939
import org.springframework.data.mongodb.test.util.MongoTestUtils;
@@ -51,7 +51,6 @@
5151
*/
5252
@ExtendWith(MongoServerCondition.class)
5353
@Testcontainers(disabledWithoutDocker = true)
54-
@Tag("collection:movie")
5554
class VectorIndexIntegrationTests {
5655

5756
private static final @Container AtlasContainer atlasLocal = AtlasContainer.bestMatch();
@@ -81,9 +80,9 @@ void cleanup() {
8180
template.dropCollection(Movie.class);
8281
}
8382

84-
@Tag("vector-search")
8583
@ParameterizedTest // GH-4706
8684
@ValueSource(strings = { "euclidean", "cosine", "dotProduct" })
85+
@EnableIfVectorSearchAvailable(collection = Movie.class)
8786
void createsSimpleVectorIndex(String similarityFunction) {
8887

8988
VectorIndex idx = new VectorIndex("vector_index").addVector("plotEmbedding",
@@ -102,8 +101,8 @@ void createsSimpleVectorIndex(String similarityFunction) {
102101
});
103102
}
104103

105-
@Tag("vector-search")
106104
@Test // GH-4706
105+
@EnableIfVectorSearchAvailable(collection = Movie.class)
107106
void dropIndex() {
108107

109108
VectorIndex idx = new VectorIndex("vector_index").addVector("plotEmbedding",
@@ -118,8 +117,8 @@ void dropIndex() {
118117
assertThat(readRawIndexInfo(idx.getName())).isNull();
119118
}
120119

121-
@Tag("vector-search")
122120
@Test // GH-4706
121+
@EnableIfVectorSearchAvailable(collection = Movie.class)
123122
void statusChanges() throws InterruptedException {
124123

125124
String indexName = "vector_index";
@@ -137,8 +136,8 @@ void statusChanges() throws InterruptedException {
137136
SearchIndexStatus.READY);
138137
}
139138

140-
@Tag("vector-search")
141139
@Test // GH-4706
140+
@EnableIfVectorSearchAvailable(collection = Movie.class)
142141
void exists() throws InterruptedException {
143142

144143
String indexName = "vector_index";
@@ -155,8 +154,8 @@ void exists() throws InterruptedException {
155154
assertThat(indexOps.exists(indexName)).isTrue();
156155
}
157156

158-
@Tag("vector-search")
159157
@Test // GH-4706
158+
@EnableIfVectorSearchAvailable(collection = Movie.class)
160159
void updatesVectorIndex() throws InterruptedException {
161160

162161
String indexName = "vector_index";
@@ -185,8 +184,8 @@ void updatesVectorIndex() throws InterruptedException {
185184
assertThatRuntimeException().isThrownBy(() -> indexOps.updateIndex(updatedIdx));
186185
}
187186

188-
@Tag("vector-search")
189187
@Test // GH-4706
188+
@EnableIfVectorSearchAvailable(collection = Movie.class)
190189
void createsVectorIndexWithFilters() throws InterruptedException {
191190

192191
VectorIndex idx = new VectorIndex("vector_index")

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/AtlasContainer.java

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@
2222
import com.github.dockerjava.api.command.InspectContainerResponse;
2323

2424
/**
25-
* Extension to MongoDBAtlasLocalContainer.
25+
* Extension to {@link MongoDBAtlasLocalContainer}. Registers mapped host an port as system properties
26+
* ({@link #ATLAS_HOST}, {@link #ATLAS_PORT}).
2627
*
2728
* @author Christoph Strobl
2829
*/
@@ -32,6 +33,9 @@ public class AtlasContainer extends MongoDBAtlasLocalContainer {
3233
private static final String DEFAULT_TAG = "8.0.0";
3334
private static final String LATEST = "latest";
3435

36+
public static final String ATLAS_HOST = "docker.mongodb.atlas.host";
37+
public static final String ATLAS_PORT = "docker.mongodb.atlas.port";
38+
3539
private AtlasContainer(String dockerImageName) {
3640
super(DockerImageName.parse(dockerImageName));
3741
}
@@ -61,7 +65,15 @@ protected void containerIsStarted(InspectContainerResponse containerInfo) {
6165

6266
super.containerIsStarted(containerInfo);
6367

64-
System.setProperty("docker.mongodb.atlas.host", getHost());
65-
System.setProperty("docker.mongodb.atlas.port", getMappedPort(27017).toString());
68+
System.setProperty(ATLAS_HOST, getHost());
69+
System.setProperty(ATLAS_PORT, getMappedPort(27017).toString());
70+
}
71+
72+
@Override
73+
protected void containerIsStopping(InspectContainerResponse containerInfo) {
74+
75+
System.clearProperty(ATLAS_HOST);
76+
System.clearProperty(ATLAS_PORT);
77+
super.containerIsStopping(containerInfo);
6678
}
6779
}

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/EnableIfVectorSearchAvailable.java

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,30 @@
2525
import org.junit.jupiter.api.extension.ExtendWith;
2626

2727
/**
28+
* {@link EnableIfVectorSearchAvailable} indicates a specific method can only be run in an environment that has a search
29+
* server available. This means that not only the mongodb instance needs to have a
30+
* {@literal searchIndexManagementHostAndPort} configured, but also that the search index sever is actually up and
31+
* running, responding to a {@literal $listSearchIndexes} aggregation.
32+
*
2833
* @author Christoph Strobl
34+
* @since 5.0
35+
* @see Tag
2936
*/
30-
@Target({ ElementType.TYPE, ElementType.METHOD })
37+
@Target({ ElementType.METHOD })
3138
@Retention(RetentionPolicy.RUNTIME)
3239
@Documented
3340
@Tag("vector-search")
3441
@ExtendWith(MongoServerCondition.class)
3542
public @interface EnableIfVectorSearchAvailable {
3643

44+
/**
45+
* @return the name of the collection used to run the {@literal $listSearchIndexes} aggregation.
46+
*/
47+
String collectionName() default "";
48+
49+
/**
50+
* @return the type for resolving the name of the collection used to run the {@literal $listSearchIndexes}
51+
* aggregation. The {@link #collectionName()} has precedence over the type.
52+
*/
53+
Class<?> collection() default Object.class;
3754
}

‎spring-data-mongodb/src/test/java/org/springframework/data/mongodb/test/util/MongoServerCondition.java

Lines changed: 45 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -16,18 +16,19 @@
1616
package org.springframework.data.mongodb.test.util;
1717

1818
import java.time.Duration;
19-
import java.util.Optional;
2019

2120
import org.junit.jupiter.api.extension.ConditionEvaluationResult;
2221
import org.junit.jupiter.api.extension.ExecutionCondition;
2322
import org.junit.jupiter.api.extension.ExtensionContext;
2423
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
2524
import org.springframework.core.annotation.AnnotatedElementUtils;
25+
import org.springframework.data.mongodb.MongoCollectionUtils;
2626
import org.springframework.data.util.Version;
2727
import org.springframework.util.NumberUtils;
2828
import org.springframework.util.StringUtils;
2929
import org.testcontainers.shaded.org.awaitility.Awaitility;
3030

31+
import com.mongodb.Function;
3132
import com.mongodb.client.MongoClient;
3233

3334
/**
@@ -54,6 +55,9 @@ public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext con
5455
if (!atlasEnvironment(context)) {
5556
return ConditionEvaluationResult.disabled("Disabled for servers not supporting Vector Search.");
5657
}
58+
if (!isSearchIndexAvailable(context)) {
59+
return ConditionEvaluationResult.disabled("Search index unavailable.");
60+
}
5761
}
5862

5963
if (context.getTags().contains("version-specific") && context.getElement().isPresent()) {
@@ -98,36 +102,55 @@ private Version serverVersion(ExtensionContext context) {
98102
Version.class);
99103
}
100104

101-
private boolean atlasEnvironment(ExtensionContext context) {
102-
103-
String host = System.getProperty("docker.mongodb.atlas.host");
104-
String port = System.getProperty("docker.mongodb.atlas.port");
105+
private boolean isSearchIndexAvailable(ExtensionContext context) {
105106

106-
return context.getStore(NAMESPACE).getOrComputeIfAbsent(Version.class, (key) -> {
107+
EnableIfVectorSearchAvailable vectorSearchAvailable = AnnotatedElementUtils
108+
.findMergedAnnotation(context.getElement().get(), EnableIfVectorSearchAvailable.class);
107109

108-
if (StringUtils.hasText(host) && StringUtils.hasText(port)) {
109-
try (MongoClient client = MongoTestUtils.client(host, NumberUtils.parseNumber(port, Integer.class))) {
110-
if (!MongoTestUtils.isVectorSearchEnabled(client)) {
111-
return false;
112-
}
110+
if (vectorSearchAvailable == null) {
111+
return true;
112+
}
113113

114-
Optional<String> collectionTag = context.getTags().stream().filter(tag -> tag.startsWith("collection"))
115-
.findFirst();
116-
if (collectionTag.isPresent()) {
114+
String collectionName = StringUtils.hasText(vectorSearchAvailable.collectionName())
115+
? vectorSearchAvailable.collectionName()
116+
: MongoCollectionUtils.getPreferredCollectionName(vectorSearchAvailable.collection());
117117

118-
String collectionName = collectionTag.get().split(":")[1];
119-
try {
118+
return context.getStore(NAMESPACE).getOrComputeIfAbsent("search-index-%s-available".formatted(collectionName),
119+
(key) -> {
120+
try {
121+
doWithClient(client -> {
120122
Awaitility.await().atMost(Duration.ofSeconds(60)).pollInterval(Duration.ofMillis(200)).until(() -> {
121123
return MongoTestUtils.isSearchIndexReady(client, null, collectionName);
122124
});
123-
} catch (Exception e) {
124-
return false;
125-
}
125+
return "done waiting for search index";
126+
});
127+
} catch (Exception e) {
128+
return false;
126129
}
127130
return true;
128-
}
131+
}, Boolean.class);
132+
133+
}
134+
135+
private boolean atlasEnvironment(ExtensionContext context) {
136+
137+
return context.getStore(NAMESPACE).getOrComputeIfAbsent("mongodb-atlas",
138+
(key) -> doWithClient(MongoTestUtils::isVectorSearchEnabled), Boolean.class);
139+
}
140+
141+
private <T> T doWithClient(Function<MongoClient, T> function) {
142+
143+
String host = System.getProperty(AtlasContainer.ATLAS_HOST);
144+
String port = System.getProperty(AtlasContainer.ATLAS_PORT);
145+
146+
if (StringUtils.hasText(host) && StringUtils.hasText(port)) {
147+
try (MongoClient client = MongoTestUtils.client(host, NumberUtils.parseNumber(port, Integer.class))) {
148+
return function.apply(client);
129149
}
130-
return MongoTestUtils.isVectorSearchEnabled();
131-
}, Boolean.class);
150+
}
151+
152+
try (MongoClient client = MongoTestUtils.client()) {
153+
return function.apply(client);
154+
}
132155
}
133156
}

0 commit comments

Comments
 (0)