From 677d17150a06dde5f50a7eb1bdb8612f230e526c Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@docker.com>
Date: Fri, 1 Sep 2017 15:30:33 -0400
Subject: [PATCH] Move common e2e things into an internal package.

Signed-off-by: Daniel Nephin <dnephin@docker.com>
---
 e2e/stack/main_test.go               | 13 ++------
 e2e/stack/remove_test.go             | 48 +++++++---------------------
 internal/test/environment/testenv.go | 21 ++++++++++++
 3 files changed, 35 insertions(+), 47 deletions(-)
 create mode 100644 internal/test/environment/testenv.go

diff --git a/e2e/stack/main_test.go b/e2e/stack/main_test.go
index 74081f457b..0d36e45bd4 100644
--- a/e2e/stack/main_test.go
+++ b/e2e/stack/main_test.go
@@ -5,22 +5,13 @@ import (
 	"os"
 	"testing"
 
-	"github.com/pkg/errors"
+	"github.com/docker/cli/internal/test/environment"
 )
 
 func TestMain(m *testing.M) {
-	if err := setupTestEnv(); err != nil {
+	if err := environment.Setup(); err != nil {
 		fmt.Println(err.Error())
 		os.Exit(3)
 	}
 	os.Exit(m.Run())
 }
-
-// TODO: move to shared internal package
-func setupTestEnv() error {
-	dockerHost := os.Getenv("TEST_DOCKER_HOST")
-	if dockerHost == "" {
-		return errors.New("$TEST_DOCKER_HOST must be set")
-	}
-	return os.Setenv("DOCKER_HOST", dockerHost)
-}
diff --git a/e2e/stack/remove_test.go b/e2e/stack/remove_test.go
index 2d39d02345..f579340793 100644
--- a/e2e/stack/remove_test.go
+++ b/e2e/stack/remove_test.go
@@ -4,14 +4,17 @@ import (
 	"fmt"
 	"strings"
 	"testing"
-	"time"
 
+	"github.com/docker/cli/internal/test/environment"
 	shlex "github.com/flynn-archive/go-shlex"
 	"github.com/gotestyourself/gotestyourself/golden"
 	"github.com/gotestyourself/gotestyourself/icmd"
+	"github.com/gotestyourself/gotestyourself/poll"
 	"github.com/stretchr/testify/require"
 )
 
+var pollSettings = environment.DefaultPollSettings
+
 func TestRemove(t *testing.T) {
 	stackname := "test-stack-remove"
 	deployFullStack(t, stackname)
@@ -29,21 +32,24 @@ func deployFullStack(t *testing.T, stackname string) {
 		"docker stack deploy --compose-file=./testdata/full-stack.yml %s", stackname))
 	result.Assert(t, icmd.Success)
 
-	waitOn(t, taskCount(stackname, 2), 0)
+	poll.WaitOn(t, taskCount(stackname, 2), pollSettings)
 }
 
 func cleanupFullStack(t *testing.T, stackname string) {
 	result := icmd.RunCmd(shell(t, "docker stack rm %s", stackname))
 	result.Assert(t, icmd.Success)
-	waitOn(t, taskCount(stackname, 0), 0)
+	poll.WaitOn(t, taskCount(stackname, 0), pollSettings)
 }
 
-func taskCount(stackname string, expected int) func() (bool, error) {
-	return func() (bool, error) {
+func taskCount(stackname string, expected int) func(t poll.LogT) poll.Result {
+	return func(poll.LogT) poll.Result {
 		result := icmd.RunCommand(
 			"docker", "stack", "ps", "-f=desired-state=running", stackname)
 		count := lines(result.Stdout()) - 1
-		return count == expected, nil
+		if count == expected {
+			return poll.Success()
+		}
+		return poll.Continue("task count is %d waiting for %d", count, expected)
 	}
 }
 
@@ -57,33 +63,3 @@ func shell(t *testing.T, format string, args ...interface{}) icmd.Cmd {
 	require.NoError(t, err)
 	return icmd.Cmd{Command: cmd}
 }
-
-// TODO: move to gotestyourself
-func waitOn(t *testing.T, check func() (bool, error), timeout time.Duration) {
-	if timeout == time.Duration(0) {
-		timeout = defaultTimeout()
-	}
-
-	after := time.After(timeout)
-	for {
-		select {
-		case <-after:
-			// TODO: include check function name in error message
-			t.Fatalf("timeout hit after %s", timeout)
-		default:
-			// TODO: maybe return a failure message as well?
-			done, err := check()
-			if done {
-				return
-			}
-			if err != nil {
-				t.Fatal(err.Error())
-			}
-		}
-	}
-}
-
-func defaultTimeout() time.Duration {
-	// TODO: support override from environment variable
-	return 10 * time.Second
-}
diff --git a/internal/test/environment/testenv.go b/internal/test/environment/testenv.go
new file mode 100644
index 0000000000..d30f22537e
--- /dev/null
+++ b/internal/test/environment/testenv.go
@@ -0,0 +1,21 @@
+package environment
+
+import (
+	"os"
+	"time"
+
+	"github.com/gotestyourself/gotestyourself/poll"
+	"github.com/pkg/errors"
+)
+
+// Setup a new environment
+func Setup() error {
+	dockerHost := os.Getenv("TEST_DOCKER_HOST")
+	if dockerHost == "" {
+		return errors.New("$TEST_DOCKER_HOST must be set")
+	}
+	return os.Setenv("DOCKER_HOST", dockerHost)
+}
+
+// DefaultPollSettings used with gotestyourself/poll
+var DefaultPollSettings = poll.WithDelay(100 * time.Millisecond)