Add getMappedPort(int, InternetProtocol) for UDP port lookups

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

@@ -158,6 +158,24 @@ default Integer getFirstMappedPort() {
      * @see #getCurrentContainerInfo()
      */
     default Integer getMappedPort(int originalPort) {
+        return getMappedPort(originalPort, InternetProtocol.TCP);
+    }
+
+    /**
+     * Get the actual mapped port for a given port exposed by the container, for a specific protocol.
+     * It should be used in conjunction with {@link #getHost()}.
+     * <p>
+     * Note: The returned port number might be outdated (for instance, after disconnecting from a network and reconnecting
+     * again). If you always need up-to-date value, override the {@link #getContainerInfo()} to return the
+     * {@link #getCurrentContainerInfo()}.
+     *
+     * @param originalPort the original port that is exposed
+     * @param protocol the protocol (TCP or UDP) that the port is exposed with
+     * @return the port that the exposed port is mapped to, or null if it is not exposed
+     * @see #getContainerInfo()
+     * @see #getCurrentContainerInfo()
+     */
+    default Integer getMappedPort(int originalPort, InternetProtocol protocol) {
         Preconditions.checkState(
             this.getContainerId() != null,
             "Mapped port can only be obtained after the container is started"
@@ -166,13 +184,17 @@ default Integer getMappedPort(int originalPort) {
         Ports.Binding[] binding = new Ports.Binding[0];
         final InspectContainerResponse containerInfo = this.getContainerInfo();
         if (containerInfo != null) {
-            binding = containerInfo.getNetworkSettings().getPorts().getBindings().get(new ExposedPort(originalPort));
+            ExposedPort exposedPort = new ExposedPort(
+                originalPort,
+                com.github.dockerjava.api.model.InternetProtocol.parse(protocol.name())
+            );
+            binding = containerInfo.getNetworkSettings().getPorts().getBindings().get(exposedPort);
         }
 
         if (binding != null && binding.length > 0 && binding[0] != null) {
             return Integer.valueOf(binding[0].getHostPortSpec());
         } else {
-            throw new IllegalArgumentException("Requested port (" + originalPort + ") is not mapped");
+            throw new IllegalArgumentException("Requested port (" + originalPort + "/" + protocol.toDockerNotation() + ") is not mapped");
         }
     }
 

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

@@ -74,7 +74,6 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -263,6 +262,23 @@ public List<Integer> getExposedPorts() {
         return exposedPorts;
     }
 
+    /**
+     * {@inheritDoc}
+     * <p>
+     * Unlike the default implementation, this preserves the protocol the first port was exposed with,
+     * so it also works for containers whose first exposed port is UDP.
+     */
+    @Override
+    public Integer getFirstMappedPort() {
+        return this.containerDef.getExposedPorts()
+            .stream()
+            .findFirst()
+            .map(exposedPort ->
+                getMappedPort(exposedPort.getPort(), InternetProtocol.fromDockerNotation(exposedPort.getProtocol().toString()))
+            )
+            .orElseThrow(() -> new IllegalStateException("Container doesn't expose any ports"));
+    }
+
     @Override
     public void setExposedPorts(List<Integer> exposedPorts) {
         this.containerDef.exposedPorts.clear();
@@ -739,9 +755,13 @@ protected void containerIsStopped(InspectContainerResponse containerInfo) {}
     @Deprecated
     protected Integer getLivenessCheckPort() {
         // legacy implementation for backwards compatibility
-        Iterator<ExposedPort> exposedPortsIterator = this.containerDef.getExposedPorts().iterator();
-        if (exposedPortsIterator.hasNext()) {
-            return getMappedPort(exposedPortsIterator.next().getPort());
+        // UDP ports cannot be checked with a TCP connection, so only TCP ports are candidates
+        Optional<ExposedPort> firstTcpPort = this.containerDef.getExposedPorts()
+            .stream()
+            .filter(exposedPort -> exposedPort.getProtocol() == com.github.dockerjava.api.model.InternetProtocol.TCP)
+            .findFirst();
+        if (firstTcpPort.isPresent()) {
+            return getMappedPort(firstTcpPort.get().getPort());
         } else if (!this.containerDef.getPortBindings().isEmpty()) {
             return Integer.valueOf(
                 this.containerDef.getPortBindings().iterator().next().getBinding().getHostPortSpec()
@@ -759,7 +779,14 @@ protected Integer getLivenessCheckPort() {
     @NonNull
     @Deprecated
     protected Set<Integer> getLivenessCheckPorts() {
-        final Set<Integer> result = WaitStrategyTarget.super.getLivenessCheckPortNumbers();
+        // Ports exposed only as UDP have no TCP mapping and cannot be checked by the default wait strategy
+        final Set<Integer> udpOnlyPorts = getUdpOnlyExposedPorts();
+        final Set<Integer> result = getExposedPorts()
+            .stream()
+            .filter(port -> !udpOnlyPorts.contains(port))
+            .map(this::getMappedPort)
+            .collect(Collectors.toSet());
+        result.addAll(getBoundPortNumbers());
 
         // for backwards compatibility
         if (this.getLivenessCheckPort() != null) {
@@ -769,6 +796,20 @@ protected Set<Integer> getLivenessCheckPorts() {
         return result;
     }
 
+    private Set<Integer> getUdpOnlyExposedPorts() {
+        final Set<Integer> udpPorts = new HashSet<>();
+        final Set<Integer> tcpPorts = new HashSet<>();
+        for (ExposedPort exposedPort : this.containerDef.getExposedPorts()) {
+            if (exposedPort.getProtocol() == com.github.dockerjava.api.model.InternetProtocol.UDP) {
+                udpPorts.add(exposedPort.getPort());
+            } else if (exposedPort.getProtocol() == com.github.dockerjava.api.model.InternetProtocol.TCP) {
+                tcpPorts.add(exposedPort.getPort());
+            }
+        }
+        udpPorts.removeAll(tcpPorts);
+        return udpPorts;
+    }
+
     @Override
     public Set<Integer> getLivenessCheckPortNumbers() {
         return this.getLivenessCheckPorts();
@@ -1048,6 +1089,30 @@ public void addExposedPorts(int... ports) {
         this.containerDef.addExposedTcpPorts(ports);
     }
 
+    /**
+     * Expose a container port using a specific protocol, so that a randomly chosen host port will be bound to it
+     * when the container starts. Use {@link ContainerState#getMappedPort(int, InternetProtocol)} to retrieve the
+     * bound host port once the container is running.
+     *
+     * @param port the container port to expose
+     * @param protocol the protocol (TCP or UDP) that the port should be exposed with
+     */
+    public void addExposedPort(int port, InternetProtocol protocol) {
+        this.containerDef.addExposedPort(port, com.github.dockerjava.api.model.InternetProtocol.parse(protocol.name()));
+    }
+
+    /**
+     * Fluent variant of {@link #addExposedPort(int, InternetProtocol)}.
+     *
+     * @param port the container port to expose
+     * @param protocol the protocol (TCP or UDP) that the port should be exposed with
+     * @return this
+     */
+    public SELF withExposedPort(int port, InternetProtocol protocol) {
+        addExposedPort(port, protocol);
+        return self();
+    }
+
     /**
      * {@inheritDoc}
      */

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

@@ -1,12 +1,20 @@
 package org.testcontainers.containers;
 
+import com.github.dockerjava.api.command.InspectContainerResponse;
+import com.github.dockerjava.api.model.ExposedPort;
+import com.github.dockerjava.api.model.Ports;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.mockito.Answers;
+import org.mockito.Mockito;
 
 import java.util.Collections;
 import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -35,4 +43,37 @@ void test(String name, String testSet, List<Integer> expectedResult) {
         List<Integer> result = containerState.getBoundPortNumbers();
         assertThat(result).hasSameElementsAs(expectedResult);
     }
+
+    @Test
+    void getMappedPortWithProtocolLooksUpTheBindingForThatProtocol() {
+        ContainerState containerState = mock(ContainerState.class);
+        doCallRealMethod().when(containerState).getMappedPort(anyInt(), any());
+        when(containerState.getContainerId()).thenReturn("container-id");
+
+        InspectContainerResponse containerInfo = Mockito.mock(InspectContainerResponse.class, Answers.RETURNS_DEEP_STUBS);
+        ExposedPort udpPort = new ExposedPort(12345, com.github.dockerjava.api.model.InternetProtocol.UDP);
+        when(containerInfo.getNetworkSettings().getPorts().getBindings())
+            .thenReturn(Collections.singletonMap(udpPort, new Ports.Binding[] { Ports.Binding.bindPort(54321) }));
+        when(containerState.getContainerInfo()).thenReturn(containerInfo);
+
+        Integer mappedPort = containerState.getMappedPort(12345, InternetProtocol.UDP);
+
+        assertThat(mappedPort).isEqualTo(54321);
+    }
+
+    @Test
+    void getMappedPortWithProtocolThrowsWhenNotMapped() {
+        ContainerState containerState = mock(ContainerState.class);
+        doCallRealMethod().when(containerState).getMappedPort(anyInt(), any());
+        when(containerState.getContainerId()).thenReturn("container-id");
+
+        InspectContainerResponse containerInfo = Mockito.mock(InspectContainerResponse.class, Answers.RETURNS_DEEP_STUBS);
+        when(containerInfo.getNetworkSettings().getPorts().getBindings()).thenReturn(Collections.emptyMap());
+        when(containerState.getContainerInfo()).thenReturn(containerInfo);
+
+        org.assertj.core.api.Assertions
+            .assertThatThrownBy(() -> containerState.getMappedPort(12345, InternetProtocol.UDP))
+            .isInstanceOf(IllegalArgumentException.class)
+            .hasMessageContaining("12345/udp");
+    }
 }

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

@@ -0,0 +1,83 @@
+package org.testcontainers.containers;
+
+import com.github.dockerjava.api.command.InspectContainerResponse;
+import com.github.dockerjava.api.model.ExposedPort;
+import com.github.dockerjava.api.model.Ports;
+import org.junit.jupiter.api.Test;
+import org.mockito.Answers;
+import org.mockito.Mockito;
+import org.testcontainers.utility.DockerImageName;
+
+import java.util.AbstractMap;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.when;
+
+class GenericContainerLivenessCheckPortsTest {
+
+    @Test
+    void udpOnlyPortsAreNotUsedForLivenessChecks() {
+        GenericContainer<?> container = startedWithBindings(
+            binding(12345, com.github.dockerjava.api.model.InternetProtocol.UDP, 54321)
+        )
+            .withExposedPort(12345, InternetProtocol.UDP);
+
+        assertThat(container.getLivenessCheckPortNumbers()).isEmpty();
+    }
+
+    @Test
+    void tcpPortsAreStillUsedWhenUdpPortsAreAlsoExposed() {
+        GenericContainer<?> container = startedWithBindings(
+            binding(12345, com.github.dockerjava.api.model.InternetProtocol.UDP, 54321),
+            binding(8080, com.github.dockerjava.api.model.InternetProtocol.TCP, 58080)
+        )
+            .withExposedPort(12345, InternetProtocol.UDP)
+            .withExposedPort(8080, InternetProtocol.TCP);
+
+        assertThat(container.getLivenessCheckPortNumbers()).containsExactly(58080);
+    }
+
+    @SafeVarargs
+    private static GenericContainer<?> startedWithBindings(Map.Entry<ExposedPort, Ports.Binding[]>... bindings) {
+        Map<ExposedPort, Ports.Binding[]> bindingMap = new HashMap<>();
+        for (Map.Entry<ExposedPort, Ports.Binding[]> binding : bindings) {
+            bindingMap.put(binding.getKey(), binding.getValue());
+        }
+        InspectContainerResponse containerInfo = Mockito.mock(InspectContainerResponse.class, Answers.RETURNS_DEEP_STUBS);
+        when(containerInfo.getNetworkSettings().getPorts().getBindings()).thenReturn(bindingMap);
+        return new StartedContainer(containerInfo);
+    }
+
+    private static class StartedContainer extends GenericContainer<StartedContainer> {
+
+        private final InspectContainerResponse containerInfo;
+
+        StartedContainer(InspectContainerResponse containerInfo) {
+            super(DockerImageName.parse("alpine:3.20"));
+            this.containerInfo = containerInfo;
+        }
+
+        @Override
+        public String getContainerId() {
+            return "container-id";
+        }
+
+        @Override
+        public InspectContainerResponse getContainerInfo() {
+            return containerInfo;
+        }
+    }
+
+    private static Map.Entry<ExposedPort, Ports.Binding[]> binding(
+        int port,
+        com.github.dockerjava.api.model.InternetProtocol protocol,
+        int hostPort
+    ) {
+        return new AbstractMap.SimpleEntry<>(
+            new ExposedPort(port, protocol),
+            new Ports.Binding[] { Ports.Binding.bindPort(hostPort) }
+        );
+    }
+}