Support multiple HTTP status codes for HttpWaitStrategy

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
+        return (!statusCodes.isEmpty() || HttpURLConnection.HTTP_OK == responseCode) &&
+            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.and(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,7 @@ 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));
         }
     }
 

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/strategy/AbstractWaitStrategyTest.java

@@ -0,0 +1,93 @@
+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.strategy.WaitStrategy;
+
+import java.time.Duration;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+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";
+
+    /**
+     * Indicates that the WaitStrategy has completed waiting successfully.
+     */
+    private AtomicBoolean ready;
+
+    /**
+     * Subclasses should return an instance of {@link W} that sets {@code ready} to {@code true},
+     * if the wait was successful.
+     *
+     * @param ready the AtomicBoolean on which to indicate success
+     * @return WaitStrategy implementation
+     */
+    @NotNull
+    protected abstract W buildWaitStrategy(final AtomicBoolean ready);
+
+    @Before
+    public void setUp() throws Exception {
+        ready = new AtomicBoolean(false);
+    }
+
+    /**
+     * Starts a GenericContainer with the {@link #IMAGE_NAME} image, passing the given {@code shellCommand} as
+     * a parameter to {@literal sh -c} (the container CMD).
+     *
+     * @param shellCommand the shell command to execute
+     * @return the (unstarted) container
+     */
+    private GenericContainer startContainerWithCommand(String shellCommand) {
+        final WaitStrategy waitStrategy = buildWaitStrategy(ready)
+                .withStartupTimeout(Duration.ofMillis(WAIT_TIMEOUT_MILLIS));
+
+        // apply WaitStrategy to container
+        return new GenericContainer(IMAGE_NAME)
+                .withExposedPorts(8080)
+                .withCommand("sh", "-c", shellCommand)
+                .waitingFor(waitStrategy);
+    }
+
+    /**
+     * Expects that the WaitStrategy returns successfully after connection to a container with a listening port.
+     *
+     * @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());
+    }
+
+    /**
+     * Expects that the WaitStrategy throws a {@link RetryCountExceededException} after unsuccessful connection
+     * to a container with a listening port.
+     *
+     * @param shellCommand the shell command to execute
+     */
+    protected void waitUntilReadyAndTimeout(String shellCommand) {
+        final GenericContainer container = startContainerWithCommand(shellCommand);
+
+        // start() blocks until successful or timeout
+        VisibleAssertions.assertThrows("an exception is thrown when timeout occurs (" + WAIT_TIMEOUT_MILLIS + "ms)",
+                ContainerLaunchException.class,
+                container::start);
+
+    }
+}

core/src/test/java/org/testcontainers/junit/wait/strategy/HttpWaitStrategyTest.java

@@ -0,0 +1,74 @@
+package org.testcontainers.junit.wait.strategy;
+
+import org.jetbrains.annotations.NotNull;
+import org.junit.Test;
+import org.rnorth.ducttape.RetryCountExceededException;
+import org.testcontainers.containers.wait.strategy.HttpWaitStrategy;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Tests for {@link HttpWaitStrategy}.
+ *
+ * @author Pete Cornish {@literal <outofcoffee@gmail.com>}
+ */
+public class HttpWaitStrategyTest extends AbstractWaitStrategyTest<HttpWaitStrategy> {
+    /**
+     * newline sequence indicating end of the HTTP header.
+     */
+    private static final String NEWLINE = "\r\n";
+
+    private static final String GOOD_RESPONSE_BODY = "Good Response Body";
+
+    /**
+     * Expects that the WaitStrategy returns successfully after receiving an HTTP 200 response from the container.
+     */
+    @Test
+    public void testWaitUntilReadyWithSuccess() {
+        waitUntilReadyAndSucceed(createShellCommand("200 OK", GOOD_RESPONSE_BODY));
+    }
+
+    /**
+     * Expects that the WaitStrategy throws a {@link RetryCountExceededException} after not receiving an HTTP 200
+     * response from the container within the timeout period.
+     */
+    @Test
+    public void testWaitUntilReadyWithTimeout() {
+        waitUntilReadyAndTimeout(createShellCommand("400 Bad Request", GOOD_RESPONSE_BODY));
+    }
+
+    /**
+     * Expects that the WaitStrategy throws a {@link RetryCountExceededException} after not the expected response body
+     * from the container within the timeout period.
+     */
+    @Test
+    public void testWaitUntilReadyWithTimeoutAndBadResponseBody() {
+        waitUntilReadyAndTimeout(createShellCommand("200 OK", "Bad Response"));
+    }
+
+    /**
+     * @param ready the AtomicBoolean on which to indicate success
+     * @return the WaitStrategy under test
+     */
+    @NotNull
+    protected HttpWaitStrategy buildWaitStrategy(final AtomicBoolean ready) {
+        return new HttpWaitStrategy() {
+            @Override
+            protected void waitUntilReady() {
+                // blocks until ready or timeout occurs
+                super.waitUntilReady();
+                ready.set(true);
+            }
+        }
+        .forResponsePredicate(s -> s.equals(GOOD_RESPONSE_BODY))
+            .forStatusCodeMatching(it -> it >= 200 && it < 300 || it == 401);
+    }
+
+    private String createShellCommand(String header, String responseBody) {
+        int length = responseBody.getBytes().length;
+        return "while true; do { echo -e \"HTTP/1.1 "+header+NEWLINE+
+                "Content-Type: text/html"+NEWLINE+
+                "Content-Length: "+length +NEWLINE+ "\";"
+                +" echo \""+responseBody+"\";} | nc -lp 8080; done";
+    }
+}

