From c105ecc25bb9e99391cec69b46740c657d7a130e Mon Sep 17 00:00:00 2001
From: Martin7-1 <yi.zheng.se@gmail.com>
Date: Fri, 15 Nov 2024 11:52:30 +0800
Subject: [PATCH] Uses official ollama testcontainers

---
 langchain4j-ollama-spring-boot-starter/pom.xml  |  2 +-
 .../langchain4j/ollama/spring/AutoConfigIT.java | 17 +++++++++++++----
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/langchain4j-ollama-spring-boot-starter/pom.xml b/langchain4j-ollama-spring-boot-starter/pom.xml
index eecb11b5..07ec72e2 100644
--- 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/langchain4j-ollama-spring-boot-starter/src/test/java/dev/langchain4j/ollama/spring/AutoConfigIT.java b/langchain4j-ollama-spring-boot-starter/src/test/java/dev/langchain4j/ollama/spring/AutoConfigIT.java
index fced9a72..4c4c1f9f 100644
--- a/langchain4j-ollama-spring-boot-starter/src/test/java/dev/langchain4j/ollama/spring/AutoConfigIT.java
+++ b/langchain4j-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