Use waitForState

Signed-off-by: graysonwu <[email protected]>

nsxt/data_source_nsxt_edge_upgrade_group.go

@@ -5,6 +5,7 @@ package nsxt
 
 import (
 	"fmt"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
 	"strings"
 	"time"
 
@@ -17,8 +18,8 @@ var (
 	edgeUpgradeGroup = "EDGE"
 	hostUpgradeGroup = "HOST"
 
-	timeoutUUGCreat  = 1 * time.Minute
-	intervalUUGCreat = 5 * time.Second
+	timeoutUUGCreate  = 1 * time.Minute
+	intervalUUGCreate = 5 * time.Second
 )
 
 func dataSourceNsxtEdgeUpgradeGroup() *schema.Resource {
@@ -113,35 +114,31 @@ func upgradeGroupRead(d *schema.ResourceData, m interface{}, groupType string) e
 }
 
 func waitUpgradeGroupCreate(client upgrade.UpgradeUnitGroupsClient, groupType string) error {
-
-	resultChan := make(chan error, 1)
-
-	poll := func(resultChan chan error) {
-		preLen := 0
-		for {
+	preLen := 0
+	stateConf := &resource.StateChangeConf{
+		Pending: []string{"Pending"},
+		Target:  []string{"Done"},
+		Refresh: func() (interface{}, string, error) {
 			objList, err := client.List(&groupType, nil, nil, nil, nil, nil, nil, nil)
 			if err != nil && !isNotFoundError(err) {
-				resultChan <- err
-				return
+				return nil, "", err
 			}
 			if err == nil && len(objList.Results) > 0 {
 				if preLen == 0 {
 					preLen = len(objList.Results)
 				} else if preLen == len(objList.Results) {
-					resultChan <- nil
-					return
+					return "Done", "Done", nil
 				}
 			}
-			time.Sleep(intervalUUGCreat)
-		}
+			return "Pending", "Pending", nil
+		},
+		Timeout:      timeoutUUGCreate,
+		PollInterval: intervalUUGCreate,
+		Delay:        intervalUUGCreate,
 	}
-
-	go poll(resultChan)
-
-	select {
-	case res := <-resultChan:
-		return res
-	case <-time.After(timeoutUUGCreat):
-		return fmt.Errorf("timeout reached")
+	_, err := stateConf.WaitForState()
+	if err != nil {
+		return fmt.Errorf("failed to wait UpgradeGroups to be created: %v", err)
 	}
+	return nil
 }