Skip to content

Commit 906b6d2

Browse files
eddumelendezeddumelendez
committed
Add Service Connection support for Testcontainers MongoDB Atlas
Testcontainers 1.20.2 provides `MongoDBAtlasLocalContainer`.
1 parent 110a520 commit 906b6d2

File tree

5 files changed

+175
-0
lines changed

5 files changed

+175
-0
lines changed

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/testcontainers.adoc‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ The following service connection factories are provided in the `spring-ai-spring
3737
| `MilvusServiceClientConnectionDetails`
3838
| Containers of type `MilvusContainer`
3939

40+
| `MongoConnectionDetails`
41+
| Containers of type `MongoDBAtlasLocalContainer`
42+
4043
| `OllamaConnectionDetails`
4144
| Containers of type `OllamaContainer`
4245

‎spring-ai-spring-boot-testcontainers/pom.xml‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,13 @@
131131
<optional>true</optional>
132132
</dependency>
133133

134+
<dependency>
135+
<groupId>org.springframework.ai</groupId>
136+
<artifactId>spring-ai-mongodb-atlas-store</artifactId>
137+
<version>${project.parent.version}</version>
138+
<optional>true</optional>
139+
</dependency>
140+
134141
<!-- test dependencies -->
135142

136143
<dependency>
@@ -218,6 +225,13 @@
218225
<optional>true</optional>
219226
</dependency>
220227

228+
<dependency>
229+
<groupId>org.testcontainers</groupId>
230+
<artifactId>mongodb</artifactId>
231+
<version>1.20.2</version>
232+
<optional>true</optional>
233+
</dependency>
234+
221235
<dependency>
222236
<groupId>org.testcontainers</groupId>
223237
<artifactId>ollama</artifactId>

‎spring-ai-spring-boot-testcontainers/src/main/java/org/springframework/ai/testcontainers/service/connection/mongo/MongoDbAtlasLocalContainerConnectionDetailsFactory.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright 2023 - 2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.ai.testcontainers.service.connection.mongo;
17+
18+
import com.mongodb.ConnectionString;
19+
import org.springframework.boot.autoconfigure.mongo.MongoConnectionDetails;
20+
import org.springframework.boot.testcontainers.service.connection.ContainerConnectionDetailsFactory;
21+
import org.springframework.boot.testcontainers.service.connection.ContainerConnectionSource;
22+
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
23+
24+
/**
25+
* @author Eddú Meléndez
26+
*/
27+
class MongoDbAtlasLocalContainerConnectionDetailsFactory
28+
extends ContainerConnectionDetailsFactory<MongoDBAtlasLocalContainer, MongoConnectionDetails> {
29+
30+
@Override
31+
protected MongoConnectionDetails getContainerConnectionDetails(
32+
ContainerConnectionSource<MongoDBAtlasLocalContainer> source) {
33+
return new MongoDbAtlasLocalContainerConnectionDetails(source);
34+
}
35+
36+
/**
37+
* {@link MongoConnectionDetails} backed by a {@link ContainerConnectionSource}.
38+
*/
39+
private static final class MongoDbAtlasLocalContainerConnectionDetails
40+
extends ContainerConnectionDetails<MongoDBAtlasLocalContainer> implements MongoConnectionDetails {
41+
42+
private MongoDbAtlasLocalContainerConnectionDetails(
43+
ContainerConnectionSource<MongoDBAtlasLocalContainer> source) {
44+
super(source);
45+
}
46+
47+
@Override
48+
public ConnectionString getConnectionString() {
49+
return new ConnectionString(getContainer().getConnectionString());
50+
}
51+
52+
}
53+
54+
}

‎spring-ai-spring-boot-testcontainers/src/main/resources/META-INF/spring.factories‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
org.springframework.boot.autoconfigure.service.connection.ConnectionDetailsFactory=\
22
org.springframework.ai.testcontainers.service.connection.chroma.ChromaContainerConnectionDetailsFactory,\
33
org.springframework.ai.testcontainers.service.connection.milvus.MilvusContainerConnectionDetailsFactory,\
4+
org.springframework.ai.testcontainers.service.connection.mongo.MongoDbAtlasLocalContainerConnectionDetailsFactory,\
45
org.springframework.ai.testcontainers.service.connection.ollama.OllamaContainerConnectionDetailsFactory,\
56
org.springframework.ai.testcontainers.service.connection.opensearch.OpenSearchContainerConnectionDetailsFactory,\
67
org.springframework.ai.testcontainers.service.connection.qdrant.QdrantContainerConnectionDetailsFactory,\

‎spring-ai-spring-boot-testcontainers/src/test/java/org/springframework/ai/testcontainers/service/connection/mongo/MongoDbAtlasLocalContainerConnectionDetailsFactoryTest.java‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/*
2+
* Copyright 2023 - 2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.ai.testcontainers.service.connection.mongo;
17+
18+
import org.junit.jupiter.api.Test;
19+
import org.springframework.ai.autoconfigure.vectorstore.mongo.MongoDBAtlasVectorStoreAutoConfiguration;
20+
import org.springframework.ai.document.Document;
21+
import org.springframework.ai.embedding.EmbeddingModel;
22+
import org.springframework.ai.openai.OpenAiEmbeddingModel;
23+
import org.springframework.ai.openai.api.OpenAiApi;
24+
import org.springframework.ai.vectorstore.SearchRequest;
25+
import org.springframework.ai.vectorstore.VectorStore;
26+
import org.springframework.beans.factory.annotation.Autowired;
27+
import org.springframework.boot.autoconfigure.ImportAutoConfiguration;
28+
import org.springframework.boot.autoconfigure.data.mongo.MongoDataAutoConfiguration;
29+
import org.springframework.boot.autoconfigure.mongo.MongoAutoConfiguration;
30+
import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
31+
import org.springframework.context.annotation.Bean;
32+
import org.springframework.context.annotation.Configuration;
33+
import org.springframework.test.context.TestPropertySource;
34+
import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
35+
import org.testcontainers.junit.jupiter.Container;
36+
import org.testcontainers.junit.jupiter.Testcontainers;
37+
import org.testcontainers.mongodb.MongoDBAtlasLocalContainer;
38+
39+
import java.util.Collections;
40+
import java.util.List;
41+
import java.util.stream.Collectors;
42+
43+
import static org.assertj.core.api.Assertions.assertThat;
44+
45+
@SpringJUnitConfig
46+
@Testcontainers
47+
@TestPropertySource(properties = { "spring.data.mongodb.database=simpleaidb",
48+
"spring.ai.vectorstore.mongodb.initialize-schema=true",
49+
"spring.ai.vectorstore.mongodb.collection-name=test_collection",
50+
"spring.ai.vectorstore.mongodb.index-name=text_index" })
51+
class MongoDbAtlasLocalContainerConnectionDetailsFactoryTest {
52+
53+
@Container
54+
@ServiceConnection
55+
private static MongoDBAtlasLocalContainer container = new MongoDBAtlasLocalContainer(
56+
"mongodb/mongodb-atlas-local:7.0.9");
57+
58+
@Autowired
59+
private VectorStore vectorStore;
60+
61+
@Test
62+
public void addAndSearch() throws InterruptedException {
63+
List<Document> documents = List.of(
64+
new Document(
65+
"Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
66+
Collections.singletonMap("meta1", "meta1")),
67+
new Document("Hello World Hello World Hello World Hello World Hello World Hello World Hello World"),
68+
new Document(
69+
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
70+
Collections.singletonMap("meta2", "meta2")));
71+
72+
vectorStore.add(documents);
73+
Thread.sleep(5000); // Await a second for the document to be indexed
74+
75+
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
76+
77+
assertThat(results).hasSize(1);
78+
Document resultDoc = results.get(0);
79+
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
80+
assertThat(resultDoc.getContent()).isEqualTo(
81+
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression");
82+
assertThat(resultDoc.getMetadata()).containsEntry("meta2", "meta2");
83+
84+
// Remove all documents from the store
85+
vectorStore.delete(documents.stream().map(Document::getId).collect(Collectors.toList()));
86+
87+
List<Document> results2 = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1));
88+
assertThat(results2).isEmpty();
89+
}
90+
91+
@Configuration(proxyBeanMethods = false)
92+
@ImportAutoConfiguration({ MongoAutoConfiguration.class, MongoDataAutoConfiguration.class,
93+
MongoDBAtlasVectorStoreAutoConfiguration.class })
94+
static class Config {
95+
96+
@Bean
97+
public EmbeddingModel embeddingModel() {
98+
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
99+
}
100+
101+
}
102+
103+
}

0 commit comments

Comments
 (0)