Refactor test context creation to be less error-prone

Change-Id: I176faf21ff3f2cac5e9d8af0e54219761693d361

Author: mattcary

test/e2e/tests/setup_e2e_test.go

@@ -65,8 +65,8 @@ var (
 	// Some architectures don't have local ssd. Give way to opt out of tests like datacache.
 	skipLocalSsdTests = flag.Bool("skip-local-ssd-tests", false, "Skip local ssd tests like datacache")
 
-	testContexts          = []*remote.TestContext{}
-	hyperdiskTestContexts = []*remote.TestContext{}
+	testContexts          []*remote.TestContext
+	hyperdiskTestContexts []*remote.TestContext
 	computeService        *compute.Service
 	computeAlphaService   *computealpha.Service
 	computeBetaService    *computebeta.Service
@@ -87,10 +87,6 @@ var _ = BeforeSuite(func() {
 	var err error
 	numberOfInstancesPerZone := 2
 	zones := strings.Split(*zones, ",")
-	tcc := make(chan *remote.TestContext, len(zones)*numberOfInstancesPerZone)
-	hdtcc := make(chan *remote.TestContext, len(zones))
-	defer close(tcc)
-	defer close(hdtcc)
 
 	rand.Seed(time.Now().UnixNano())
 
@@ -116,42 +112,43 @@ var _ = BeforeSuite(func() {
 
 	klog.Infof("Running in project %v with service account %v", *project, *serviceAccount)
 
-	setupContext := func(zone string) {
+	testContexts = make([]*remote.TestContext, numberOfInstancesPerZone*len(zones))
+	if *hdMinCpuPlatform != noMachineType {
+		hyperdiskTestContexts = make([]*remote.TestContext, len(zones))
+	}
+	var zoneGroup sync.WaitGroup
+	setupContext := func(idx int, zone string) {
+		defer zoneGroup.Done()
 		var wg sync.WaitGroup
 		// Create 2 instances for each zone as we need 2 instances each zone for certain test cases
 		for j := 0; j < numberOfInstancesPerZone; j++ {
 			wg.Add(1)
 			go func(curZone string, randInt int) {
 				defer GinkgoRecover()
 				defer wg.Done()
-				tcc <- NewDefaultTestContext(curZone, strconv.Itoa(randInt))
+				tc := NewDefaultTestContext(curZone, strconv.Itoa(randInt))
+				testContexts[j+idx*numberOfInstancesPerZone] = tc
+				klog.Infof("Added TestContext for node %s", tc.Instance.GetName())
 			}(zone, j)
 		}
 		if *hdMachineType != noMachineType {
 			wg.Add(1)
 			go func(curZone string) {
 				defer GinkgoRecover()
 				defer wg.Done()
-				hdtcc <- NewTestContext(curZone, *hdMinCpuPlatform, *hdMachineType, "0")
+				tc := NewTestContext(curZone, *hdMinCpuPlatform, *hdMachineType, "0")
+				hyperdiskTestContexts[idx] = tc
+				klog.Infof("Added hyperdisk TestContext for node %s", tc.Instance.GetName())
 			}(zone)
 		}
 		wg.Wait()
 	}
 
-	for _, zone := range zones {
-		setupContext(zone)
-	}
-
-	for i := 0; i < len(zones)*numberOfInstancesPerZone; i++ {
-		tc := <-tcc
-		testContexts = append(testContexts, tc)
-		klog.Infof("Added TestContext for node %s", tc.Instance.GetName())
-	}
-	for i := 0; i < len(zones); i++ {
-		tc := <-hdtcc
-		hyperdiskTestContexts = append(hyperdiskTestContexts, tc)
-		klog.Infof("Added TestContext for node %s", tc.Instance.GetName())
+	for i, zone := range zones {
+		zoneGroup.Add(1)
+		go setupContext(i, zone)
 	}
+	zoneGroup.Wait()
 })
 
 var _ = AfterSuite(func() {