From 9e86f4544fd1c5d212daf56f5b770c0ece330bdb Mon Sep 17 00:00:00 2001
From: GaruGaru <garuglieritommaso.dev@gmail.com>
Date: Sat, 10 Aug 2019 15:42:54 +0200
Subject: [PATCH] Added multi wait strategy

---
 docker_test.go | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++
 wait/multi.go  | 47 ++++++++++++++++++++++++++++++
 2 files changed, 124 insertions(+)
 create mode 100644 wait/multi.go

diff --git a/docker_test.go b/docker_test.go
index a0051fedea..abffeb818e 100644
--- a/docker_test.go
+++ b/docker_test.go
@@ -592,3 +592,80 @@ func TestContainerCreationWaitsForLog(t *testing.T) {
 		t.Errorf("error creating table: %+v\n", err)
 	}
 }
+
+func TestContainerCreationWaitsForLogAndPortContextTimeout(t *testing.T) {
+	ctx := context.Background()
+	req := ContainerRequest{
+		Image:        "mysql:latest",
+		ExposedPorts: []string{"3306/tcp", "33060/tcp"},
+		Env: map[string]string{
+			"MYSQL_ROOT_PASSWORD": "password",
+			"MYSQL_DATABASE":      "database",
+		},
+		WaitingFor: wait.ForAll(
+			wait.ForLog("I love testcontainers-go"),
+			wait.ForListeningPort("3306/tcp"),
+		),
+	}
+	_, err := GenericContainer(ctx, GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+
+	if err == nil {
+		t.Fatal("Expected timeout")
+	}
+
+}
+
+func TestContainerCreationWaitsForLogAndPort(t *testing.T) {
+	ctx := context.Background()
+	req := ContainerRequest{
+		Image:        "mysql:latest",
+		ExposedPorts: []string{"3306/tcp", "33060/tcp"},
+		Env: map[string]string{
+			"MYSQL_ROOT_PASSWORD": "password",
+			"MYSQL_DATABASE":      "database",
+		},
+		WaitingFor: wait.ForAll(
+			wait.ForLog("port: 3306  MySQL Community Server - GPL"),
+			wait.ForListeningPort("3306/tcp"),
+		),
+	}
+
+	mysqlC, err := GenericContainer(ctx, GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer func() {
+		t.Log("terminating container")
+		err := mysqlC.Terminate(ctx)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	host, _ := mysqlC.Host(ctx)
+	p, _ := mysqlC.MappedPort(ctx, "3306/tcp")
+	port := p.Int()
+	connectionString := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?tls=skip-verify",
+		"root", "password", host, port, "database")
+
+	db, err := sql.Open("mysql", connectionString)
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer db.Close()
+
+	if err = db.Ping(); err != nil {
+		t.Errorf("error pinging db: %+v\n", err)
+	}
+
+}
diff --git a/wait/multi.go b/wait/multi.go
new file mode 100644
index 0000000000..468a5a1932
--- /dev/null
+++ b/wait/multi.go
@@ -0,0 +1,47 @@
+package wait
+
+import (
+	"context"
+	"fmt"
+	"time"
+)
+
+// Implement interface
+var _ Strategy = (*MultiStrategy)(nil)
+
+type MultiStrategy struct {
+	// all Strategies should have a startupTimeout to avoid waiting infinitely
+	startupTimeout time.Duration
+
+	// additional properties
+	Strategies []Strategy
+}
+
+func (ms *MultiStrategy) WithStartupTimeout(startupTimeout time.Duration) *MultiStrategy {
+	ms.startupTimeout = startupTimeout
+	return ms
+}
+
+func ForAll(strategies ...Strategy) *MultiStrategy {
+	return &MultiStrategy{
+		startupTimeout: defaultStartupTimeout(),
+		Strategies:     strategies,
+	}
+}
+
+func (ms *MultiStrategy) WaitUntilReady(ctx context.Context, target StrategyTarget) (err error) {
+	ctx, cancelContext := context.WithTimeout(ctx, ms.startupTimeout)
+	defer cancelContext()
+
+	if len(ms.Strategies) == 0 {
+		return fmt.Errorf("no wait strategy supplied")
+	}
+
+	for _, strategy := range ms.Strategies {
+		err := strategy.WaitUntilReady(ctx, target)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}