diff --git a/cli/command/container/client_test.go b/cli/command/container/client_test.go
index 692ed3115a..8fa5abf0e8 100644
--- a/cli/command/container/client_test.go
+++ b/cli/command/container/client_test.go
@@ -21,6 +21,8 @@ type fakeClient struct {
 	containerStatPathFunc func(container, path string) (types.ContainerPathStat, error)
 	containerCopyFromFunc func(container, srcPath string) (io.ReadCloser, types.ContainerPathStat, error)
 	logFunc               func(string, types.ContainerLogsOptions) (io.ReadCloser, error)
+	waitFunc              func(string) (<-chan container.ContainerWaitOKBody, <-chan error)
+	Version               string
 }
 
 func (f *fakeClient) ContainerInspect(_ context.Context, containerID string) (types.ContainerJSON, error) {
@@ -95,3 +97,14 @@ func (f *fakeClient) ContainerLogs(_ context.Context, container string, options
 	}
 	return nil, nil
 }
+
+func (f *fakeClient) ClientVersion() string {
+	return f.Version
+}
+
+func (f *fakeClient) ContainerWait(_ context.Context, container string, _ container.WaitCondition) (<-chan container.ContainerWaitOKBody, <-chan error) {
+	if f.waitFunc != nil {
+		return f.waitFunc(container)
+	}
+	return nil, nil
+}
diff --git a/cli/command/container/utils_test.go b/cli/command/container/utils_test.go
new file mode 100644
index 0000000000..13c43ff074
--- /dev/null
+++ b/cli/command/container/utils_test.go
@@ -0,0 +1,69 @@
+package container
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/docker/cli/internal/test"
+	"github.com/docker/docker/api"
+	"github.com/docker/docker/api/types/container"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/net/context"
+)
+
+func waitFn(cid string) (<-chan container.ContainerWaitOKBody, <-chan error) {
+	resC := make(chan container.ContainerWaitOKBody)
+	errC := make(chan error, 1)
+	var res container.ContainerWaitOKBody
+
+	go func() {
+		switch {
+		case strings.Contains(cid, "exit-code-42"):
+			res.StatusCode = 42
+			resC <- res
+		case strings.Contains(cid, "non-existent"):
+			err := errors.Errorf("No such container: %v", cid)
+			errC <- err
+		case strings.Contains(cid, "wait-error"):
+			res.Error = &container.ContainerWaitOKBodyError{Message: "removal failed"}
+			resC <- res
+		default:
+			// normal exit
+			resC <- res
+		}
+	}()
+
+	return resC, errC
+}
+
+func TestWaitExitOrRemoved(t *testing.T) {
+	testcases := []struct {
+		cid      string
+		exitCode int
+	}{
+		{
+			cid:      "normal-container",
+			exitCode: 0,
+		},
+		{
+			cid:      "give-me-exit-code-42",
+			exitCode: 42,
+		},
+		{
+			cid:      "i-want-a-wait-error",
+			exitCode: 125,
+		},
+		{
+			cid:      "non-existent-container-id",
+			exitCode: 125,
+		},
+	}
+
+	client := test.NewFakeCli(&fakeClient{waitFunc: waitFn, Version: api.DefaultVersion})
+	for _, testcase := range testcases {
+		statusC := waitExitOrRemoved(context.Background(), client, testcase.cid, true)
+		exitCode := <-statusC
+		assert.Equal(t, testcase.exitCode, exitCode)
+	}
+}