From 2d0c45b37fe48299be53abc19986c86e9b2f170c Mon Sep 17 00:00:00 2001
From: Dimitri Koshkin <dimitri.koshkin@nutanix.com>
Date: Thu, 4 Jan 2024 18:27:29 -0800
Subject: [PATCH] fix: revert to previous behaviod polling forever

The ctx passed into WaitForTaskToSucceed is only used to cancel HTTP reqests and not to cancel the wait.
---
 pkg/client/status.go      |  5 +++--
 pkg/client/status_test.go | 11 +++++++----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/pkg/client/status.go b/pkg/client/status.go
index a56ca4069f..1bdc58e21c 100644
--- a/pkg/client/status.go
+++ b/pkg/client/status.go
@@ -33,9 +33,10 @@ const (
 )
 
 // WaitForTaskToSucceed will poll indefinitely every 2 seconds for the task with uuid to have status of "SUCCEEDED".
-// Returns an error from GetTaskStatus or a timeout error if the context is cancelled.
+// The polling will not stop if the ctx is cancelled, it's only used for HTTP requests in the client.
+// WaitForTaskToSucceed will exit immediately on an error getting the task.
 func WaitForTaskToSucceed(ctx context.Context, conn *nutanixClientV3.Client, uuid string) error {
-	return wait.PollImmediateInfiniteWithContext(ctx, pollingInterval, func(ctx context.Context) (done bool, err error) {
+	return wait.PollImmediateInfinite(pollingInterval, func() (done bool, err error) {
 		status, getErr := GetTaskStatus(ctx, conn, uuid)
 		return status == statusSucceeded, getErr
 	})
diff --git a/pkg/client/status_test.go b/pkg/client/status_test.go
index 63dc5c8491..6e2a8a0aab 100644
--- a/pkg/client/status_test.go
+++ b/pkg/client/status_test.go
@@ -8,7 +8,6 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"k8s.io/apimachinery/pkg/util/wait"
 
 	nutanixtestclient "github.com/nutanix-cloud-native/cluster-api-provider-nutanix/test/helpers/prism-go-client/v3"
 )
@@ -90,7 +89,10 @@ func Test_WaitForTaskCompletion(t *testing.T) {
 		client.Close()
 	})
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
+	const (
+		timeout = time.Second * 1
+	)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	t.Cleanup(func() {
 		cancel()
 	})
@@ -124,10 +126,11 @@ func Test_WaitForTaskCompletion(t *testing.T) {
 			name:     "timeout",
 			taskUUID: "timeout",
 			handler: func(w http.ResponseWriter, r *http.Request) {
-				fmt.Fprint(w, `{"status": "PENDING"}`)
+				// always wait 1 second longer than the timeout to force the context to cancel
+				time.Sleep(timeout + time.Second)
 			},
 			ctx:         ctx,
-			expectedErr: wait.ErrWaitTimeout,
+			expectedErr: context.DeadlineExceeded,
 		},
 	}
 	for _, tt := range tests {