docs/usage/docker_compose.md

@@ -72,14 +72,15 @@ public static DockerComposeContainer environment =
                 Wait.forListeningPort().withStartupTimeout(Duration.ofSeconds(30)));
 ```
 
-Wait for arbitrary status code on an HTTPS endpoint:
+Wait for arbitrary status codes on an HTTPS endpoint:
 ```java
 @ClassRule
 public static DockerComposeContainer environment =
     new DockerComposeContainer(new File("src/test/resources/compose-test.yml"))
             .withExposedService("elasticsearch_1", ELASTICSEARCH_PORT, 
                 Wait.forHttp("/all")
-                    .forStatusCode(301)
+                    .forStatusCode(200)
+                    .forStatusCode(401)
                     .usingTls());
 ```
 
@@ -91,7 +92,8 @@ public static DockerComposeContainer environment =
             .withExposedService("redis_1", REDIS_PORT, Wait.forListeningPort())
             .withExposedService("elasticsearch_1", ELASTICSEARCH_PORT, 
                 Wait.forHttp("/all")
-                    .forStatusCode(301)
+                    .forStatusCode(200)
+                    .forStatusCode(401)
                     .usingTls());
 ```
 

docs/usage/options.md

@@ -97,17 +97,30 @@ public static GenericContainer elasticsearch =
                .waitingFor(Wait.forHttp("/all"));
 ```
 
-Wait for arbitrary status code on an HTTPS endpoint:
+Wait for 200 or 401 status codes on an HTTPS endpoint:
 ```java
 @ClassRule
 public static GenericContainer elasticsearch =
     new GenericContainer("elasticsearch:2.3")
                .withExposedPorts(9200)
                .waitingFor(
                		Wait.forHttp("/all")
-               			 .forStatusCode(301)
+               			 .forStatusCode(200)
+               			 .forStatusCode(401)
                			 .usingTls());
- ```
+```
+
+Wait for 200...299 or 401 status codes on an HTTPS endpoint:
+```java
+@ClassRule
+public static GenericContainer elasticsearch =
+    new GenericContainer("elasticsearch:2.3")
+               .withExposedPorts(9200)
+               .waitingFor(
+               		Wait.forHttp("/all")
+               			 .forStatusCodeMatching(it -> it >= 200 && it < 300 || it == 401)
+               			 .usingTls());
+```
 
 If the used image supports Docker's [Healthcheck](https://docs.docker.com/engine/reference/builder/#healthcheck) feature, you can directly leverage the `healthy` state of the container as your wait condition:
 ```java

modules/vault/src/test/java/org/testcontainers/vault/VaultContainerTest.java

@@ -3,7 +3,7 @@
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.testcontainers.containers.GenericContainer;
-import org.testcontainers.containers.wait.Wait;
+import org.testcontainers.containers.wait.strategy.Wait;
 
 import java.io.IOException;