From 17a09ddee4dc8fa6d6d4ff95e280c797a9b024ee Mon Sep 17 00:00:00 2001
From: Matthew Cary <mattcary@google.com>
Date: Wed, 7 Oct 2020 11:42:56 -0700
Subject: [PATCH] Wait for windows prepull images rather than sleeping

---
 test/k8s-integration/main.go          |  4 +--
 test/k8s-integration/prepull-image.sh | 38 ++++++++++++++++++++-------
 2 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/test/k8s-integration/main.go b/test/k8s-integration/main.go
index 9d51016a5..75a0f6c97 100644
--- a/test/k8s-integration/main.go
+++ b/test/k8s-integration/main.go
@@ -22,7 +22,6 @@ import (
 	"path/filepath"
 	"strings"
 	"syscall"
-	"time"
 
 	"k8s.io/apimachinery/pkg/util/uuid"
 	apimachineryversion "k8s.io/apimachinery/pkg/util/version"
@@ -327,7 +326,7 @@ func handle() error {
 		// It typically takes 5+ minutes to download Windows container image. To avoid tests being timed out,
 		// pre-pulling the test images as best effort.
 		klog.Infof("Prepulling test images.")
-		err = os.Setenv("PREPULL_IMAGE", filepath.Join(pkgDir, "test", "k8s-integration", "prepull.yaml"))
+		err = os.Setenv("PREPULL_YAML", filepath.Join(pkgDir, "test", "k8s-integration", "prepull.yaml"))
 		if err != nil {
 			return err
 		}
@@ -335,7 +334,6 @@ func handle() error {
 		if err != nil {
 			return fmt.Errorf("failed to prepull images: %s, err: %v", out, err)
 		}
-		time.Sleep(10 * time.Minute)
 		out, err = exec.Command("kubectl", "describe", "pods", "-n", getDriverNamespace()).CombinedOutput()
 		klog.Infof("describe pods \n %s", string(out))
 
diff --git a/test/k8s-integration/prepull-image.sh b/test/k8s-integration/prepull-image.sh
index 6fc047d0f..810d6010e 100755
--- a/test/k8s-integration/prepull-image.sh
+++ b/test/k8s-integration/prepull-image.sh
@@ -4,22 +4,40 @@ set -o nounset
 set -o pipefail
 set -o xtrace
 
-if [[ -z "${PREPULL_IMAGE}" ]]; then
+# This is taken from prepull.yaml.
+readonly prepull_daemonset=prepull-test-containers
+
+wait_on_prepull()
+{
+    # Wait up to 15 minutes for the test images to be pulled onto the nodes.
+    retries=90
+    while [[ $retries -ge 0 ]];do
+        ready=$(kubectl get daemonset "${prepull_daemonset}" -o jsonpath="{.status.numberReady}")
+        required=$(kubectl get daemonset "${prepull_daemonset}" -o jsonpath="{.status.desiredNumberScheduled}")
+        if [[ $ready -eq $required ]];then
+            echo "Daemonset $prepull_daemonset ready"
+            return 0
+        fi
+        ((retries--))
+        sleep 10s
+    done
+    echo "Timeout waiting for daemonset $prepull_daemonset"
+    return -1
+
+}
+
+if [[ -z "${PREPULL_YAML}" ]]; then
     # Pre-pull all the test images. The images are currently hard-coded.
     # Eventually, we should get the list directly from
     # https://github.com/kubernetes-sigs/windows-testing/blob/master/images/PullImages.ps1
     curl https://raw.githubusercontent.com/kubernetes-sigs/windows-testing/master/gce/prepull-1.18.yaml -o prepull.yaml
-    PREPULL_IMAGE=prepull.yaml
-    echo ${PREPULL_IMAGE}
+    PREPULL_YAML=prepull.yaml
+    echo ${PREPULL_YAML}
 fi
 
-kubectl create -f ${PREPULL_IMAGE}
-# Wait 10 minutes for the test images to be pulled onto the nodes.
-sleep 15m
-echo "sleep 15m"
+kubectl create -f ${PREPULL_YAML}
+wait_on_prepull || exit -1  # Error already printed
 # Check the status of the pods.
 kubectl get pods -o wide
 # Delete the pods anyway since pre-pulling is best-effort
-kubectl delete -f ${PREPULL_IMAGE}
-# Wait a few more minutes for the pod to be cleaned up.
-sleep 5m
\ No newline at end of file
+kubectl delete -f ${PREPULL_YAML} --wait=true