From fed5ad4bc5f94daa03b40b2ab86eaaec0b78f2d7 Mon Sep 17 00:00:00 2001
From: Kenfe-Mickael Laventure <mickael.laventure@gmail.com>
Date: Fri, 11 Aug 2017 10:30:38 -0700
Subject: [PATCH] Update tests to properly use IsServing

Signed-off-by: Kenfe-Mickael Laventure <mickael.laventure@gmail.com>
---
 client_test.go          | 11 +++++++----
 container_linux_test.go | 11 +++--------
 daemon_test.go          | 25 ++++++++++++-------------
 3 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/client_test.go b/client_test.go
index 23e99875c..4fd72652c 100644
--- a/client_test.go
+++ b/client_test.go
@@ -11,6 +11,7 @@ import (
 	"os/exec"
 	"runtime"
 	"testing"
+	"time"
 
 	"google.golang.org/grpc/grpclog"
 
@@ -73,18 +74,20 @@ func TestMain(m *testing.M) {
 		}
 	}
 
-	client, err := ctrd.waitForStart(ctx)
+	waitCtx, waitCancel := context.WithTimeout(ctx, 2*time.Second)
+	client, err := ctrd.waitForStart(waitCtx)
+	waitCancel()
 	if err != nil {
 		ctrd.Kill()
 		ctrd.Wait()
-		fmt.Fprintf(os.Stderr, "%s: %s", err, buf.String())
+		fmt.Fprintf(os.Stderr, "%s: %s\n", err, buf.String())
 		os.Exit(1)
 	}
 
 	// print out the version in information
 	version, err := client.Version(ctx)
 	if err != nil {
-		fmt.Fprintf(os.Stderr, "error getting version: %s", err)
+		fmt.Fprintf(os.Stderr, "error getting version: %s\n", err)
 		os.Exit(1)
 	}
 
@@ -99,7 +102,7 @@ func TestMain(m *testing.M) {
 		if _, err = client.Pull(ctx, testImage, WithPullUnpack); err != nil {
 			ctrd.Stop()
 			ctrd.Wait()
-			fmt.Fprintf(os.Stderr, "%s: %s", err, buf.String())
+			fmt.Fprintf(os.Stderr, "%s: %s\n", err, buf.String())
 			os.Exit(1)
 		}
 	}
diff --git a/container_linux_test.go b/container_linux_test.go
index 2744f7e99..5d6270707 100644
--- a/container_linux_test.go
+++ b/container_linux_test.go
@@ -240,14 +240,9 @@ func TestDaemonRestart(t *testing.T) {
 
 	<-statusC
 
-	serving := false
-	for i := 0; i < 20; i++ {
-		serving, err = client.IsServing(ctx)
-		if serving {
-			break
-		}
-		time.Sleep(100 * time.Millisecond)
-	}
+	waitCtx, waitCancel := context.WithTimeout(ctx, 2*time.Second)
+	serving, err := client.IsServing(waitCtx)
+	waitCancel()
 	if !serving {
 		t.Fatalf("containerd did not start within 2s: %v", err)
 	}
diff --git a/daemon_test.go b/daemon_test.go
index de57b9fd9..577d4550e 100644
--- a/daemon_test.go
+++ b/daemon_test.go
@@ -2,7 +2,6 @@ package containerd
 
 import (
 	"context"
-	"fmt"
 	"io"
 	"os/exec"
 	"sync"
@@ -44,19 +43,19 @@ func (d *daemon) waitForStart(ctx context.Context) (*Client, error) {
 		err     error
 	)
 
-	for i := 0; i < 20; i++ {
-		if client == nil {
-			client, err = New(d.addr)
-		}
-		if err == nil {
-			serving, err = client.IsServing(ctx)
-			if serving {
-				return client, nil
-			}
-		}
-		time.Sleep(100 * time.Millisecond)
+	client, err = New(address)
+	if err != nil {
+		return nil, err
 	}
-	return nil, fmt.Errorf("containerd did not start within 2s: %v", err)
+	serving, err = client.IsServing(ctx)
+	if !serving {
+		client.Close()
+		if err == nil {
+			err = errors.New("connection was successful but service is not available")
+		}
+		return nil, err
+	}
+	return client, err
 }
 
 func (d *daemon) Stop() error {