From f5aa6bad2515665f21597a449f618f17062d0fe6 Mon Sep 17 00:00:00 2001
From: Andrew Lavery <laverya@umich.edu>
Date: Tue, 2 May 2023 12:19:48 -0400
Subject: [PATCH] use wait.PollUntilContextTimeout instead of wait.Poll

Signed-off-by: Andrew Lavery <laverya@umich.edu>
---
 go.mod               | 2 ++
 pkg/client/delete.go | 4 ++--
 pkg/client/run.go    | 5 +++--
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/go.mod b/go.mod
index 18ee6b708..3253e9192 100644
--- a/go.mod
+++ b/go.mod
@@ -60,10 +60,12 @@ require (
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/pelletier/go-toml v1.9.5 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.6 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
 	github.com/spf13/afero v1.9.3 // indirect
 	github.com/spf13/cast v1.5.0 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
+	github.com/stretchr/testify v1.8.1 // indirect
 	github.com/subosito/gotenv v1.4.1 // indirect
 	golang.org/x/net v0.8.0 // indirect
 	golang.org/x/oauth2 v0.3.0 // indirect
diff --git a/pkg/client/delete.go b/pkg/client/delete.go
index ec404e9b7..a5bd46759 100644
--- a/pkg/client/delete.go
+++ b/pkg/client/delete.go
@@ -85,7 +85,7 @@ func (c *SonobuoyClient) Delete(cfg *DeleteConfig) error {
 
 	if cfg.Wait > time.Duration(0) {
 		lastProgress := ""
-		allConditions := func() (bool, error) {
+		allConditions := func(ctx context.Context) (bool, error) {
 			for _, condition := range conditions {
 				progress, done, err := condition()
 				if cfg.WaitOutput == progressMode {
@@ -111,7 +111,7 @@ func (c *SonobuoyClient) Delete(cfg *DeleteConfig) error {
 			s.Start()
 			defer s.Stop()
 		}
-		if err := wait.Poll(pollFreq, cfg.Wait, allConditions); err != nil {
+		if err := wait.PollUntilContextTimeout(context.TODO(), pollFreq, cfg.Wait, false, allConditions); err != nil {
 			return errors.Wrap(err, "waiting for delete conditions to be met")
 		}
 	}
diff --git a/pkg/client/run.go b/pkg/client/run.go
index 392fd8e35..e77bbec0e 100644
--- a/pkg/client/run.go
+++ b/pkg/client/run.go
@@ -18,6 +18,7 @@ package client
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -138,7 +139,7 @@ func (c *SonobuoyClient) WaitForRun(cfg *RunConfig) error {
 			lastMsg = s
 		}
 	}
-	runCondition := func() (bool, error) {
+	runCondition := func(ctx context.Context) (bool, error) {
 
 		// Get the Aggregator pod and check if its status is completed or terminated.
 		status, pod, err := c.GetStatusPod(&StatusConfig{Namespace: ns})
@@ -186,7 +187,7 @@ func (c *SonobuoyClient) WaitForRun(cfg *RunConfig) error {
 		// handled by conditionFunc since it has to be part of the polling.
 		// Could use channels but that will lead to future issues.
 	}
-	err := wait.Poll(pollInterval, cfg.Wait, runCondition)
+	err := wait.PollUntilContextTimeout(context.TODO(), pollInterval, cfg.Wait, false, runCondition)
 	if err != nil {
 		return errors.Wrap(err, "waiting for run to finish")
 	}