Skip to content

Commit d43d6b0

Browse files
christophstroblmp911de
authored andcommitted
Introduce @EnableIfVectorSearchAvailable to wait and conditionally skip tests.
We now wait until a search index becomes available. If the search index doesn't come alive within 60 seconds, we skip that test (or test class). Closes: #5013 Original pull request: #5014
1 parent 8fc6df0 commit d43d6b0

File tree

12 files changed

+239
-55
lines changed

12 files changed

+239
-55
lines changed

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/aggregation/VectorSearchTests.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,8 +167,8 @@ static void initIndexes() {
167167
template.searchIndexOps(WithVectorFields.class).createIndex(rawIndex);
168168
template.searchIndexOps(WithVectorFields.class).createIndex(wrapperIndex);
169169

170-
template.awaitIndexCreation(WithVectorFields.class, rawIndex.getName());
171-
template.awaitIndexCreation(WithVectorFields.class, wrapperIndex.getName());
170+
template.awaitSearchIndexCreation(WithVectorFields.class, rawIndex.getName());
171+
template.awaitSearchIndexCreation(WithVectorFields.class, wrapperIndex.getName());
172172
}
173173

174174
private static void assertScoreIsDecreasing(Iterable<Document> documents) {

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

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,27 +15,29 @@
1515
*/
1616
package org.springframework.data.mongodb.core.index;
1717

18-
import static org.assertj.core.api.Assertions.*;
19-
import static org.awaitility.Awaitility.*;
18+
import static org.assertj.core.api.Assertions.assertThatRuntimeException;
19+
import static org.awaitility.Awaitility.await;
2020
import static org.springframework.data.mongodb.test.util.Assertions.assertThat;
2121

22+
import java.time.Duration;
2223
import java.util.List;
2324

2425
import org.bson.Document;
2526
import org.jspecify.annotations.Nullable;
2627
import org.junit.jupiter.api.AfterEach;
2728
import org.junit.jupiter.api.BeforeEach;
2829
import org.junit.jupiter.api.Test;
30+
import org.junit.jupiter.api.extension.ExtendWith;
2931
import org.junit.jupiter.params.ParameterizedTest;
3032
import org.junit.jupiter.params.provider.ValueSource;
31-
3233
import org.springframework.data.annotation.Id;
3334
import org.springframework.data.mongodb.core.index.VectorIndex.SimilarityFunction;
3435
import org.springframework.data.mongodb.core.mapping.Field;
3536
import org.springframework.data.mongodb.test.util.AtlasContainer;
37+
import org.springframework.data.mongodb.test.util.EnableIfVectorSearchAvailable;
38+
import org.springframework.data.mongodb.test.util.MongoServerCondition;
3639
import org.springframework.data.mongodb.test.util.MongoTestTemplate;
3740
import org.springframework.data.mongodb.test.util.MongoTestUtils;
38-
3941
import org.testcontainers.junit.jupiter.Container;
4042
import org.testcontainers.junit.jupiter.Testcontainers;
4143

@@ -48,6 +50,7 @@
4850
* @author Christoph Strobl
4951
* @author Mark Paluch
5052
*/
53+
@ExtendWith(MongoServerCondition.class)
5154
@Testcontainers(disabledWithoutDocker = true)
5255
class VectorIndexIntegrationTests {
5356

@@ -66,19 +69,22 @@ class VectorIndexIntegrationTests {
6669

6770
@BeforeEach
6871
void init() {
69-
template.createCollection(Movie.class);
72+
73+
template.createCollectionIfNotExists(Movie.class);
7074
indexOps = template.searchIndexOps(Movie.class);
7175
}
7276

7377
@AfterEach
7478
void cleanup() {
7579

80+
template.flush(Movie.class);
7681
template.searchIndexOps(Movie.class).dropAllIndexes();
77-
template.dropCollection(Movie.class);
82+
template.awaitNoSearchIndexAvailable(Movie.class, Duration.ofSeconds(30));
7883
}
7984

8085
@ParameterizedTest // GH-4706
8186
@ValueSource(strings = { "euclidean", "cosine", "dotProduct" })
87+
@EnableIfVectorSearchAvailable(collection = Movie.class)
8288
void createsSimpleVectorIndex(String similarityFunction) {
8389

8490
VectorIndex idx = new VectorIndex("vector_index").addVector("plotEmbedding",
@@ -98,21 +104,23 @@ void createsSimpleVectorIndex(String similarityFunction) {
98104
}
99105

100106
@Test // GH-4706
107+
@EnableIfVectorSearchAvailable(collection = Movie.class)
101108
void dropIndex() {
102109

103110
VectorIndex idx = new VectorIndex("vector_index").addVector("plotEmbedding",
104111
builder -> builder.dimensions(1536).similarity("cosine"));
105112

106113
indexOps.createIndex(idx);
107114

108-
template.awaitIndexCreation(Movie.class, idx.getName());
115+
template.awaitSearchIndexCreation(Movie.class, idx.getName());
109116

110117
indexOps.dropIndex(idx.getName());
111118

112119
assertThat(readRawIndexInfo(idx.getName())).isNull();
113120
}
114121

115122
@Test // GH-4706
123+
@EnableIfVectorSearchAvailable(collection = Movie.class)
116124
void statusChanges() throws InterruptedException {
117125

118126
String indexName = "vector_index";
@@ -131,6 +139,7 @@ void statusChanges() throws InterruptedException {
131139
}
132140

133141
@Test // GH-4706
142+
@EnableIfVectorSearchAvailable(collection = Movie.class)
134143
void exists() throws InterruptedException {
135144

136145
String indexName = "vector_index";
@@ -148,6 +157,7 @@ void exists() throws InterruptedException {
148157
}
149158

150159
@Test // GH-4706
160+
@EnableIfVectorSearchAvailable(collection = Movie.class)
151161
void updatesVectorIndex() throws InterruptedException {
152162

153163
String indexName = "vector_index";
@@ -177,6 +187,7 @@ void updatesVectorIndex() throws InterruptedException {
177187
}
178188

179189
@Test // GH-4706
190+
@EnableIfVectorSearchAvailable(collection = Movie.class)
180191
void createsVectorIndexWithFilters() throws InterruptedException {
181192

182193
VectorIndex idx = new VectorIndex("vector_index")

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/ReactiveVectorSearchTests.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -167,9 +167,9 @@ static void initIndexes() {
167167
template.searchIndexOps(WithVectorFields.class).createIndex(cosIndex);
168168
template.searchIndexOps(WithVectorFields.class).createIndex(euclideanIndex);
169169
template.searchIndexOps(WithVectorFields.class).createIndex(inner);
170-
template.awaitIndexCreation(WithVectorFields.class, cosIndex.getName());
171-
template.awaitIndexCreation(WithVectorFields.class, euclideanIndex.getName());
172-
template.awaitIndexCreation(WithVectorFields.class, inner.getName());
170+
template.awaitSearchIndexCreation(WithVectorFields.class, cosIndex.getName());
171+
template.awaitSearchIndexCreation(WithVectorFields.class, euclideanIndex.getName());
172+
template.awaitSearchIndexCreation(WithVectorFields.class, inner.getName());
173173
}
174174

175175
interface ReactiveVectorSearchRepository extends CrudRepository<WithVectorFields, String> {

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/repository/VectorSearchTests.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -211,9 +211,9 @@ static void initIndexes() {
211211
template.searchIndexOps(WithVectorFields.class).createIndex(cosIndex);
212212
template.searchIndexOps(WithVectorFields.class).createIndex(euclideanIndex);
213213
template.searchIndexOps(WithVectorFields.class).createIndex(inner);
214-
template.awaitIndexCreation(WithVectorFields.class, cosIndex.getName());
215-
template.awaitIndexCreation(WithVectorFields.class, euclideanIndex.getName());
216-
template.awaitIndexCreation(WithVectorFields.class, inner.getName());
214+
template.awaitSearchIndexCreation(WithVectorFields.class, cosIndex.getName());
215+
template.awaitSearchIndexCreation(WithVectorFields.class, euclideanIndex.getName());
216+
template.awaitSearchIndexCreation(WithVectorFields.class, inner.getName());
217217
}
218218

219219
interface VectorSearchRepository extends CrudRepository<WithVectorFields, String> {

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

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,14 @@
1616
package org.springframework.data.mongodb.test.util;
1717

1818
import org.springframework.core.env.StandardEnvironment;
19-
2019
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
2120
import org.testcontainers.utility.DockerImageName;
2221

22+
import com.github.dockerjava.api.command.InspectContainerResponse;
23+
2324
/**
24-
* Extension to MongoDBAtlasLocalContainer.
25+
* Extension to {@link MongoDBAtlasLocalContainer}. Registers mapped host an port as system properties
26+
* ({@link #ATLAS_HOST}, {@link #ATLAS_PORT}).
2527
*
2628
* @author Christoph Strobl
2729
*/
@@ -31,6 +33,9 @@ public class AtlasContainer extends MongoDBAtlasLocalContainer {
3133
private static final String DEFAULT_TAG = "8.0.0";
3234
private static final String LATEST = "latest";
3335

36+
public static final String ATLAS_HOST = "docker.mongodb.atlas.host";
37+
public static final String ATLAS_PORT = "docker.mongodb.atlas.port";
38+
3439
private AtlasContainer(String dockerImageName) {
3540
super(DockerImageName.parse(dockerImageName));
3641
}
@@ -55,4 +60,20 @@ public static AtlasContainer tagged(String tag) {
5560
return new AtlasContainer(DEFAULT_IMAGE_NAME.withTag(tag));
5661
}
5762

63+
@Override
64+
protected void containerIsStarted(InspectContainerResponse containerInfo) {
65+
66+
super.containerIsStarted(containerInfo);
67+
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);
78+
}
5879
}

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/MongoExtensions.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ static class Client {
3131
static final String REACTIVE_REPLSET_KEY = "mongo.client.replset.reactive";
3232
}
3333

34-
static class Termplate {
34+
static class Template {
3535

3636
static final Namespace NAMESPACE = Namespace.create(MongoTemplateExtension.class);
3737
static final String SYNC = "mongo.template.sync";

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

Lines changed: 63 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,21 @@
1515
*/
1616
package org.springframework.data.mongodb.test.util;
1717

18+
import java.time.Duration;
19+
1820
import org.junit.jupiter.api.extension.ConditionEvaluationResult;
1921
import org.junit.jupiter.api.extension.ExecutionCondition;
2022
import org.junit.jupiter.api.extension.ExtensionContext;
2123
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
2224
import org.springframework.core.annotation.AnnotatedElementUtils;
25+
import org.springframework.data.mongodb.MongoCollectionUtils;
2326
import org.springframework.data.util.Version;
27+
import org.springframework.util.NumberUtils;
28+
import org.springframework.util.StringUtils;
29+
import org.testcontainers.shaded.org.awaitility.Awaitility;
30+
31+
import com.mongodb.Function;
32+
import com.mongodb.client.MongoClient;
2433

2534
/**
2635
* @author Christoph Strobl
@@ -42,10 +51,13 @@ public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext con
4251
}
4352
}
4453

45-
if(context.getTags().contains("vector-search")) {
46-
if(!atlasEnvironment(context)) {
54+
if (context.getTags().contains("vector-search")) {
55+
if (!atlasEnvironment(context)) {
4756
return ConditionEvaluationResult.disabled("Disabled for servers not supporting Vector Search.");
4857
}
58+
if (!isSearchIndexAvailable(context)) {
59+
return ConditionEvaluationResult.disabled("Search index unavailable.");
60+
}
4961
}
5062

5163
if (context.getTags().contains("version-specific") && context.getElement().isPresent()) {
@@ -90,8 +102,55 @@ private Version serverVersion(ExtensionContext context) {
90102
Version.class);
91103
}
92104

105+
private boolean isSearchIndexAvailable(ExtensionContext context) {
106+
107+
EnableIfVectorSearchAvailable vectorSearchAvailable = AnnotatedElementUtils
108+
.findMergedAnnotation(context.getElement().get(), EnableIfVectorSearchAvailable.class);
109+
110+
if (vectorSearchAvailable == null) {
111+
return true;
112+
}
113+
114+
String collectionName = StringUtils.hasText(vectorSearchAvailable.collectionName())
115+
? vectorSearchAvailable.collectionName()
116+
: MongoCollectionUtils.getPreferredCollectionName(vectorSearchAvailable.collection());
117+
118+
return context.getStore(NAMESPACE).getOrComputeIfAbsent("search-index-%s-available".formatted(collectionName),
119+
(key) -> {
120+
try {
121+
doWithClient(client -> {
122+
Awaitility.await().atMost(Duration.ofSeconds(60)).pollInterval(Duration.ofMillis(200)).until(() -> {
123+
return MongoTestUtils.isSearchIndexReady(client, null, collectionName);
124+
});
125+
return "done waiting for search index";
126+
});
127+
} catch (Exception e) {
128+
return false;
129+
}
130+
return true;
131+
}, Boolean.class);
132+
133+
}
134+
93135
private boolean atlasEnvironment(ExtensionContext context) {
94-
return context.getStore(NAMESPACE).getOrComputeIfAbsent(Version.class, (key) -> MongoTestUtils.isVectorSearchEnabled(),
95-
Boolean.class);
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);
149+
}
150+
}
151+
152+
try (MongoClient client = MongoTestUtils.client()) {
153+
return function.apply(client);
154+
}
96155
}
97156
}

0 commit comments

Comments
 (0)