Polish

Author: eddumelendez

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

@@ -1,6 +1,8 @@
 package org.testcontainers.containers;
 
+import com.github.dockerjava.api.command.InspectContainerResponse;
 import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.utility.DockerImageName;
 
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -10,29 +12,40 @@
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
 
-public class DockerModelRunnerContainer extends GenericContainer<DockerModelRunnerContainer> {
+/**
+ * Testcontainers proxy container for the Docker Model Runner service
+ * provided by Docker Desktop.
+ * <p>
+ * Supported images: {@code alpine/socat}
+ * <p>
+ * Exposed ports: 80
+ */
+public class DockerModelRunnerContainer extends SocatContainer {
 
     public static final String MODEL_RUNNER_ENDPOINT = "model-runner.docker.internal";
 
-    private SocatContainer socat;
-
     private String model;
 
+    public DockerModelRunnerContainer(String image) {
+        this(DockerImageName.parse(image));
+    }
+
+    public DockerModelRunnerContainer(DockerImageName image) {
+        super(image);
+        withTarget(80, MODEL_RUNNER_ENDPOINT);
+        waitingFor(Wait.forHttp("/").forResponsePredicate(res -> res.contains("The service is running")));
+    }
+
     @Override
-    public void start() {
-        this.socat =
-            new SocatContainer()
-                .withTarget(80, MODEL_RUNNER_ENDPOINT)
-                .waitingFor(Wait.forHttp("/").forResponsePredicate(res -> res.contains("The service is running")));
-        this.socat.start();
+    protected void containerIsStarted(InspectContainerResponse containerInfo) {
         pullModel();
     }
 
     private void pullModel() {
         logger().info("Pulling model: {}. Please be patient, no progress bar yet!", this.model);
         try {
             String json = String.format("{\"from\":\"%s\"}", this.model);
-            String endpoint = "http://" + this.socat.getHost() + ":" + this.socat.getMappedPort(80) + "/models/create";
+            String endpoint = "http://" + getHost() + ":" + getMappedPort(80) + "/models/create";
 
             URL url = new URL(endpoint);
             HttpURLConnection connection = (HttpURLConnection) url.openConnection();
@@ -63,13 +76,8 @@ private void pullModel() {
         logger().info("Finished pulling model: {}", model);
     }
 
-    @Override
-    public void stop() {
-        this.socat.stop();
-    }
-
     public String getBaseEndpoint() {
-        return "http://" + this.socat.getHost() + ":" + this.socat.getMappedPort(80);
+        return "http://" + getHost() + ":" + getMappedPort(80);
     }
 
     public String getOpenAIEndpoint() {