From fe592e7f804c6b0db5465ad6bd0cab6df64b18b3 Mon Sep 17 00:00:00 2001
From: Michael Gugino <mgugino@redhat.com>
Date: Sat, 23 Nov 2019 13:58:59 -0500
Subject: [PATCH] kubectl/drain: Add context support

This commits allows specifying a context.Context
in the Helper type. This context is utilized to
cancel waitForDelete.

Kubernetes-commit: 8682e902f5487e04b893da7230125db0d7ae66b4
---
 pkg/drain/drain.go      | 15 +++++++++++++--
 pkg/drain/drain_test.go | 31 ++++++++++++++++++++++++++++---
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/pkg/drain/drain.go b/pkg/drain/drain.go
index d123dd33..c922c6aa 100644
--- a/pkg/drain/drain.go
+++ b/pkg/drain/drain.go
@@ -43,6 +43,7 @@ const (
 
 // Helper contains the parameters to control the behaviour of drainer
 type Helper struct {
+	Ctx                 context.Context
 	Client              kubernetes.Interface
 	Force               bool
 	GracePeriodSeconds  int
@@ -203,7 +204,7 @@ func (d *Helper) evictPods(pods []corev1.Pod, policyGroupVersion string, getPodF
 	} else {
 		globalTimeout = d.Timeout
 	}
-	ctx, cancel := context.WithTimeout(context.TODO(), globalTimeout)
+	ctx, cancel := context.WithTimeout(d.getContext(), globalTimeout)
 	defer cancel()
 	for _, pod := range pods {
 		go func(pod corev1.Pod, returnCh chan error) {
@@ -271,7 +272,7 @@ func (d *Helper) deletePods(pods []corev1.Pod, getPodFn func(namespace, name str
 			return err
 		}
 	}
-	ctx := context.TODO()
+	ctx := d.getContext()
 	_, err := waitForDelete(ctx, pods, 1*time.Second, globalTimeout, false, getPodFn, d.OnPodDeletedOrEvicted, globalTimeout)
 	return err
 }
@@ -306,3 +307,13 @@ func waitForDelete(ctx context.Context, pods []corev1.Pod, interval, timeout tim
 	})
 	return pods, err
 }
+
+// Since Helper does not have a constructor, we can't enforce Helper.Ctx != nil
+// Multiple public methods prevent us from initializing the context in a single
+// place as well.
+func (d *Helper) getContext() context.Context {
+	if d.Ctx != nil {
+		return d.Ctx
+	}
+	return context.Background()
+}
diff --git a/pkg/drain/drain_test.go b/pkg/drain/drain_test.go
index d6ecca17..bdbf8d43 100644
--- a/pkg/drain/drain_test.go
+++ b/pkg/drain/drain_test.go
@@ -46,6 +46,7 @@ func TestDeletePods(t *testing.T) {
 		description       string
 		interval          time.Duration
 		timeout           time.Duration
+		ctxTimeoutEarly   bool
 		expectPendingPods bool
 		expectError       bool
 		expectedError     *error
@@ -91,6 +92,22 @@ func TestDeletePods(t *testing.T) {
 				return nil, fmt.Errorf("%q: not found", name)
 			},
 		},
+		{
+			description:       "Context Canceled",
+			interval:          1000 * time.Millisecond,
+			timeout:           5 * time.Second,
+			ctxTimeoutEarly:   true,
+			expectPendingPods: true,
+			expectError:       true,
+			expectedError:     &wait.ErrWaitTimeout,
+			getPodFn: func(namespace, name string) (*corev1.Pod, error) {
+				oldPodMap, _ := createPods(false)
+				if oldPod, found := oldPodMap[name]; found {
+					return &oldPod, nil
+				}
+				return nil, fmt.Errorf("%q: not found", name)
+			},
+		},
 		{
 			description:       "Client error could be passed out",
 			interval:          200 * time.Millisecond,
@@ -107,14 +124,22 @@ func TestDeletePods(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.description, func(t *testing.T) {
 			_, pods := createPods(false)
-			ctx := context.TODO()
+			ctx := context.Background()
+			if test.ctxTimeoutEarly {
+				ctx, _ = context.WithTimeout(ctx, 100*time.Millisecond)
+			}
+			start := time.Now()
 			pendingPods, err := waitForDelete(ctx, pods, test.interval, test.timeout, false, test.getPodFn, nil, time.Duration(math.MaxInt64))
-
+			elapsed := time.Since(start)
 			if test.expectError {
 				if err == nil {
 					t.Fatalf("%s: unexpected non-error", test.description)
 				} else if test.expectedError != nil {
-					if *test.expectedError != err {
+					if test.ctxTimeoutEarly {
+						if elapsed >= test.timeout {
+							t.Fatalf("%s: the supplied context did not effectively cancel the waitForDelete", test.description)
+						}
+					} else if *test.expectedError != err {
 						t.Fatalf("%s: the error does not match expected error", test.description)
 					}
 				}