Support build args in FROM statement (#6119)

Fixes #3238

---------

Co-authored-by: Eddú Meléndez Gonzales <eddu.melendez@gmail.com>

core/src/main/java/org/testcontainers/images/ParsedDockerfile.java

@@ -29,17 +29,17 @@ public class ParsedDockerfile {
     private final Path dockerFilePath;
 
     @Getter
-    private Set<String> dependencyImageNames = Collections.emptySet();
+    private final Set<String> dependencyImageNames;
 
     public ParsedDockerfile(Path dockerFilePath) {
         this.dockerFilePath = dockerFilePath;
-        parse(read());
+        this.dependencyImageNames = parse(read());
     }
 
     @VisibleForTesting
     ParsedDockerfile(List<String> lines) {
         this.dockerFilePath = Paths.get("dummy.Dockerfile");
-        parse(lines);
+        this.dependencyImageNames = parse(lines);
     }
 
     private List<String> read() {
@@ -56,17 +56,17 @@ private List<String> read() {
         }
     }
 
-    private void parse(List<String> lines) {
-        dependencyImageNames =
-            lines
-                .stream()
-                .map(FROM_LINE_PATTERN::matcher)
-                .filter(Matcher::matches)
-                .map(matcher -> matcher.group("image"))
-                .collect(Collectors.toSet());
+    private Set<String> parse(List<String> lines) {
+        Set<String> imageNames = lines
+            .stream()
+            .map(FROM_LINE_PATTERN::matcher)
+            .filter(Matcher::matches)
+            .map(matcher -> matcher.group("image"))
+            .collect(Collectors.toSet());
 
-        if (!dependencyImageNames.isEmpty()) {
-            log.debug("Found dependency images in Dockerfile {}: {}", dockerFilePath, dependencyImageNames);
+        if (!imageNames.isEmpty()) {
+            log.debug("Found dependency images in Dockerfile {}: {}", dockerFilePath, imageNames);
         }
+        return imageNames;
     }
 }

core/src/main/java/org/testcontainers/images/builder/ImageFromDockerfile.java

@@ -37,6 +37,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Consumer;
+import java.util.regex.Matcher;
 import java.util.zip.GZIPOutputStream;
 
 @Slf4j
@@ -96,6 +97,7 @@ public ImageFromDockerfile withFileFromTransferable(String path, Transferable tr
     protected final String resolve() {
         Logger logger = DockerLoggerFactory.getLogger(dockerImageName);
 
+        //noinspection resource
         DockerClient dockerClient = DockerClientFactory.instance().client();
 
         try {
@@ -107,12 +109,12 @@ public void onNext(BuildResponseItem item) {
                     if (item.isErrorIndicated()) {
                         logger.error(item.getErrorDetail().getMessage());
                     } else {
-                        logger.debug(StringUtils.chomp(item.getStream(), "\n"));
+                        logger.debug(StringUtils.removeEnd(item.getStream(), "\n"));
                     }
                 }
             };
 
-            // We have to use pipes to avoid high memory consumption since users might want to build really big images
+            // We have to use pipes to avoid high memory consumption since users might want to build huge images
             @Cleanup
             PipedInputStream in = new PipedInputStream();
             @Cleanup
@@ -169,7 +171,7 @@ public void onNext(BuildResponseItem item) {
     }
 
     protected void configure(BuildImageCmd buildImageCmd) {
-        buildImageCmd.withTag(this.getDockerImageName());
+        buildImageCmd.withTags(Collections.singleton(getDockerImageName()));
         this.dockerFilePath.ifPresent(buildImageCmd::withDockerfilePath);
         this.dockerfile.ifPresent(p -> {
                 buildImageCmd.withDockerfile(p.toFile());
@@ -188,27 +190,40 @@ protected void configure(BuildImageCmd buildImageCmd) {
     }
 
     private void prePullDependencyImages(Set<String> imagesToPull) {
-        final DockerClient dockerClient = DockerClientFactory.instance().client();
-
         imagesToPull.forEach(imageName -> {
+            String resolvedImageName = applyBuildArgsToImageName(imageName);
             try {
                 log.info(
                     "Pre-emptively checking local images for '{}', referenced via a Dockerfile. If not available, it will be pulled.",
-                    imageName
+                    resolvedImageName
                 );
-                new RemoteDockerImage(DockerImageName.parse(imageName))
+                new RemoteDockerImage(DockerImageName.parse(resolvedImageName))
                     .withImageNameSubstitutor(ImageNameSubstitutor.noop())
                     .get();
             } catch (Exception e) {
                 log.warn(
                     "Unable to pre-fetch an image ({}) depended upon by Dockerfile - image build will continue but may fail. Exception message was: {}",
-                    imageName,
+                    resolvedImageName,
                     e.getMessage()
                 );
             }
         });
     }
 
+    /**
+     * See {@code filterForEnvironmentVars()} in {@link com.github.dockerjava.core.dockerfile.DockerfileStatement}.
+     */
+    private String applyBuildArgsToImageName(String imageName) {
+        for (Map.Entry<String, String> entry : buildArgs.entrySet()) {
+            String value = Matcher.quoteReplacement(entry.getValue());
+            // handle: $VARIABLE case
+            imageName = imageName.replace("$" + entry.getKey(), value);
+            // handle ${VARIABLE} case
+            imageName = imageName.replace("${" + entry.getKey() + "}", value);
+        }
+        return imageName;
+    }
+
     public ImageFromDockerfile withBuildArg(final String key, final String value) {
         this.buildArgs.put(key, value);
         return this;

core/src/test/java/org/testcontainers/images/builder/DockerfileBuildTest.java

@@ -8,6 +8,8 @@
 
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.HashMap;
+import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -22,6 +24,13 @@ public class DockerfileBuildTest {
 
     @Parameterized.Parameters
     public static Object[][] parameters() {
+        Map<String, String> buildArgs = new HashMap<>(4);
+        buildArgs.put("BUILD_IMAGE", "alpine:3.16");
+        buildArgs.put("BASE_IMAGE", "alpine");
+        buildArgs.put("BASE_IMAGE_TAG", "3.12");
+        buildArgs.put("UNUSED", "ignored");
+
+        //noinspection deprecation
         return new Object[][] {
             // Dockerfile build without explicit per-file inclusion
             new Object[] {
@@ -38,14 +47,22 @@ public static Object[][] parameters() {
                 "test4567",
                 new ImageFromDockerfile().withFileFromPath(".", RESOURCE_PATH).withDockerfilePath("./Dockerfile-alt"),
             },
-            // Dockerfile build using build args
+            // Dockerfile build using withBuildArg()
             new Object[] {
                 "test7890",
                 new ImageFromDockerfile()
                     .withFileFromPath(".", RESOURCE_PATH)
                     .withDockerfilePath("./Dockerfile-buildarg")
                     .withBuildArg("CUSTOM_ARG", "test7890"),
             },
+            // Dockerfile build using withBuildArgs() with build args in FROM statement
+            new Object[] {
+                "test1234",
+                new ImageFromDockerfile()
+                    .withFileFromPath(".", RESOURCE_PATH)
+                    .withDockerfile(RESOURCE_PATH.resolve("Dockerfile-from-buildarg"))
+                    .withBuildArgs(buildArgs),
+            },
             // Dockerfile build using withDockerfile(File)
             new Object[] {
                 "test4567",
@@ -64,7 +81,7 @@ public DockerfileBuildTest(String expectedFileContent, ImageFromDockerfile image
     @Test
     public void performTest() {
         try (
-            final GenericContainer container = new GenericContainer(image)
+            final GenericContainer<?> container = new GenericContainer<>(image)
                 .withStartupCheckStrategy(new OneShotStartupCheckStrategy())
                 .withCommand("cat", "/test.txt")
         ) {

core/src/test/java/org/testcontainers/images/builder/ImageFromDockerfileTest.java

@@ -32,6 +32,7 @@ public void shouldNotAddSessionLabelIfDeleteOnExitIsFalse() {
         )
             .withDockerfileFromBuilder(it -> it.from("scratch"));
         String imageId = image.resolve();
+
         DockerClient dockerClient = DockerClientFactory.instance().client();
 
         try {

core/src/test/resources/dockerfile-build-test/Dockerfile-from-buildarg

@@ -0,0 +1,9 @@
+ARG BUILD_IMAGE
+ARG BASE_IMAGE
+ARG BASE_IMAGE_TAG
+
+FROM ${BUILD_IMAGE} AS build
+COPY localfile.txt /test-build.txt
+
+FROM $BASE_IMAGE:${BASE_IMAGE_TAG} AS base
+COPY --from=build /test-build.txt /test.txt