From ce8f82bb5a77f3e3013c005192a4bfc3a6dd9f39 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Tue, 7 May 2019 16:07:10 +0800
Subject: [PATCH] support timeout for NewLogStrategy

Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
---
 docker_test.go | 21 +++++++++++++++++++++
 wait/log.go    | 31 +++++++++++++++++++------------
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/docker_test.go b/docker_test.go
index ea9f508b5b..d7854644dc 100644
--- a/docker_test.go
+++ b/docker_test.go
@@ -328,6 +328,27 @@ func TestContainerCreationTimesOutWithHttp(t *testing.T) {
 	}
 }
 
+func TestContainerCreationWaitsForLogContextTimeout(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.ForLog("test context timeout").WithStartupTimeout(1 * time.Second),
+	}
+	_, err := GenericContainer(ctx, GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+
+	if err == nil {
+		t.Error("Expected timeout")
+	}
+}
+
 func TestContainerCreationWaitsForLog(t *testing.T) {
 	ctx := context.Background()
 	req := ContainerRequest{
diff --git a/wait/log.go b/wait/log.go
index 2409f32dca..06fe2de97c 100644
--- a/wait/log.go
+++ b/wait/log.go
@@ -62,19 +62,26 @@ func (ws *LogStrategy) WaitUntilReady(ctx context.Context, target StrategyTarget
 	ctx, cancelContext := context.WithTimeout(ctx, ws.startupTimeout)
 	defer cancelContext()
 
+LOOP:
 	for {
-		reader, err := target.Logs(ctx)
-		if err != nil {
-			time.Sleep(ws.PollInterval)
-			continue
-		}
-		b, err := ioutil.ReadAll(reader)
-		logs := string(b)
-		if strings.Contains(logs, ws.Log) {
-			break
-		} else {
-			time.Sleep(ws.PollInterval)
-			continue
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			reader, err := target.Logs(ctx)
+
+			if err != nil {
+				time.Sleep(ws.PollInterval)
+				continue
+			}
+			b, err := ioutil.ReadAll(reader)
+			logs := string(b)
+			if strings.Contains(logs, ws.Log) {
+				break LOOP
+			} else {
+				time.Sleep(ws.PollInterval)
+				continue
+			}
 		}
 	}