Uses official ollama testcontainers

langchain4j · Nov 15, 2024 · c105ecc · c105ecc
1 parent c6c183b
commit c105ecc
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/langchain4j-ollama-spring-boot-starter/pom.xml b/langchain4j-ollama-spring-boot-starter/pom.xml
@@ -59,7 +59,7 @@
 
         <dependency>
             <groupId>org.testcontainers</groupId>
-            <artifactId>testcontainers</artifactId>
+            <artifactId>ollama</artifactId>
             <scope>test</scope>
         </dependency>
 

diff --git a/...-ollama-spring-boot-starter/src/test/java/dev/langchain4j/ollama/spring/AutoConfigIT.java b/...-ollama-spring-boot-starter/src/test/java/dev/langchain4j/ollama/spring/AutoConfigIT.java
@@ -9,16 +9,17 @@
 import dev.langchain4j.model.language.StreamingLanguageModel;
 import dev.langchain4j.model.ollama.*;
 import dev.langchain4j.model.output.Response;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
-import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.junit.jupiter.Container;
 import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.ollama.OllamaContainer;
+import org.testcontainers.utility.DockerImageName;
 
 import java.util.concurrent.CompletableFuture;
 
-import static java.lang.String.format;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -28,14 +29,22 @@ class AutoConfigIT {
     private static final String MODEL_NAME = "phi";
 
     @Container
-    static GenericContainer<?> ollama = new GenericContainer<>("langchain4j/ollama-" + MODEL_NAME)
+    static OllamaContainer ollama = new OllamaContainer(
+            DockerImageName.parse("alpine/ollama:latest")
+                    .asCompatibleSubstituteFor("ollama/ollama")
+    )
             .withExposedPorts(11434);
 
+    @BeforeAll
+    static void beforeAll() throws Exception {
+        ollama.execInContainer("ollama", "pull", MODEL_NAME);
+    }
+
     ApplicationContextRunner contextRunner = new ApplicationContextRunner()
             .withConfiguration(AutoConfigurations.of(AutoConfig.class));
 
     private static String baseUrl() {
-        return format("http://%s:%s", ollama.getHost(), ollama.getFirstMappedPort());
+        return ollama.getEndpoint();
     }
 
     @Test