From dcdc081a6a5fbaea867fdf96871243dace56a461 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 +-
 .../dev/langchain4j/ollama/spring/AutoConfigIT.java   | 11 +++++++----
 2 files changed, 8 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..8f0ca1cf 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
@@ -12,13 +12,13 @@
 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 +28,17 @@ 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("langchain4j/ollama-" + MODEL_NAME)
+                    .asCompatibleSubstituteFor("ollama/ollama")
+    )
             .withExposedPorts(11434);
 
     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