tidy codes

wait/http.go

@@ -25,6 +25,7 @@ type HTTPStrategy struct {
 	Port              nat.Port
 	Path              string
 	StatusCodeMatcher func(status int) bool
+	ResponseMatcher   func(body io.Reader) bool
 	UseTLS            bool
 	AllowInsecure     bool
 	TLSConfig         *tls.Config // TLS config for HTTPS
@@ -39,6 +40,7 @@ func NewHTTPStrategy(path string) *HTTPStrategy {
 		Port:              "80/tcp",
 		Path:              path,
 		StatusCodeMatcher: defaultStatusCodeMatcher,
+		ResponseMatcher:   func(body io.Reader) bool { return true },
 		UseTLS:            false,
 		TLSConfig:         nil,
 		Method:            http.MethodGet,
@@ -69,6 +71,11 @@ func (ws *HTTPStrategy) WithStatusCodeMatcher(statusCodeMatcher func(status int)
 	return ws
 }
 
+func (ws *HTTPStrategy) WithResponseMatcher(matcher func(body io.Reader) bool) *HTTPStrategy {
+	ws.ResponseMatcher = matcher
+	return ws
+}
+
 func (ws *HTTPStrategy) WithTLS(useTLS bool, tlsconf ...*tls.Config) *HTTPStrategy {
 	ws.UseTLS = useTLS
 	if useTLS && len(tlsconf) > 0 {
@@ -175,8 +182,13 @@ func (ws *HTTPStrategy) WaitUntilReady(ctx context.Context, target StrategyTarge
 			if err != nil {
 				continue
 			}
-			_ = resp.Body.Close()
-			if !ws.StatusCodeMatcher(resp.StatusCode) {
+			if ws.StatusCodeMatcher != nil && !ws.StatusCodeMatcher(resp.StatusCode) {
+				continue
+			}
+			if ws.ResponseMatcher != nil && !ws.ResponseMatcher(resp.Body) {
+				continue
+			}
+			if err := resp.Body.Close(); err != nil {
 				continue
 			}
 			return nil

wait/http_test.go

@@ -6,11 +6,11 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net"
 	"net/http"
 	"os"
-	"strconv"
 	"testing"
 	"time"
 
@@ -36,7 +36,7 @@ func ExampleHTTPStrategy() {
 		panic(err)
 	}
 
-	defer gogs.Terminate(ctx)
+	defer gogs.Terminate(ctx) // nolint: errcheck
 	// Here you have a running container
 
 }
@@ -67,8 +67,12 @@ func TestHTTPStrategyWaitUntilReady(t *testing.T) {
 			Context: workdir + "/testdata",
 		},
 		ExposedPorts: []string{"6443/tcp"},
-		WaitingFor: wait.NewHTTPStrategy("/").WithTLS(true, tlsconfig).
+		WaitingFor: wait.NewHTTPStrategy("/ping").WithTLS(true, tlsconfig).
 			WithStartupTimeout(time.Second * 10).WithPort("6443/tcp").
+			WithResponseMatcher(func(body io.Reader) bool {
+				data, _ := ioutil.ReadAll(body)
+				return bytes.Equal(data, []byte("pong"))
+			}).
 			WithMethod(http.MethodPost).WithBody(bytes.NewReader([]byte("ping"))),
 	}
 
@@ -80,7 +84,7 @@ func TestHTTPStrategyWaitUntilReady(t *testing.T) {
 	}
 	defer container.Terminate(context.Background()) // nolint: errcheck
 
-	ipAddress, err := container.Host(context.Background())
+	host, err := container.Host(context.Background())
 	if err != nil {
 		t.Error(err)
 		return
@@ -106,7 +110,7 @@ func TestHTTPStrategyWaitUntilReady(t *testing.T) {
 			ExpectContinueTimeout: 1 * time.Second,
 		},
 	}
-	resp, err := client.Get(fmt.Sprintf("https://%s/ping", net.JoinHostPort(ipAddress, strconv.Itoa(port.Int()))))
+	resp, err := client.Get(fmt.Sprintf("https://%s:%s", host, port.Port()))
 	if err != nil {
 		t.Error(err)
 		return

wait/testdata/main.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"bytes"
 	"context"
+	"io/ioutil"
 	"log"
 	"net/http"
 	"os"
@@ -17,8 +19,13 @@ func main() {
 	})
 
 	mux.HandleFunc("/ping", func(w http.ResponseWriter, req *http.Request) {
-		w.WriteHeader(http.StatusOK)
-		_, _ = w.Write([]byte("pong"))
+		data, _ := ioutil.ReadAll(req.Body)
+		if bytes.Equal(data, []byte("ping")) {
+			w.WriteHeader(http.StatusOK)
+			_, _ = w.Write([]byte("pong"))
+		} else {
+			w.WriteHeader(http.StatusBadRequest)
+		}
 	})
 
 	server := http.Server{Addr: ":6443", Handler: mux}