From 8ba2f2f51694910d9b6561a58f1702a2d184c1f0 Mon Sep 17 00:00:00 2001
From: graysonwu <wgrayson@vmware.com>
Date: Tue, 6 Feb 2024 07:13:13 -0800
Subject: [PATCH] Use waitForState

Signed-off-by: graysonwu <wgrayson@vmware.com>
---
 nsxt/data_source_nsxt_edge_upgrade_group.go | 41 ++++++++++-----------
 1 file changed, 19 insertions(+), 22 deletions(-)

diff --git a/nsxt/data_source_nsxt_edge_upgrade_group.go b/nsxt/data_source_nsxt_edge_upgrade_group.go
index 06cc045b0..ebe2ccff6 100644
--- a/nsxt/data_source_nsxt_edge_upgrade_group.go
+++ b/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
 }