Support multiple HTTP status codes for HttpWaitStrategy

CHANGELOG.md

@@ -7,6 +7,8 @@ All notable changes to this project will be documented in this file.
 
 ### Changed
 
+- Support multiple HTTP status codes for HttpWaitStrategy ([\#630](https://github.com/testcontainers/testcontainers-java/issues/630))
+
 ## [1.7.0] - 2018-04-07
 
 ### Fixed

core/src/main/java/org/testcontainers/containers/wait/Wait.java

@@ -39,8 +39,7 @@ public static HostPortWaitStrategy forListeningPort() {
      */
     public static HttpWaitStrategy forHttp(String path) {
         return new HttpWaitStrategy()
-                .forPath(path)
-                .forStatusCode(HttpURLConnection.HTTP_OK);
+                .forPath(path);
     }
 
     /**

core/src/main/java/org/testcontainers/containers/wait/strategy/HttpWaitStrategy.java

@@ -12,6 +12,7 @@
 import java.net.HttpURLConnection;
 import java.net.URI;
 import java.net.URL;
+import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Predicate;
@@ -32,11 +33,16 @@ public class HttpWaitStrategy extends AbstractWaitStrategy {
     private static final String AUTH_BASIC = "Basic ";
 
     private String path = "/";
-    private int statusCode = HttpURLConnection.HTTP_OK;
+    private Set<Integer> statusCodes = new HashSet<>();
     private boolean tlsEnabled;
     private String username;
     private String password;
     private Predicate<String> responsePredicate;
+    private Predicate<Integer> statusCodePredicate = responseCode -> {
+        // If we did not provide any status code, we assume by default HttpURLConnection.HTTP_OK
+        if (statusCodes.isEmpty() && HttpURLConnection.HTTP_OK == responseCode) return true;
+        return statusCodes.contains(responseCode);
+    };
 
     /**
      * Waits for the given status code.
@@ -45,7 +51,17 @@ public class HttpWaitStrategy extends AbstractWaitStrategy {
      * @return this
      */
     public HttpWaitStrategy forStatusCode(int statusCode) {
-        this.statusCode = statusCode;
+        statusCodes.add(statusCode);
+        return this;
+    }
+
+    /**
+     * Waits for the status code to pass the given predicate
+     * @param statusCodePredicate The predicate to test the response against
+     * @return this
+     */
+    public HttpWaitStrategy forStatusCodeMatching(Predicate<Integer> statusCodePredicate) {
+        this.statusCodePredicate = this.statusCodePredicate.or(statusCodePredicate);
         return this;
     }
 
@@ -122,7 +138,7 @@ protected void waitUntilReady() {
                         connection.setRequestMethod("GET");
                         connection.connect();
 
-                        if (statusCode != connection.getResponseCode()) {
+                        if (!statusCodePredicate.test(connection.getResponseCode())) {
                             throw new RuntimeException(String.format("HTTP response code was: %s",
                                 connection.getResponseCode()));
                         }
@@ -144,7 +160,8 @@ protected void waitUntilReady() {
 
         } catch (TimeoutException e) {
             throw new ContainerLaunchException(String.format(
-                "Timed out waiting for URL to be accessible (%s should return HTTP %s)", uri, statusCode));
+                "Timed out waiting for URL to be accessible (%s should return HTTP %s)", uri, statusCodes.isEmpty() ?
+                    HttpURLConnection.HTTP_OK : statusCodes));
         }
     }
 

core/src/main/java/org/testcontainers/containers/wait/strategy/Wait.java

@@ -35,8 +35,7 @@ public static HostPortWaitStrategy forListeningPort() {
      */
     public static HttpWaitStrategy forHttp(String path) {
         return new HttpWaitStrategy()
-                .forPath(path)
-                .forStatusCode(HttpURLConnection.HTTP_OK);
+                .forPath(path);
     }
 
     /**

core/src/test/java/org/testcontainers/junit/wait/AbstractWaitStrategyTest.java → core/src/test/java/org/testcontainers/junit/wait/strategy/AbstractWaitStrategyTest.java

@@ -1,31 +1,31 @@
-package org.testcontainers.junit.wait;
+package org.testcontainers.junit.wait.strategy;
 
 import org.jetbrains.annotations.NotNull;
 import org.junit.Before;
 import org.rnorth.ducttape.RetryCountExceededException;
 import org.rnorth.visibleassertions.VisibleAssertions;
 import org.testcontainers.containers.ContainerLaunchException;
 import org.testcontainers.containers.GenericContainer;
-import org.testcontainers.containers.wait.WaitStrategy;
+import org.testcontainers.containers.wait.strategy.WaitStrategy;
 
 import java.time.Duration;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertTrue;
 
 /**
  * Common test methods for {@link WaitStrategy} implementations.
  *
  * @author Pete Cornish {@literal <outofcoffee@gmail.com>}
  */
 public abstract class AbstractWaitStrategyTest<W extends WaitStrategy> {
-    private static final long WAIT_TIMEOUT_MILLIS = 3000;
-    private static final String IMAGE_NAME = "alpine:latest";
+    static final long WAIT_TIMEOUT_MILLIS = 3000;
+    static final String IMAGE_NAME = "alpine:latest";
 
     /**
      * Indicates that the WaitStrategy has completed waiting successfully.
      */
-    private AtomicBoolean ready;
+    AtomicBoolean ready;
 
     /**
      * Subclasses should return an instance of {@link W} that sets {@code ready} to {@code true},
@@ -50,14 +50,23 @@ public void setUp() throws Exception {
      * @return the (unstarted) container
      */
     private GenericContainer startContainerWithCommand(String shellCommand) {
-        final WaitStrategy waitStrategy = buildWaitStrategy(ready)
-                .withStartupTimeout(Duration.ofMillis(WAIT_TIMEOUT_MILLIS));
+        return startContainerWithCommand(shellCommand, buildWaitStrategy(ready));
+    }
 
+    /**
+     * Starts a GenericContainer with the {@link #IMAGE_NAME} image, passing the given {@code shellCommand} as
+     * a parameter to {@literal sh -c} (the container CMD) and apply a give wait strategy.
+     * Note that the timeout will be overwritten if any with {@link #WAIT_TIMEOUT_MILLIS}.
+     * @param shellCommand the shell command to execute
+     * @param waitStrategy The wait strategy to apply
+     * @return the (unstarted) container
+     */
+    protected GenericContainer startContainerWithCommand(String shellCommand, WaitStrategy waitStrategy) {
         // apply WaitStrategy to container
         return new GenericContainer(IMAGE_NAME)
                 .withExposedPorts(8080)
                 .withCommand("sh", "-c", shellCommand)
-                .waitingFor(waitStrategy);
+                .waitingFor(waitStrategy.withStartupTimeout(Duration.ofMillis(WAIT_TIMEOUT_MILLIS)));
     }
 
     /**
@@ -66,13 +75,7 @@ private GenericContainer startContainerWithCommand(String shellCommand) {
      * @param shellCommand the shell command to execute
      */
     protected void waitUntilReadyAndSucceed(String shellCommand) {
-        final GenericContainer container = startContainerWithCommand(shellCommand);
-
-        // start() blocks until successful or timeout
-        container.start();
-
-        assertTrue(String.format("Expected container to be ready after timeout of %sms",
-                WAIT_TIMEOUT_MILLIS), ready.get());
+        waitUntilReadyAndSucceed(startContainerWithCommand(shellCommand));
     }
 
     /**
@@ -82,12 +85,33 @@ protected void waitUntilReadyAndSucceed(String shellCommand) {
      * @param shellCommand the shell command to execute
      */
     protected void waitUntilReadyAndTimeout(String shellCommand) {
-        final GenericContainer container = startContainerWithCommand(shellCommand);
+        waitUntilReadyAndTimeout(startContainerWithCommand(shellCommand));
+    }
 
+    /**
+     * Expects that the WaitStrategy throws a {@link RetryCountExceededException} after unsuccessful connection
+     * to a container with a listening port.
+     *
+     * @param container the container to start
+     */
+    protected void waitUntilReadyAndTimeout(GenericContainer container) {
         // start() blocks until successful or timeout
         VisibleAssertions.assertThrows("an exception is thrown when timeout occurs (" + WAIT_TIMEOUT_MILLIS + "ms)",
                 ContainerLaunchException.class,
                 container::start);
 
     }
+
+    /**
+     * Expects that the WaitStrategy returns successfully after connection to a container with a listening port.
+     *
+     * @param container the container to start
+     */
+    protected void waitUntilReadyAndSucceed(GenericContainer container) {
+        // start() blocks until successful or timeout
+        container.start();
+
+        assertTrue(String.format("Expected container to be ready after timeout of %sms",
+            WAIT_TIMEOUT_MILLIS), ready.get());
+    }
 }