Skip to content

Commit ad07a81

Browse files
ilayaperumalgilayaperumalg
committed
Fix MongoDB Atlas vector store ITs
 - Use MongoConnectionDetails in the auto-configuration test configuration when using testcontainers - Use MongoClient to create MongoTemplate for the vector store ITs test configuration Signed-off-by: Ilayaperumal Gopinathan <[email protected]>
1 parent a7b837a commit ad07a81

File tree

3 files changed

+73
-63
lines changed

3 files changed

+73
-63
lines changed

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-mongodb-atlas/src/test/java/org/springframework/ai/vectorstore/mongodb/autoconfigure/MongoDBAtlasVectorStoreAutoConfigurationIT.java‎

Lines changed: 33 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.List;
2121
import java.util.stream.Collectors;
2222

23+
import com.mongodb.ConnectionString;
2324
import io.micrometer.observation.tck.TestObservationRegistry;
2425
import org.junit.jupiter.api.Test;
2526
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -40,6 +41,7 @@
4041
import org.springframework.boot.autoconfigure.AutoConfigurations;
4142
import org.springframework.boot.data.mongodb.autoconfigure.DataMongoAutoConfiguration;
4243
import org.springframework.boot.mongodb.autoconfigure.MongoAutoConfiguration;
44+
import org.springframework.boot.mongodb.autoconfigure.MongoConnectionDetails;
4345
import org.springframework.boot.restclient.autoconfigure.RestClientAutoConfiguration;
4446
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4547
import org.springframework.context.annotation.Bean;
@@ -53,26 +55,25 @@
5355
* @author Christian Tzolov
5456
* @author Thomas Vitale
5557
* @author Ignacio López
58+
* @author Ilayaperumal Gopinathan
5659
*/
5760
@Testcontainers
5861
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
5962
class MongoDBAtlasVectorStoreAutoConfigurationIT {
6063

6164
@Container
62-
static MongoDBAtlasLocalContainer mongo = new MongoDBAtlasLocalContainer("mongodb/mongodb-atlas-local:7.0.9");
63-
64-
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
65-
.withUserConfiguration(Config.class)
66-
.withConfiguration(AutoConfigurations.of(MongoAutoConfiguration.class, DataMongoAutoConfiguration.class,
67-
MongoDBAtlasVectorStoreAutoConfiguration.class, RestClientAutoConfiguration.class,
68-
SpringAiRetryAutoConfiguration.class, OpenAiEmbeddingAutoConfiguration.class))
69-
.withPropertyValues("spring.data.mongodb.database=springaisample",
70-
"spring.ai.vectorstore.mongodb.initialize-schema=true",
71-
"spring.ai.vectorstore.mongodb.collection-name=test_collection",
72-
// "spring.ai.vectorstore.mongodb.path-name=testembedding",
73-
"spring.ai.vectorstore.mongodb.index-name=text_index",
74-
"spring.ai.openai.api-key=" + System.getenv("OPENAI_API_KEY"),
75-
String.format("spring.data.mongodb.uri=" + mongo.getConnectionString()));
65+
static MongoDBAtlasLocalContainer mongo = new MongoDBAtlasLocalContainer("mongodb/mongodb-atlas-local:8.0.0");
66+
67+
private ApplicationContextRunner getContextRunner() {
68+
return new ApplicationContextRunner().withUserConfiguration(Config.class)
69+
.withConfiguration(AutoConfigurations.of(MongoAutoConfiguration.class, DataMongoAutoConfiguration.class,
70+
MongoDBAtlasVectorStoreAutoConfiguration.class, RestClientAutoConfiguration.class,
71+
SpringAiRetryAutoConfiguration.class, OpenAiEmbeddingAutoConfiguration.class))
72+
.withPropertyValues("spring.ai.vectorstore.mongodb.initialize-schema=true",
73+
"spring.ai.vectorstore.mongodb.collection-name=test_collection",
74+
"spring.ai.vectorstore.mongodb.index-name=text_index",
75+
"spring.ai.openai.api-key=" + System.getenv("OPENAI_API_KEY"));
76+
}
7677

7778
List<Document> documents = List.of(
7879
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
@@ -90,7 +91,7 @@ class MongoDBAtlasVectorStoreAutoConfigurationIT {
9091

9192
@Test
9293
public void addAndSearch() {
93-
this.contextRunner.run(context -> {
94+
getContextRunner().run(context -> {
9495

9596
VectorStore vectorStore = context.getBean(VectorStore.class);
9697
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
@@ -133,7 +134,7 @@ public void addAndSearch() {
133134

134135
@Test
135136
public void addAndSearchWithFilters() {
136-
this.contextRunner.withPropertyValues("spring.ai.vectorstore.mongodb.metadata-fields-to-filter=foo")
137+
getContextRunner().withPropertyValues("spring.ai.vectorstore.mongodb.metadata-fields-to-filter=foo")
137138
.run(context -> {
138139

139140
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -165,7 +166,7 @@ public void addAndSearchWithFilters() {
165166

166167
@Test
167168
public void autoConfigurationDisabledWhenTypeIsNone() {
168-
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
169+
getContextRunner().withPropertyValues("spring.ai.vectorstore.type=none").run(context -> {
169170
assertThat(context.getBeansOfType(MongoDBAtlasVectorStoreProperties.class)).isEmpty();
170171
assertThat(context.getBeansOfType(MongoDBAtlasVectorStore.class)).isEmpty();
171172
assertThat(context.getBeansOfType(VectorStore.class)).isEmpty();
@@ -174,7 +175,7 @@ public void autoConfigurationDisabledWhenTypeIsNone() {
174175

175176
@Test
176177
public void autoConfigurationEnabledByDefault() {
177-
this.contextRunner.run(context -> {
178+
getContextRunner().run(context -> {
178179
assertThat(context.getBeansOfType(MongoDBAtlasVectorStoreProperties.class)).isNotEmpty();
179180
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
180181
assertThat(context.getBean(VectorStore.class)).isInstanceOf(MongoDBAtlasVectorStore.class);
@@ -183,7 +184,7 @@ public void autoConfigurationEnabledByDefault() {
183184

184185
@Test
185186
public void autoConfigurationEnabledWhenTypeIsMongodbAtlas() {
186-
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=mongodb-atlas").run(context -> {
187+
getContextRunner().withPropertyValues("spring.ai.vectorstore.type=mongodb-atlas").run(context -> {
187188
assertThat(context.getBeansOfType(MongoDBAtlasVectorStoreProperties.class)).isNotEmpty();
188189
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
189190
assertThat(context.getBean(VectorStore.class)).isInstanceOf(MongoDBAtlasVectorStore.class);
@@ -198,6 +199,19 @@ public TestObservationRegistry observationRegistry() {
198199
return TestObservationRegistry.create();
199200
}
200201

202+
@Bean
203+
public MongoConnectionDetails mongoConnectionDetails() {
204+
return new MongoConnectionDetails() {
205+
@Override
206+
public ConnectionString getConnectionString() {
207+
// Add database name to the connection string
208+
String baseUri = mongo.getConnectionString();
209+
String uriWithDb = baseUri.replace("/?", "/springaisample?");
210+
return new ConnectionString(uriWithDb);
211+
}
212+
};
213+
}
214+
201215
}
202216

203217
}

‎vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java‎

Lines changed: 24 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import java.util.stream.Collectors;
2929

3030
import com.mongodb.client.MongoClient;
31+
import com.mongodb.client.MongoClients;
3132
import org.junit.jupiter.api.BeforeEach;
3233
import org.junit.jupiter.api.Test;
3334
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -45,15 +46,12 @@
4546
import org.springframework.ai.vectorstore.VectorStore;
4647
import org.springframework.ai.vectorstore.filter.Filter;
4748
import org.springframework.boot.SpringBootConfiguration;
48-
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
4949
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
5050
import org.springframework.context.annotation.Bean;
5151
import org.springframework.core.convert.converter.Converter;
5252
import org.springframework.core.io.DefaultResourceLoader;
5353
import org.springframework.data.mongodb.core.MongoTemplate;
54-
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
5554
import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
56-
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
5755
import org.springframework.util.MimeType;
5856

5957
import static org.assertj.core.api.Assertions.assertThat;
@@ -63,6 +61,7 @@
6361
* @author Soby Chacko
6462
* @author Eddú Meléndez
6563
* @author Thomas Vitale
64+
* @author Ilayaperumal Gopinathan
6665
*/
6766
@Testcontainers
6867
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -71,30 +70,29 @@ class MongoDBAtlasVectorStoreIT extends BaseVectorStoreTests {
7170
@Container
7271
private static MongoDBAtlasLocalContainer container = new MongoDBAtlasLocalContainer(MongoDbImage.DEFAULT_IMAGE);
7372

74-
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
75-
.withUserConfiguration(TestApplication.class)
76-
.withPropertyValues("spring.data.mongodb.database=springaisample",
77-
String.format("spring.data.mongodb.uri=" + container.getConnectionString()));
73+
private ApplicationContextRunner getContextRunner() {
74+
return new ApplicationContextRunner().withUserConfiguration(TestApplication.class);
75+
}
7876

7977
@BeforeEach
8078
public void beforeEach() {
81-
this.contextRunner.run(context -> {
79+
getContextRunner().run(context -> {
8280
MongoTemplate mongoTemplate = context.getBean(MongoTemplate.class);
8381
mongoTemplate.getCollection("vector_store").deleteMany(new org.bson.Document());
8482
});
8583
}
8684

8785
@Override
8886
protected void executeTest(Consumer<VectorStore> testFunction) {
89-
this.contextRunner.run(context -> {
87+
getContextRunner().run(context -> {
9088
VectorStore vectorStore = context.getBean(VectorStore.class);
9189
testFunction.accept(vectorStore);
9290
});
9391
}
9492

9593
@Test
9694
void vectorStoreTest() {
97-
this.contextRunner.run(context -> {
95+
getContextRunner().run(context -> {
9896
VectorStore vectorStore = context.getBean(VectorStore.class);
9997

10098
List<Document> documents = List.of(
@@ -131,7 +129,7 @@ void vectorStoreTest() {
131129

132130
@Test
133131
void documentUpdateTest() {
134-
this.contextRunner.run(context -> {
132+
getContextRunner().run(context -> {
135133
VectorStore vectorStore = context.getBean(VectorStore.class);
136134

137135
Document document = new Document(UUID.randomUUID().toString(), "Spring AI rocks!!",
@@ -167,7 +165,7 @@ void documentUpdateTest() {
167165

168166
@Test
169167
void searchWithFilters() {
170-
this.contextRunner.run(context -> {
168+
getContextRunner().run(context -> {
171169
VectorStore vectorStore = context.getBean(VectorStore.class);
172170

173171
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
@@ -230,7 +228,7 @@ void searchWithFilters() {
230228

231229
@Test
232230
public void searchWithThreshold() {
233-
this.contextRunner.run(context -> {
231+
getContextRunner().run(context -> {
234232
VectorStore vectorStore = context.getBean(VectorStore.class);
235233

236234
var documents = List.of(
@@ -269,7 +267,7 @@ public void searchWithThreshold() {
269267

270268
@Test
271269
void deleteWithComplexFilterExpression() {
272-
this.contextRunner.run(context -> {
270+
getContextRunner().run(context -> {
273271
VectorStore vectorStore = context.getBean(VectorStore.class);
274272

275273
var doc1 = new Document("Content 1", Map.of("type", "A", "priority", 1));
@@ -303,7 +301,7 @@ void deleteWithComplexFilterExpression() {
303301

304302
@Test
305303
void getNativeClientTest() {
306-
this.contextRunner.run(context -> {
304+
getContextRunner().run(context -> {
307305
MongoDBAtlasVectorStore vectorStore = context.getBean(MongoDBAtlasVectorStore.class);
308306
Optional<MongoTemplate> nativeClient = vectorStore.getNativeClient();
309307
assertThat(nativeClient).isPresent();
@@ -321,7 +319,6 @@ public static String getText(String uri) {
321319
}
322320

323321
@SpringBootConfiguration
324-
@EnableAutoConfiguration
325322
public static class TestApplication {
326323

327324
@Bean
@@ -333,19 +330,20 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed
333330
}
334331

335332
@Bean
336-
public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) {
337-
MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample");
338-
MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter();
339-
converter.setCustomConversions(mongoCustomConversions);
340-
((MongoMappingContext) converter.getMappingContext())
341-
.setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder());
342-
converter.afterPropertiesSet();
343-
return mongoTemplate;
333+
public EmbeddingModel embeddingModel() {
334+
return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build());
344335
}
345336

346337
@Bean
347-
public EmbeddingModel embeddingModel() {
348-
return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build());
338+
public MongoClient mongoClient() {
339+
String baseUri = container.getConnectionString();
340+
String uriWithDb = baseUri.replace("/?", "/springaisample?");
341+
return MongoClients.create(uriWithDb);
342+
}
343+
344+
@Bean
345+
public MongoTemplate mongoTemplate(MongoClient mongoClient) {
346+
return new MongoTemplate(mongoClient, "springaisample");
349347
}
350348

351349
@Bean

‎vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDbVectorStoreObservationIT.java‎

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.Map;
2424

2525
import com.mongodb.client.MongoClient;
26+
import com.mongodb.client.MongoClients;
2627
import io.micrometer.observation.ObservationRegistry;
2728
import io.micrometer.observation.tck.TestObservationRegistry;
2829
import io.micrometer.observation.tck.TestObservationRegistryAssert;
@@ -46,15 +47,12 @@
4647
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames;
4748
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames;
4849
import org.springframework.boot.SpringBootConfiguration;
49-
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
5050
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
5151
import org.springframework.context.annotation.Bean;
5252
import org.springframework.core.convert.converter.Converter;
5353
import org.springframework.core.io.DefaultResourceLoader;
5454
import org.springframework.data.mongodb.core.MongoTemplate;
55-
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
5655
import org.springframework.data.mongodb.core.convert.MongoCustomConversions;
57-
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
5856
import org.springframework.util.MimeType;
5957

6058
import static org.assertj.core.api.Assertions.assertThat;
@@ -64,6 +62,7 @@
6462
* @author Soby Chacko
6563
* @author Thomas Vitale
6664
* @author Eddú Meléndez
65+
* @author Ilayaperumal Gopinathan
6766
*/
6867
@Testcontainers
6968
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -72,10 +71,9 @@ public class MongoDbVectorStoreObservationIT {
7271
@Container
7372
private static MongoDBAtlasLocalContainer container = new MongoDBAtlasLocalContainer(MongoDbImage.DEFAULT_IMAGE);
7473

75-
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
76-
.withUserConfiguration(Config.class)
77-
.withPropertyValues("spring.data.mongodb.database=springaisample",
78-
String.format("spring.data.mongodb.uri=" + container.getConnectionString()));
74+
private ApplicationContextRunner getContextRunner() {
75+
return new ApplicationContextRunner().withUserConfiguration(Config.class);
76+
}
7977

8078
List<Document> documents = List.of(
8179
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
@@ -94,7 +92,7 @@ public static String getText(String uri) {
9492

9593
@BeforeEach
9694
public void beforeEach() {
97-
this.contextRunner.run(context -> {
95+
getContextRunner().run(context -> {
9896
MongoTemplate mongoTemplate = context.getBean(MongoTemplate.class);
9997
mongoTemplate.getCollection("vector_store").deleteMany(new org.bson.Document());
10098
});
@@ -103,7 +101,7 @@ public void beforeEach() {
103101
@Test
104102
void observationVectorStoreAddAndQueryOperations() {
105103

106-
this.contextRunner.run(context -> {
104+
getContextRunner().run(context -> {
107105

108106
VectorStore vectorStore = context.getBean(VectorStore.class);
109107

@@ -176,7 +174,6 @@ void observationVectorStoreAddAndQueryOperations() {
176174
}
177175

178176
@SpringBootConfiguration
179-
@EnableAutoConfiguration
180177
public static class Config {
181178

182179
@Bean
@@ -197,14 +194,15 @@ public VectorStore vectorStore(MongoTemplate mongoTemplate, EmbeddingModel embed
197194
}
198195

199196
@Bean
200-
public MongoTemplate mongoTemplate(MongoClient mongoClient, MongoCustomConversions mongoCustomConversions) {
201-
MongoTemplate mongoTemplate = new MongoTemplate(mongoClient, "springaisample");
202-
MappingMongoConverter converter = (MappingMongoConverter) mongoTemplate.getConverter();
203-
converter.setCustomConversions(mongoCustomConversions);
204-
((MongoMappingContext) converter.getMappingContext())
205-
.setSimpleTypeHolder(mongoCustomConversions.getSimpleTypeHolder());
206-
converter.afterPropertiesSet();
207-
return mongoTemplate;
197+
public MongoClient mongoClient() {
198+
String baseUri = container.getConnectionString();
199+
String uriWithDb = baseUri.replace("/?", "/springaisample?");
200+
return MongoClients.create(uriWithDb);
201+
}
202+
203+
@Bean
204+
public MongoTemplate mongoTemplate(MongoClient mongoClient) {
205+
return new MongoTemplate(mongoClient, "springaisample");
208206
}
209207

210208
@Bean

0 commit comments

Comments
 (0)