3 files changed, 33 insertions, 12 deletions

diff --git a/runsc/container/container_test.go b/runsc/container/container_test.go
index 34febe038..50f038450 100644
--- a/runsc/container/container_test.go
+++ b/runsc/container/container_test.go
@@ -122,18 +122,24 @@ func createWriteableOutputFile(path string) (*os.File, error) {
 	return outputFile, nil
 }
 
-func readOutputNum(f *os.File, first bool) (int, error) {
-	var num int
-	time.Sleep(1 * time.Second)
-
-	// Check that f exists and contains counting data.
-	fileInfo, err := f.Stat()
-	if err != nil {
-		return 0, fmt.Errorf("error creating output file: %v", err)
+func waitForFile(f *os.File) error {
+	op := func() error {
+		fi, err := f.Stat()
+		if err != nil {
+			return err
+		}
+		if fi.Size() == 0 {
+			return fmt.Errorf("file %q is empty", f.Name())
+		}
+		return nil
 	}
+	return testutil.Poll(op, 5*time.Second)
+}
 
-	if fileInfo.Size() == 0 {
-		return 0, fmt.Errorf("failed to write to file, file still appears empty")
+func readOutputNum(f *os.File, first bool) (int, error) {
+	// Wait until file has contents.
+	if err := waitForFile(f); err != nil {
+		return 0, err
 	}
 
 	// Read the first number in the new file
@@ -147,6 +153,7 @@ func readOutputNum(f *os.File, first bool) (int, error) {
 
 	nums := strings.Split(string(b), "\n")
 
+	var num int
 	if first {
 		num, err = strconv.Atoi(nums[0])
 	} else {
@@ -579,7 +586,10 @@ func TestCheckpointRestore(t *testing.T) {
 		}
 		defer file.Close()
 
-		time.Sleep(1 * time.Second)
+		// Wait until application has ran.
+		if err := waitForFile(outputFile); err != nil {
+			t.Fatalf("Failed to wait for output file: %v", err)
+		}
 
 		// Checkpoint running container; save state into new file.
 		if err := cont.Checkpoint(file); err != nil {
@@ -727,7 +737,7 @@ func TestPauseResume(t *testing.T) {
 			t.Errorf("container status got %v, want %v", got, want)
 		}
 
-		time.Sleep(10 * time.Second)
+		time.Sleep(6 * time.Second)
 
 		// Verify that the two processes still exist. Sleep 5 is paused so
 		// it should still be in the process list after 10 seconds.
diff --git a/runsc/test/testutil/BUILD b/runsc/test/testutil/BUILD
index 3ebcc1362..03ab3c4ac 100644
--- a/runsc/test/testutil/BUILD
+++ b/runsc/test/testutil/BUILD
@@ -16,6 +16,7 @@ go_library(
     deps = [
         "//runsc/boot",
         "//runsc/specutils",
+        "@com_github_cenkalti_backoff//:go_default_library",
         "@com_github_opencontainers_runtime-spec//specs-go:go_default_library",
     ],
 )
diff --git a/runsc/test/testutil/testutil.go b/runsc/test/testutil/testutil.go
index c7cef9c75..721478353 100644
--- a/runsc/test/testutil/testutil.go
+++ b/runsc/test/testutil/testutil.go
@@ -16,6 +16,7 @@
 package testutil
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -24,6 +25,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/cenkalti/backoff"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"gvisor.googlesource.com/gvisor/runsc/boot"
 	"gvisor.googlesource.com/gvisor/runsc/specutils"
@@ -172,3 +174,11 @@ func Copy(src, dst string) error {
 	_, err = io.Copy(out, in)
 	return err
 }
+
+// Poll is a shorthand function to poll for something with given timeout.
+func Poll(cb func() error, timeout time.Duration) error {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	b := backoff.WithContext(backoff.NewConstantBackOff(100*time.Millisecond), ctx)
+	return backoff.Retry(cb, b)
+}