From 6bcad9cf7177bbcea6c89fbe01ff2e9472238db4 Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Tue, 21 May 2019 15:41:48 -0700
Subject: [PATCH] Support image family in node e2e image config file.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 test/e2e_node/runner/remote/run_remote.go | 51 ++++++++++++++---------
 1 file changed, 31 insertions(+), 20 deletions(-)

diff --git a/test/e2e_node/runner/remote/run_remote.go b/test/e2e_node/runner/remote/run_remote.go
index 407a31eaec..479641e37a 100644
--- a/test/e2e_node/runner/remote/run_remote.go
+++ b/test/e2e_node/runner/remote/run_remote.go
@@ -21,6 +21,7 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"io/ioutil"
@@ -152,6 +153,8 @@ type GCEImage struct {
 	// Defaults to using only the latest image. Acceptable values are [0, # of images that match the regex).
 	// If the number of existing previous images is lesser than what is desired, the test will use that is available.
 	PreviousImages int `json:"previous_images,omitempty"`
+	// ImageFamily is the image family to use. The latest image from the image family will be used.
+	ImageFamily string `json:"image_family,omitempty"`
 
 	Machine   string    `json:"machine,omitempty"`
 	Resources Resources `json:"resources,omitempty"`
@@ -229,11 +232,12 @@ func main() {
 		for shortName, imageConfig := range externalImageConfig.Images {
 			var images []string
 			isRegex, name := false, shortName
-			if imageConfig.ImageRegex != "" && imageConfig.Image == "" {
+			if (imageConfig.ImageRegex != "" || imageConfig.ImageFamily != "") && imageConfig.Image == "" {
 				isRegex = true
-				images, err = getGCEImages(imageConfig.ImageRegex, imageConfig.Project, imageConfig.PreviousImages)
+				images, err = getGCEImages(imageConfig.ImageRegex, imageConfig.ImageFamily, imageConfig.Project, imageConfig.PreviousImages)
 				if err != nil {
-					klog.Fatalf("Could not retrieve list of images based on image prefix %q: %v", imageConfig.ImageRegex, err)
+					klog.Fatalf("Could not retrieve list of images based on image prefix %q and family %q: %v",
+						imageConfig.ImageRegex, imageConfig.ImageFamily, err)
 				}
 			} else {
 				images = []string{imageConfig.Image}
@@ -468,26 +472,33 @@ func (a byCreationTime) Less(i, j int) bool { return a[i].creationTime.After(a[j
 func (a byCreationTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
 
 // Returns a list of image names based on regex and number of previous images requested.
-func getGCEImages(imageRegex, project string, previousImages int) ([]string, error) {
-	ilc, err := computeService.Images.List(project).Do()
-	if err != nil {
-		return nil, fmt.Errorf("Failed to list images in project %q: %v", project, err)
-	}
+func getGCEImages(imageRegex, imageFamily string, project string, previousImages int) ([]string, error) {
 	imageObjs := []imageObj{}
 	imageRe := regexp.MustCompile(imageRegex)
-	for _, instance := range ilc.Items {
-		if imageRe.MatchString(instance.Name) {
-			creationTime, err := time.Parse(time.RFC3339, instance.CreationTimestamp)
-			if err != nil {
-				return nil, fmt.Errorf("Failed to parse instance creation timestamp %q: %v", instance.CreationTimestamp, err)
+	if err := computeService.Images.List(project).Pages(context.Background(),
+		func(ilc *compute.ImageList) error {
+			for _, instance := range ilc.Items {
+				if imageRegex != "" && !imageRe.MatchString(instance.Name) {
+					continue
+				}
+				if imageFamily != "" && instance.Family != imageFamily {
+					continue
+				}
+				creationTime, err := time.Parse(time.RFC3339, instance.CreationTimestamp)
+				if err != nil {
+					return fmt.Errorf("failed to parse instance creation timestamp %q: %v", instance.CreationTimestamp, err)
+				}
+				io := imageObj{
+					creationTime: creationTime,
+					name:         instance.Name,
+				}
+				klog.V(4).Infof("Found image %q based on regex %q and family %q in project %q", io.string(), imageRegex, imageFamily, project)
+				imageObjs = append(imageObjs, io)
 			}
-			io := imageObj{
-				creationTime: creationTime,
-				name:         instance.Name,
-			}
-			klog.V(4).Infof("Found image %q based on regex %q in project %q", io.string(), imageRegex, project)
-			imageObjs = append(imageObjs, io)
-		}
+			return nil
+		},
+	); err != nil {
+		return nil, fmt.Errorf("failed to list images in project %q: %v", project, err)
 	}
 	sort.Sort(byCreationTime(imageObjs))
 	images := []string{}