address PR comments

test/framework/resource/k8s/node/wrapper.go

@@ -46,9 +46,9 @@ func GetNodeAndWaitTillCapacityPresent(manager Manager, os string, expectedResou
 	return observedNodeList
 }
 
-// VerifyCNINodeCount checks if the number of CNINodes is equal to number of nodes in the cluster, and verifies 1:1 mapping between CNINode and Node objects
+// VerifyCNINode checks if the number of CNINodes is equal to number of nodes in the cluster, and verifies 1:1 mapping between CNINode and Node objects
 // Returns nil if count and 1:1 mapping exists, else returns error
-func VerifyCNINodeCount(manager Manager) error {
+func VerifyCNINode(manager Manager) error {
 	cniNodes, err := manager.GetCNINodeList()
 	Expect(err).NotTo(HaveOccurred())
 	nodes, err := manager.GetNodeList()

test/integration/cninode/cninode_suite_test.go

@@ -38,13 +38,13 @@ var _ = BeforeSuite(func() {
 	Expect(len(nodeList.Items)).To(BeNumerically(">", 1))
 
 	By("verify CNINode count")
-	err = node.VerifyCNINodeCount(frameWork.NodeManager)
+	err = node.VerifyCNINode(frameWork.NodeManager)
 	Expect(err).ToNot(HaveOccurred())
 })
 
 // Verify CNINode count before and after test remains same
 var _ = AfterSuite(func() {
 	By("verify CNINode count")
-	err := node.VerifyCNINodeCount(frameWork.NodeManager)
+	err := node.VerifyCNINode(frameWork.NodeManager)
 	Expect(err).ToNot(HaveOccurred())
 })

test/integration/cninode/cninode_test.go

@@ -14,14 +14,15 @@
 package cninode_test
 
 import (
-	"time"
+	"context"
 
 	"github.com/aws/amazon-vpc-resource-controller-k8s/pkg/config"
 	"github.com/aws/amazon-vpc-resource-controller-k8s/pkg/utils"
 	"github.com/aws/amazon-vpc-resource-controller-k8s/test/framework/resource/k8s/node"
 	testUtils "github.com/aws/amazon-vpc-resource-controller-k8s/test/framework/utils"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"k8s.io/apimachinery/pkg/util/wait"
 )
 
 var _ = Describe("[CANARY]CNINode test", func() {
@@ -44,8 +45,7 @@ var _ = Describe("[CANARY]CNINode test", func() {
 			By("restoring ASG minSize & maxSize after test")
 			err := frameWork.AutoScalingManager.UpdateAutoScalingGroup(asgName, oldDesiredSize, oldMinSize, oldMaxSize)
 			Expect(err).ToNot(HaveOccurred())
-			// sleep to allow ASG to be updated
-			time.Sleep(testUtils.ResourceCreationTimeout)
+			Expect(WaitTillNodeSizeUpdated(int(oldDesiredSize))).Should(Succeed())
 		})
 
 		Context("when new node is added", func() {
@@ -55,11 +55,8 @@ var _ = Describe("[CANARY]CNINode test", func() {
 				By("adding new node")
 				err := frameWork.AutoScalingManager.UpdateAutoScalingGroup(asgName, newSize, oldMinSize, newSize)
 				Expect(err).ToNot(HaveOccurred())
-				// sleep to allow ASG to be updated
-				time.Sleep(testUtils.ResourceCreationTimeout)
-
-				err = node.VerifyCNINodeCount(frameWork.NodeManager)
-				Expect(err).ToNot(HaveOccurred())
+				Expect(WaitTillNodeSizeUpdated(int(newSize))).Should(Succeed())
+				Expect(node.VerifyCNINode(frameWork.NodeManager)).Should(Succeed())
 			})
 		})
 		Context("when existing node is removed", func() {
@@ -69,11 +66,8 @@ var _ = Describe("[CANARY]CNINode test", func() {
 				By("removing existing node")
 				err := frameWork.AutoScalingManager.UpdateAutoScalingGroup(asgName, newSize, newSize, oldMaxSize)
 				Expect(err).ToNot(HaveOccurred())
-				// sleep to allow ASG to be updated
-				time.Sleep(testUtils.ResourceCreationTimeout)
-				err = node.VerifyCNINodeCount(frameWork.NodeManager)
-				Expect(err).ToNot(HaveOccurred())
-
+				Expect(WaitTillNodeSizeUpdated(int(newSize))).Should(Succeed())
+				Expect(node.VerifyCNINode(frameWork.NodeManager)).Should(Succeed())
 			})
 		})
 	})
@@ -93,3 +87,19 @@ func ListNodesAndGetAutoScalingGroupName() string {
 	Expect(ok).To(BeTrue())
 	return val
 }
+
+// Verifies (linux) node size is updated after ASG is updated
+func WaitTillNodeSizeUpdated(desiredSize int) error {
+	err := wait.PollUntilContextTimeout(context.Background(), testUtils.PollIntervalShort, testUtils.ResourceCreationTimeout, true,
+		func(ctx context.Context) (bool, error) {
+			nodes, err := frameWork.NodeManager.GetNodesWithOS(config.OSLinux) // since we are only updating the linux ASG in the test
+			if err != nil {
+				return false, nil
+			}
+			if len(nodes.Items) != desiredSize {
+				return false, nil
+			}
+			return true, nil
+		})
+	return err
+}

test/integration/perpodsg/perpodsg_suite_test.go

@@ -54,7 +54,7 @@ var _ = BeforeSuite(func() {
 
 	nodeList = node.GetNodeAndWaitTillCapacityPresent(frameWork.NodeManager, "linux",
 		config.ResourceNamePodENI)
-	err = node.VerifyCNINodeCount(frameWork.NodeManager)
+	err = node.VerifyCNINode(frameWork.NodeManager)
 	Expect(err).ToNot(HaveOccurred())
 })