Polish

Author: eddumelendez

core/build.gradle

@@ -123,9 +123,6 @@ dependencies {
     testImplementation 'org.assertj:assertj-core:3.26.3'
     testImplementation 'io.rest-assured:rest-assured:5.5.0'
 
-    // DockerModelRunnerContainer tests
-    testImplementation("com.openai:openai-java:1.3.0")
-
     jarFileTestCompileOnly "org.projectlombok:lombok:${lombok.version}"
     jarFileTestAnnotationProcessor "org.projectlombok:lombok:${lombok.version}"
     jarFileTestImplementation 'junit:junit:4.13.2'

core/src/main/java/org/testcontainers/containers/DockerModelRunnerContainer.java

@@ -1,5 +1,7 @@
 package org.testcontainers.containers;
 
+import org.testcontainers.containers.wait.strategy.Wait;
+
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStreamReader;
@@ -8,26 +10,29 @@
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 
-public class DockerModelRunnerContainer extends GenericContainer {
+public class DockerModelRunnerContainer extends GenericContainer<DockerModelRunnerContainer> {
 
     public static final String MODEL_RUNNER_ENDPOINT = "model-runner.docker.internal";
+
     private SocatContainer socat;
+
     private String model;
 
     @Override
     public void start() {
-        socat = new SocatContainer()
-            .withTarget(80, MODEL_RUNNER_ENDPOINT, 80);
-        socat.start();
+        this.socat =
+            new SocatContainer()
+                .withTarget(80, MODEL_RUNNER_ENDPOINT)
+                .waitingFor(Wait.forHttp("/").forResponsePredicate(res -> res.contains("The service is running")));
+        this.socat.start();
         pullModel();
     }
 
     private void pullModel() {
-        logger().info("Pulling model: {}. Please be patient, no progress bar yet!", model);
+        logger().info("Pulling model: {}. Please be patient, no progress bar yet!", this.model);
         try {
-            // Construct JSON payload
-            String json = String.format("{\"from\":\"%s\"}", model);
-            String endpoint = "http://" + socat.getHost() + ":" + socat.getMappedPort(80) + "/models/create";
+            String json = String.format("{\"from\":\"%s\"}", this.model);
+            String endpoint = "http://" + this.socat.getHost() + ":" + this.socat.getMappedPort(80) + "/models/create";
 
             URL url = new URL(endpoint);
             HttpURLConnection connection = (HttpURLConnection) url.openConnection();
@@ -40,7 +45,11 @@ private void pullModel() {
                 os.write(input, 0, input.length);
             }
 
-            try (BufferedReader br = new BufferedReader(new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8))) {
+            try (
+                BufferedReader br = new BufferedReader(
+                    new InputStreamReader(connection.getInputStream(), StandardCharsets.UTF_8)
+                )
+            ) {
                 StringBuilder response = new StringBuilder();
                 String responseLine;
                 while ((responseLine = br.readLine()) != null) {
@@ -56,16 +65,19 @@ private void pullModel() {
 
     @Override
     public void stop() {
-        socat.stop();
+        this.socat.stop();
+    }
+
+    public String getBaseEndpoint() {
+        return "http://" + this.socat.getHost() + ":" + this.socat.getMappedPort(80);
     }
 
     public String getOpenAIEndpoint() {
-        return "http://" + socat.getHost() + ":" + socat.getMappedPort(80) + "/engines";
+        return getBaseEndpoint() + "/engines";
     }
 
     public DockerModelRunnerContainer withModel(String modelName) {
         this.model = modelName;
         return this;
     }
-
 }

core/src/test/java/org/testcontainers/containers/DockerModelRunnerContainerTest.java

@@ -1,35 +1,25 @@
 package org.testcontainers.containers;
 
-import com.openai.client.OpenAIClient;
-import com.openai.client.okhttp.OpenAIOkHttpClient;
-import com.openai.models.chat.completions.ChatCompletion;
-import com.openai.models.chat.completions.ChatCompletionCreateParams;
+import io.restassured.RestAssured;
+import io.restassured.response.Response;
 import org.junit.Test;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 public class DockerModelRunnerContainerTest {
 
     @Test
     public void pullsModelAndExposesInference() {
         String modelName = "ai/smollm2:360M-Q4_K_M";
 
-        try (DockerModelRunnerContainer dmr = new DockerModelRunnerContainer()
-            .withModel(modelName);) {
+        try (DockerModelRunnerContainer dmr = new DockerModelRunnerContainer().withModel(modelName)) {
             dmr.start();
 
-            OpenAIClient client = OpenAIOkHttpClient.builder()
-                .baseUrl(dmr.getOpenAIEndpoint())
-                .build();
-
-            ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
-                .addUserMessage("Say this is a test")
-                .model(modelName)
-                .build();
-            ChatCompletion chatCompletion = client.chat().completions().create(params);
-
-            String answer = chatCompletion.toString();
-            System.out.println(answer);
+            Response response = RestAssured.get(dmr.getBaseEndpoint() + "/models").thenReturn();
+            assertThat(response.body().jsonPath().getList("tags.flatten()")).contains(modelName);
 
+            Response openAiResponse = RestAssured.get(dmr.getOpenAIEndpoint() + "/v1/models").prettyPeek().thenReturn();
+            assertThat(openAiResponse.body().jsonPath().getList("data.id")).contains(modelName);
         }
     }
-
 }