Issue 554 added udp support for generic container and docker compose container

Author: mkrzywanski

core/build.gradle

@@ -98,6 +98,8 @@ task japicmp(type: me.champeau.gradle.japicmp.JapicmpTask) {
         "org.testcontainers.containers.wait.LogMessageWaitStrategy",
         "org.testcontainers.containers.wait.Wait",
         "org.testcontainers.containers.wait.WaitAllStrategy",
+        "org.testcontainers.containers.ContainerState",
+        "org.testcontainers.containers.wait.strategy.WaitStrategyTarget",
         "org.testcontainers.containers.wait.WaitStrategy",
         "org.testcontainers.dockerclient.AuditLoggingDockerClient",
         "org.testcontainers.dockerclient.LogToStringContainerCallback",

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

@@ -8,10 +8,11 @@
 import org.testcontainers.DockerClientFactory;
 import org.testcontainers.containers.wait.strategy.WaitStrategyTarget;
 
-import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 
 /**
@@ -23,31 +24,44 @@ class ComposeServiceWaitStrategyTarget implements WaitStrategyTarget {
     private final Container container;
     private final GenericContainer proxyContainer;
     @NonNull
-    private Map<Integer, Integer> mappedPorts;
-    @Getter(lazy=true)
+    private Map<Port, Integer> mappedPorts;
+    @Getter(lazy = true)
     private final InspectContainerResponse containerInfo = DockerClientFactory.instance().client().inspectContainerCmd(getContainerId()).exec();
 
     ComposeServiceWaitStrategyTarget(Container container, GenericContainer proxyContainer,
-                                     @NonNull Map<Integer, Integer> mappedPorts) {
+                                     @NonNull Map<Port, Integer> mappedPorts) {
         this.container = container;
         this.proxyContainer = proxyContainer;
-        this.mappedPorts = new HashMap<>(mappedPorts);
+        this.mappedPorts = mappedPorts;
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
     public List<Integer> getExposedPorts() {
-        return new ArrayList<>(this.mappedPorts.keySet());
+        return this.mappedPorts.keySet()
+            .stream()
+            .map(Port::getValue)
+            .collect(Collectors.toList());
+    }
+
+    @Override
+    public Set<Port> exposedPorts() {
+        return new HashSet<>(mappedPorts.keySet());
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
     public Integer getMappedPort(int originalPort) {
-        return this.proxyContainer.getMappedPort(this.mappedPorts.get(originalPort));
+        return this.getMappedPort(originalPort, InternetProtocol.TCP);
+    }
+
+    @Override
+    public Integer getMappedPort(int originalPort, InternetProtocol internetProtocol) {
+        return this.proxyContainer.getMappedPort(this.mappedPorts.get(Port.of(originalPort, internetProtocol)), internetProtocol);
     }
 
     /**

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

@@ -7,18 +7,19 @@
 import lombok.NonNull;
 import lombok.Value;
 import org.testcontainers.DockerClientFactory;
-import org.testcontainers.images.ImagePullPolicy;
 import org.testcontainers.containers.output.OutputFrame;
 import org.testcontainers.containers.startupcheck.StartupCheckStrategy;
 import org.testcontainers.containers.traits.LinkableContainer;
 import org.testcontainers.containers.wait.strategy.WaitStrategy;
+import org.testcontainers.images.ImagePullPolicy;
 import org.testcontainers.utility.LogUtils;
 import org.testcontainers.utility.MountableFile;
 
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.Future;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -85,9 +86,9 @@ default void addFileSystemBind(final String hostPath, final String containerPath
      * Adds a file system binding. Consider using {@link #withFileSystemBind(String, String, BindMode)}
      * for building a container in a fluent style.
      *
-     * @param hostPath      the file system path on the host
-     * @param containerPath the file system path inside the container
-     * @param mode          the bind mode
+     * @param hostPath       the file system path on the host
+     * @param containerPath  the file system path inside the container
+     * @param mode           the bind mode
      * @param selinuxContext selinux context argument to use for this file
      */
     void addFileSystemBind(String hostPath, String containerPath, BindMode mode, SelinuxContext selinuxContext);
@@ -96,7 +97,7 @@ default void addFileSystemBind(final String hostPath, final String containerPath
      * Add a link to another container.
      *
      * @param otherContainer the other container object to link to
-     * @param alias the alias (for the other container) that this container should be able to use
+     * @param alias          the alias (for the other container) that this container should be able to use
      * @deprecated Links are deprecated (see <a href="https://github.com/testcontainers/testcontainers-java/issues/465">#465</a>). Please use {@link Network} features instead.
      */
     @Deprecated
@@ -110,6 +111,8 @@ default void addFileSystemBind(final String hostPath, final String containerPath
      */
     void addExposedPort(Integer port);
 
+    void addExposedPort(Integer port, InternetProtocol internetProtocol);
+
     /**
      * Add exposed ports. Consider using {@link #withExposedPorts(Integer...)}
      * for building a container in a fluent style.
@@ -118,19 +121,21 @@ default void addFileSystemBind(final String hostPath, final String containerPath
      */
     void addExposedPorts(int... ports);
 
+    void addExposedPorts(Set<Integer> ports, InternetProtocol internetProtocol);
+
     /**
      * Specify the {@link WaitStrategy} to use to determine if the container is ready.
      *
-     * @see org.testcontainers.containers.wait.strategy.Wait#defaultWaitStrategy()
      * @param waitStrategy the WaitStrategy to use
      * @return this
+     * @see org.testcontainers.containers.wait.strategy.Wait#defaultWaitStrategy()
      */
     SELF waitingFor(@NonNull WaitStrategy waitStrategy);
 
     /**
      * Adds a file system binding.
      *
-     * @param hostPath the file system path on the host
+     * @param hostPath      the file system path on the host
      * @param containerPath the file system path inside the container
      * @return this
      */
@@ -141,9 +146,9 @@ default SELF withFileSystemBind(String hostPath, String containerPath) {
     /**
      * Adds a file system binding.
      *
-     * @param hostPath the file system path on the host
+     * @param hostPath      the file system path on the host
      * @param containerPath the file system path inside the container
-     * @param mode the bind mode
+     * @param mode          the bind mode
      * @return this
      */
     SELF withFileSystemBind(String hostPath, String containerPath, BindMode mode);
@@ -152,7 +157,7 @@ default SELF withFileSystemBind(String hostPath, String containerPath) {
      * Adds container volumes.
      *
      * @param container the container to add volumes from
-     * @param mode the bind mode
+     * @param mode      the bind mode
      * @return this
      */
     SELF withVolumesFrom(Container container, BindMode mode);
@@ -165,6 +170,8 @@ default SELF withFileSystemBind(String hostPath, String containerPath) {
      */
     SELF withExposedPorts(Integer... ports);
 
+    SELF withExposedPorts(Set<Integer> ports, InternetProtocol internetProtocol);
+
     /**
      * Set the file to be copied before starting a created container
      *
@@ -186,7 +193,7 @@ default SELF withFileSystemBind(String hostPath, String containerPath) {
     /**
      * Add an environment variable to be passed to the container.
      *
-     * @param key   environment variable key
+     * @param key    environment variable key
      * @param mapper environment variable value mapper, accepts old value as an argument
      * @return this
      */
@@ -214,6 +221,7 @@ default SELF withEnv(String key, Function<Optional<String>, String> mapper) {
 
     /**
      * Add labels to the container.
+     *
      * @param labels map of labels
      * @return this
      */
@@ -237,7 +245,8 @@ default SELF withEnv(String key, Function<Optional<String>, String> mapper) {
 
     /**
      * Add an extra host entry to be passed to the container
-     * @param hostname hostname to use for this hosts file entry
+     *
+     * @param hostname  hostname to use for this hosts file entry
      * @param ipAddress IP address to use for this hosts file entry
      * @return this
      */
@@ -272,6 +281,7 @@ default SELF withEnv(String key, Function<Optional<String>, String> mapper) {
 
     /**
      * Set the image pull policy of the container
+     *
      * @return
      */
     SELF withImagePullPolicy(ImagePullPolicy policy);
@@ -304,15 +314,16 @@ default SELF withClasspathResourceMapping(final String resourcePath, final Strin
 
     /**
      * Set the duration of waiting time until container treated as started.
-     * @see WaitStrategy#waitUntilReady(org.testcontainers.containers.wait.strategy.WaitStrategyTarget)
      *
      * @param startupTimeout timeout
      * @return this
+     * @see WaitStrategy#waitUntilReady(org.testcontainers.containers.wait.strategy.WaitStrategyTarget)
      */
     SELF withStartupTimeout(Duration startupTimeout);
 
     /**
      * Set the privilegedMode mode for the container
+     *
      * @param mode boolean
      * @return this
      */
@@ -406,7 +417,6 @@ default void followOutput(Consumer<OutputFrame> consumer, OutputFrame.OutputType
     Future<String> getImage();
 
     /**
-     *
      * @deprecated use getEnvMap
      */
     @Deprecated
@@ -428,6 +438,8 @@ default void followOutput(Consumer<OutputFrame> consumer, OutputFrame.OutputType
 
     void setExposedPorts(List<Integer> exposedPorts);
 
+    void setExposedPorts(Set<Port> exposedPorts);
+
     void setPortBindings(List<String> portBindings);
 
     void setExtraHosts(List<String> extraHosts);

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

@@ -12,6 +12,7 @@
 import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
 import org.apache.commons.compress.archivers.tar.TarArchiveOutputStream;
 import org.apache.commons.compress.utils.IOUtils;
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testcontainers.DockerClientFactory;
@@ -30,9 +31,11 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 public interface ContainerState {
@@ -138,12 +141,17 @@ default Integer getFirstMappedPort() {
      * @return the port that the exposed port is mapped to, or null if it is not exposed
      */
     default Integer getMappedPort(int originalPort) {
+        return getMappedPort(originalPort, InternetProtocol.TCP);
+    }
+
+    default Integer getMappedPort(int originalPort, InternetProtocol internetProtocol) {
         Preconditions.checkState(this.getContainerId() != null, "Mapped port can only be obtained after the container is started");
 
         Ports.Binding[] binding = new Ports.Binding[0];
         final InspectContainerResponse containerInfo = this.getContainerInfo();
         if (containerInfo != null) {
-            binding = containerInfo.getNetworkSettings().getPorts().getBindings().get(new ExposedPort(originalPort));
+            com.github.dockerjava.api.model.InternetProtocol internetProtocol1 = com.github.dockerjava.api.model.InternetProtocol.parse(internetProtocol.toDockerNotation());
+            binding = containerInfo.getNetworkSettings().getPorts().getBindings().get(new ExposedPort(originalPort, internetProtocol1));
         }
 
         if (binding != null && binding.length > 0 && binding[0] != null) {
@@ -158,6 +166,8 @@ default Integer getMappedPort(int originalPort) {
      */
     List<Integer> getExposedPorts();
 
+    Set<Port> exposedPorts();
+
     /**
      * @return the port bindings
      */
@@ -185,6 +195,21 @@ default List<Integer> getBoundPortNumbers() {
             .collect(Collectors.toList());
     }
 
+    default Set<Port> getBoundPorts() {
+        final Ports hostPortBindings = this.getContainerInfo().getHostConfig().getPortBindings();
+        Set<Port> ports = new LinkedHashSet<>();
+        for (Map.Entry<ExposedPort, Ports.Binding[]> binding : hostPortBindings.getBindings().entrySet()) {
+            for (Ports.Binding portBinding : binding.getValue()) {
+                String hostPortSpec = portBinding.getHostPortSpec();
+                if(StringUtils.isNotEmpty(hostPortSpec)) {
+                    InternetProtocol internetProtocol = InternetProtocol.fromDockerNotation(binding.getKey().getProtocol().toString());
+                    ports.add(Port.of(Integer.parseInt(hostPortSpec), internetProtocol));
+                }
+            }
+        }
+        return ports;
+    }
+
 
     /**
      * @return all log output from the container from start until the current instant (both stdout and stderr)

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

@@ -86,7 +86,7 @@ public class DockerComposeContainer<SELF extends DockerComposeContainer<SELF>> e
     private String project;
 
     private final AtomicInteger nextAmbassadorPort = new AtomicInteger(2000);
-    private final Map<String, Map<Integer, Integer>> ambassadorPortMappings = new ConcurrentHashMap<>();
+    private final Map<String, Map<Port, Integer>> ambassadorPortMappings = new ConcurrentHashMap<>();
     private final Map<String, ComposeServiceWaitStrategyTarget> serviceInstanceMap = new ConcurrentHashMap<>();
     private final Map<String, WaitAllStrategy> waitStrategyMap = new ConcurrentHashMap<>();
     private final SocatContainer ambassadorContainer = new SocatContainer();
@@ -252,8 +252,11 @@ private void waitUntilServiceStarted() {
 
     private void createServiceInstance(Container container) {
         String serviceName = getServiceNameFromContainer(container);
-        final ComposeServiceWaitStrategyTarget containerInstance = new ComposeServiceWaitStrategyTarget(container,
-            ambassadorContainer, ambassadorPortMappings.getOrDefault(serviceName, new HashMap<>()));
+        final ComposeServiceWaitStrategyTarget containerInstance = new ComposeServiceWaitStrategyTarget(
+            container,
+            ambassadorContainer,
+            ambassadorPortMappings.getOrDefault(serviceName, new HashMap<>())
+        );
 
         String containerId = containerInstance.getContainerId();
         if (tailChildContainers) {
@@ -348,7 +351,7 @@ public DockerComposeContainer withExposedService(String serviceName, int instanc
         return withExposedService(serviceName + "_" + instance, servicePort, waitStrategy);
     }
 
-    public SELF withExposedService(String serviceName, int servicePort, @NonNull WaitStrategy waitStrategy) {
+    public SELF withExposedService(String serviceName, int servicePort, InternetProtocol internetProtocol,@NonNull WaitStrategy waitStrategy) {
 
         String serviceInstanceName = getServiceInstanceName(serviceName);
 
@@ -368,13 +371,17 @@ public SELF withExposedService(String serviceName, int servicePort, @NonNull Wai
 
         // Ambassador container will be started together after docker compose has started
         int ambassadorPort = nextAmbassadorPort.getAndIncrement();
-        ambassadorPortMappings.computeIfAbsent(serviceInstanceName, __ -> new ConcurrentHashMap<>()).put(servicePort, ambassadorPort);
-        ambassadorContainer.withTarget(ambassadorPort, serviceInstanceName, servicePort);
+        ambassadorPortMappings.computeIfAbsent(serviceInstanceName, __ -> new ConcurrentHashMap<>()).put(Port.of(servicePort, internetProtocol), ambassadorPort);
+        ambassadorContainer.withTarget(ambassadorPort, serviceInstanceName, servicePort, internetProtocol);
         ambassadorContainer.addLink(new FutureContainer(this.project + "_" + serviceInstanceName), serviceInstanceName);
         addWaitStrategy(serviceInstanceName, waitStrategy);
         return self();
     }
 
+    public SELF withExposedService(String serviceName, int servicePort, @NonNull WaitStrategy waitStrategy) {
+        return withExposedService(serviceName, servicePort, InternetProtocol.TCP, waitStrategy);
+    }
+
     private String getServiceInstanceName(String serviceName) {
         String serviceInstanceName = serviceName;
         if (!serviceInstanceName.matches(".*_[0-9]+")) {
@@ -433,14 +440,19 @@ public String getServiceHost(String serviceName, Integer servicePort) {
      * @return a port that can be used for accessing the service container.
      */
     public Integer getServicePort(String serviceName, Integer servicePort) {
-        Map<Integer, Integer> portMap = ambassadorPortMappings.get(getServiceInstanceName(serviceName));
+        return getServicePort(serviceName, servicePort, InternetProtocol.TCP);
+    }
+
+    public Integer getServicePort(String serviceName, Integer servicePort, InternetProtocol internetProtocol) {
+        Map<Port, Integer> portMap = ambassadorPortMappings.get(getServiceInstanceName(serviceName));
 
         if (portMap == null) {
             throw new IllegalArgumentException("Could not get a port for '" + serviceName + "'. " +
-                                               "Testcontainers does not have an exposed port configured for '" + serviceName + "'. " +
-                                               "To fix, please ensure that the service '" + serviceName + "' has ports exposed using .withExposedService(...)");
+                "Testcontainers does not have an exposed port configured for '" + serviceName + "'. " +
+                "To fix, please ensure that the service '" + serviceName + "' has ports exposed using .withExposedService(...)");
         } else {
-            return ambassadorContainer.getMappedPort(portMap.get(servicePort));
+            Port svcPort = Port.of(servicePort, internetProtocol);
+            return ambassadorContainer.getMappedPort(portMap.get(svcPort), internetProtocol);
         }
     }