Skip to content

Commit d2c2fca

Browse files
christophstroblchristophstrobl
committed
Use vector-serach test tag to indicate potential slow environment
1 parent 01cd906 commit d2c2fca

File tree

4 files changed

+68
-25
lines changed

4 files changed

+68
-25
lines changed

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

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
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

2222
import java.util.List;
@@ -25,17 +25,18 @@
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;
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.MongoServerCondition;
3638
import org.springframework.data.mongodb.test.util.MongoTestTemplate;
3739
import org.springframework.data.mongodb.test.util.MongoTestUtils;
38-
3940
import org.testcontainers.junit.jupiter.Container;
4041
import org.testcontainers.junit.jupiter.Testcontainers;
4142

@@ -48,6 +49,7 @@
4849
* @author Christoph Strobl
4950
* @author Mark Paluch
5051
*/
52+
@ExtendWith(MongoServerCondition.class)
5153
@Testcontainers(disabledWithoutDocker = true)
5254
class VectorIndexIntegrationTests {
5355

@@ -66,6 +68,7 @@ class VectorIndexIntegrationTests {
6668

6769
@BeforeEach
6870
void init() {
71+
6972
template.createCollection(Movie.class);
7073
indexOps = template.searchIndexOps(Movie.class);
7174
}
@@ -77,6 +80,7 @@ void cleanup() {
7780
template.dropCollection(Movie.class);
7881
}
7982

83+
@Tag("vector-search")
8084
@ParameterizedTest // GH-4706
8185
@ValueSource(strings = { "euclidean", "cosine", "dotProduct" })
8286
void createsSimpleVectorIndex(String similarityFunction) {
@@ -97,6 +101,7 @@ void createsSimpleVectorIndex(String similarityFunction) {
97101
});
98102
}
99103

104+
@Tag("vector-search")
100105
@Test // GH-4706
101106
void dropIndex() {
102107

@@ -112,6 +117,7 @@ void dropIndex() {
112117
assertThat(readRawIndexInfo(idx.getName())).isNull();
113118
}
114119

120+
@Tag("vector-search")
115121
@Test // GH-4706
116122
void statusChanges() throws InterruptedException {
117123

@@ -130,6 +136,7 @@ void statusChanges() throws InterruptedException {
130136
SearchIndexStatus.READY);
131137
}
132138

139+
@Tag("vector-search")
133140
@Test // GH-4706
134141
void exists() throws InterruptedException {
135142

@@ -147,6 +154,7 @@ void exists() throws InterruptedException {
147154
assertThat(indexOps.exists(indexName)).isTrue();
148155
}
149156

157+
@Tag("vector-search")
150158
@Test // GH-4706
151159
void updatesVectorIndex() throws InterruptedException {
152160

@@ -176,6 +184,7 @@ void updatesVectorIndex() throws InterruptedException {
176184
assertThatRuntimeException().isThrownBy(() -> indexOps.updateIndex(updatedIdx));
177185
}
178186

187+
@Tag("vector-search")
179188
@Test // GH-4706
180189
void createsVectorIndexWithFilters() throws InterruptedException {
181190

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,11 @@
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
/**
2425
* Extension to MongoDBAtlasLocalContainer.
2526
*
@@ -55,4 +56,12 @@ public static AtlasContainer tagged(String tag) {
5556
return new AtlasContainer(DEFAULT_IMAGE_NAME.withTag(tag));
5657
}
5758

59+
@Override
60+
protected void containerIsStarted(InspectContainerResponse containerInfo) {
61+
62+
super.containerIsStarted(containerInfo);
63+
64+
System.setProperty("docker.mongodb.atlas.host", getHost());
65+
System.setProperty("docker.mongodb.atlas.port", getMappedPort(27017).toString());
66+
}
5867
}

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

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,10 @@
2121
import org.junit.jupiter.api.extension.ExtensionContext.Namespace;
2222
import org.springframework.core.annotation.AnnotatedElementUtils;
2323
import org.springframework.data.util.Version;
24+
import org.springframework.util.NumberUtils;
25+
import org.springframework.util.StringUtils;
26+
27+
import com.mongodb.client.MongoClient;
2428

2529
/**
2630
* @author Christoph Strobl
@@ -42,8 +46,8 @@ public ConditionEvaluationResult evaluateExecutionCondition(ExtensionContext con
4246
}
4347
}
4448

45-
if(context.getTags().contains("vector-search")) {
46-
if(!atlasEnvironment(context)) {
49+
if (context.getTags().contains("vector-search")) {
50+
if (!atlasEnvironment(context)) {
4751
return ConditionEvaluationResult.disabled("Disabled for servers not supporting Vector Search.");
4852
}
4953
}
@@ -91,7 +95,18 @@ private Version serverVersion(ExtensionContext context) {
9195
}
9296

9397
private boolean atlasEnvironment(ExtensionContext context) {
94-
return context.getStore(NAMESPACE).getOrComputeIfAbsent(Version.class, (key) -> MongoTestUtils.isVectorSearchEnabled(),
95-
Boolean.class);
98+
99+
String host = System.getProperty("docker.mongodb.atlas.host");
100+
String port = System.getProperty("docker.mongodb.atlas.port");
101+
102+
return context.getStore(NAMESPACE).getOrComputeIfAbsent(Version.class, (key) -> {
103+
104+
if (StringUtils.hasText(host) && StringUtils.hasText(port)) {
105+
try (MongoClient client = MongoTestUtils.client(host, NumberUtils.parseNumber(port, Integer.class))) {
106+
return MongoTestUtils.isVectorSearchEnabled(client);
107+
}
108+
}
109+
return MongoTestUtils.isVectorSearchEnabled();
110+
}, Boolean.class);
96111
}
97112
}

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

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import java.time.Duration;
2323
import java.util.List;
24+
import java.util.concurrent.TimeUnit;
2425

2526
import org.bson.Document;
2627
import org.springframework.core.env.Environment;
@@ -30,6 +31,7 @@
3031
import org.springframework.util.ObjectUtils;
3132

3233
import com.mongodb.ConnectionString;
34+
import com.mongodb.MongoClientSettings;
3335
import com.mongodb.ReadPreference;
3436
import com.mongodb.WriteConcern;
3537
import com.mongodb.client.MongoClient;
@@ -68,6 +70,10 @@ public static MongoClient client(String host, int port) {
6870
}
6971

7072
public static MongoClient client(ConnectionString connectionString) {
73+
MongoClientSettings settings = MongoClientSettings.builder().applyConnectionString(connectionString)
74+
.applyToSocketSettings(builder -> {
75+
builder.connectTimeout(120, TimeUnit.SECONDS);
76+
}).build();
7177
return com.mongodb.client.MongoClients.create(connectionString, SpringDataMongoDB.driverInformation());
7278
}
7379

@@ -176,11 +182,10 @@ public static void dropCollectionNow(String dbName, String collectionName,
176182
* @param collectionName must not be {@literal null}.
177183
* @param client must not be {@literal null}.
178184
*/
179-
public static void dropCollectionNow(String dbName, String collectionName,
180-
com.mongodb.client.MongoClient client) {
185+
public static void dropCollectionNow(String dbName, String collectionName, com.mongodb.client.MongoClient client) {
181186

182-
com.mongodb.client.MongoDatabase database = client.getDatabase(dbName)
183-
.withWriteConcern(WriteConcern.MAJORITY).withReadPreference(ReadPreference.primary());
187+
com.mongodb.client.MongoDatabase database = client.getDatabase(dbName).withWriteConcern(WriteConcern.MAJORITY)
188+
.withReadPreference(ReadPreference.primary());
184189

185190
database.getCollection(collectionName).drop();
186191
}
@@ -205,11 +210,10 @@ public static void flushCollection(String dbName, String collectionName,
205210
.verifyComplete();
206211
}
207212

208-
public static void flushCollection(String dbName, String collectionName,
209-
com.mongodb.client.MongoClient client) {
213+
public static void flushCollection(String dbName, String collectionName, com.mongodb.client.MongoClient client) {
210214

211-
com.mongodb.client.MongoDatabase database = client.getDatabase(dbName)
212-
.withWriteConcern(WriteConcern.MAJORITY).withReadPreference(ReadPreference.primary());
215+
com.mongodb.client.MongoDatabase database = client.getDatabase(dbName).withWriteConcern(WriteConcern.MAJORITY)
216+
.withReadPreference(ReadPreference.primary());
213217

214218
database.getCollection(collectionName).deleteMany(new Document());
215219
}
@@ -267,14 +271,19 @@ public static boolean serverIsReplSet() {
267271
@SuppressWarnings("unchecked")
268272
public static boolean isVectorSearchEnabled() {
269273
try (MongoClient client = MongoTestUtils.client()) {
274+
return isVectorSearchEnabled(client);
275+
}
276+
}
270277

278+
public static boolean isVectorSearchEnabled(MongoClient client) {
279+
try {
271280
return client.getDatabase("admin").runCommand(new Document("getCmdLineOpts", "1")).get("argv", List.class)
272-
.stream().anyMatch(it -> {
273-
if(it instanceof String cfgString) {
274-
return cfgString.startsWith("searchIndexManagementHostAndPort");
275-
}
276-
return false;
277-
});
281+
.stream().anyMatch(it -> {
282+
if (it instanceof String cfgString) {
283+
return cfgString.startsWith("searchIndexManagementHostAndPort");
284+
}
285+
return false;
286+
});
278287
} catch (Exception e) {
279288
return false;
280289
}
@@ -297,10 +306,11 @@ private static void giveTheServerALittleTimeToThink() {
297306

298307
public static CollectionInfo readCollectionInfo(MongoDatabase db, String collectionName) {
299308

300-
List<Document> list = db.runCommand(new Document().append("listCollections", 1).append("filter", new Document("name", collectionName)))
309+
List<Document> list = db
310+
.runCommand(new Document().append("listCollections", 1).append("filter", new Document("name", collectionName)))
301311
.get("cursor", Document.class).get("firstBatch", List.class);
302312

303-
if(list.isEmpty()) {
313+
if (list.isEmpty()) {
304314
throw new IllegalStateException(String.format("Collection %s not found.", collectionName));
305315
}
306316
return CollectionInfo.from(list.get(0));

0 commit comments

Comments
 (0)