Update tests to properly use IsServing

Signed-off-by: Kenfe-Mickael Laventure <mickael.laventure@gmail.com>

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)
 		}
 	}

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)
 	}

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)
+	client, err = New(address)
+	if err != nil {
+		return nil, err
 	}
-		if err == nil {
 	serving, err = client.IsServing(ctx)
-			if serving {
-				return client, nil
+	if !serving {
+		client.Close()
+		if err == nil {
+			err = errors.New("connection was successful but service is not available")
 		}
+		return nil, err
 	}
-		time.Sleep(100 * time.Millisecond)
-	}
-	return nil, fmt.Errorf("containerd did not start within 2s: %v", err)
+	return client, err
 }
 
 func (d *daemon) Stop() error {