From c298d8fa7155d744a4636b3164e759ea2924bc95 Mon Sep 17 00:00:00 2001
From: Richard North <rich.north@gmail.com>
Date: Mon, 4 Oct 2021 11:59:04 +0100
Subject: [PATCH] Prevent Compose image pre-fetching from pulling all tags when
 tag is absent

---
 .../containers/DockerComposeFiles.java        | 24 +++++++++----------
 .../containers/DockerComposeFilesTest.java    |  4 ++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/core/src/main/java/org/testcontainers/containers/DockerComposeFiles.java b/core/src/main/java/org/testcontainers/containers/DockerComposeFiles.java
index 15d035e3673..76754df5118 100644
--- a/core/src/main/java/org/testcontainers/containers/DockerComposeFiles.java
+++ b/core/src/main/java/org/testcontainers/containers/DockerComposeFiles.java
@@ -1,43 +1,43 @@
 package org.testcontainers.containers;
 
+import org.testcontainers.utility.DockerImageName;
+
 import java.io.File;
+import java.util.Collection;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Map.Entry;
 import java.util.Set;
 import java.util.stream.Collectors;
 
 public class DockerComposeFiles {
 
-    private List<ParsedDockerComposeFile> parsedComposeFiles;
+    private final List<ParsedDockerComposeFile> parsedComposeFiles;
 
     public DockerComposeFiles(List<File> composeFiles) {
         this.parsedComposeFiles = composeFiles.stream().map(ParsedDockerComposeFile::new).collect(Collectors.toList());
     }
 
     public Set<String> getDependencyImages() {
-
         Map<String, Set<String>> mergedServiceNameToImageNames = mergeServiceDependencyImageNames();
 
         return getImageNames(mergedServiceNameToImageNames);
     }
 
     private Map<String, Set<String>> mergeServiceDependencyImageNames() {
-        Map<String, Set<String>> mergedServiceNameToImageNames = new HashMap();
+        Map<String, Set<String>> mergedServiceNameToImageNames = new HashMap<>();
         for (ParsedDockerComposeFile parsedComposeFile : parsedComposeFiles) {
-            for (Entry<String, Set<String>> entry : parsedComposeFile.getServiceNameToImageNames().entrySet()) {
-                mergedServiceNameToImageNames.put(entry.getKey(), entry.getValue());
-            }
+            mergedServiceNameToImageNames.putAll(parsedComposeFile.getServiceNameToImageNames());
         }
         return mergedServiceNameToImageNames;
     }
 
     private Set<String> getImageNames(Map<String, Set<String>> serviceToImageNames) {
-        Set<String> imageNames = new HashSet<>();
-        serviceToImageNames.values().stream().forEach(imageNames::addAll);
-        return imageNames;
+        return serviceToImageNames.values().stream()
+            .flatMap(Collection::stream)
+            // Pass through DockerImageName to convert image names to canonical form (e.g. making implicit latest tag explicit)
+            .map(DockerImageName::parse)
+            .map(DockerImageName::asCanonicalNameString)
+            .collect(Collectors.toSet());
     }
-
 }
diff --git a/core/src/test/java/org/testcontainers/containers/DockerComposeFilesTest.java b/core/src/test/java/org/testcontainers/containers/DockerComposeFilesTest.java
index 81586c8290b..c47db0c4b66 100644
--- a/core/src/test/java/org/testcontainers/containers/DockerComposeFilesTest.java
+++ b/core/src/test/java/org/testcontainers/containers/DockerComposeFilesTest.java
@@ -11,7 +11,7 @@ public class DockerComposeFilesTest {
     public void shouldGetDependencyImages() {
         DockerComposeFiles dockerComposeFiles = new DockerComposeFiles(Lists.newArrayList(new File("src/test/resources/docker-compose-imagename-parsing-v2.yml")));
         Assertions.assertThat(dockerComposeFiles.getDependencyImages())
-            .containsExactlyInAnyOrder("postgres", "redis", "mysql");
+            .containsExactlyInAnyOrder("postgres:latest", "redis:latest", "mysql:latest");
     }
 
     @Test
@@ -20,6 +20,6 @@ public void shouldGetDependencyImagesWhenOverriding() {
             Lists.newArrayList(new File("src/test/resources/docker-compose-imagename-overriding-a.yml"),
                                new File("src/test/resources/docker-compose-imagename-overriding-b.yml")));
         Assertions.assertThat(dockerComposeFiles.getDependencyImages())
-            .containsExactlyInAnyOrder("alpine:3.14", "redis:b", "mysql:b", "aservice");
+            .containsExactlyInAnyOrder("alpine:3.14", "redis:b", "mysql:b", "aservice:latest");
     }
 }