Fix context usage

Author: Jefftree

test/integration/apiserver/coordinated_leader_election_test.go

@@ -63,12 +63,13 @@ func TestSingleLeaseCandidate(t *testing.T) {
 	}
 	defer server.TearDownFn()
 	config := server.ClientConfig
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			cletest := setupCLE(config, ctx, cancel, t)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			cletest := setupCLE(config, ctx, t)
 			defer cletest.cleanup()
 			go cletest.createAndRunFakeController("foo1", "default", "foo", "1.20.0", "1.20.0", tc.preferredStrategy)
 			cletest.pollForLease(ctx, "foo", "default", tc.expectedHolderIdentity)
@@ -101,12 +102,13 @@ func TestSingleLeaseCandidateUsingThirdPartyStrategy(t *testing.T) {
 	}
 	defer server.TearDownFn()
 	config := server.ClientConfig
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			cletest := setupCLE(config, ctx, cancel, t)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			cletest := setupCLE(config, ctx, t)
 			defer cletest.cleanup()
 			go cletest.createAndRunFakeController("foo1", "default", "foo", "1.20.0", "1.20.0", tc.preferredStrategy)
 			cletest.pollForLease(ctx, "foo", "default", tc.expectedHolderIdentity)
@@ -139,12 +141,12 @@ func TestMultipleLeaseCandidate(t *testing.T) {
 	defer server.TearDownFn()
 	config := server.ClientConfig
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
 
-			cletest := setupCLE(config, ctx, cancel, t)
+			cletest := setupCLE(config, ctx, t)
 			defer cletest.cleanup()
 			go cletest.createAndRunFakeController("baz1", "default", "baz", "1.20.0", "1.20.0", tc.preferredStrategy)
 			go cletest.createAndRunFakeController("baz2", "default", "baz", "1.20.0", "1.19.0", tc.preferredStrategy)
@@ -182,12 +184,13 @@ func TestMultipleLeaseCandidateUsingThirdPartyStrategy(t *testing.T) {
 	}
 	defer server.TearDownFn()
 	config := server.ClientConfig
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			cletest := setupCLE(config, ctx, cancel, t)
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			cletest := setupCLE(config, ctx, t)
 			defer cletest.cleanup()
 			go cletest.createAndRunFakeController("baz1", "default", "baz", "1.20.0", "1.20.0", tc.preferredStrategy)
 			go cletest.createAndRunFakeController("baz2", "default", "baz", "1.20.0", "1.19.0", tc.preferredStrategy)
@@ -211,7 +214,8 @@ func TestLeaseSwapIfBetterAvailable(t *testing.T) {
 	config := server.ClientConfig
 
 	ctx, cancel := context.WithCancel(context.Background())
-	cletest := setupCLE(config, ctx, cancel, t)
+	defer cancel()
+	cletest := setupCLE(config, ctx, t)
 	defer cletest.cleanup()
 
 	go cletest.createAndRunFakeController("bar1", "default", "bar", "1.20.0", "1.20.0", v1.OldestEmulationVersion)
@@ -234,7 +238,7 @@ func TestUpgradeSkew(t *testing.T) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	cletest := setupCLE(config, ctx, cancel, t)
+	cletest := setupCLE(config, ctx, t)
 	defer cletest.cleanup()
 
 	go cletest.createAndRunFakeLegacyController("foo1-130", "default", "foobar")
@@ -307,11 +311,12 @@ type cleTest struct {
 	clientset *kubernetes.Clientset
 	t         *testing.T
 	mu        sync.Mutex
+	ctx       context.Context
 	ctxList   map[string]ctxCancelPair
 }
 
 func (t *cleTest) createAndRunFakeLegacyController(name string, namespace string, targetLease string) {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.ctx)
 	t.mu.Lock()
 	defer t.mu.Unlock()
 	t.ctxList[name+"/"+namespace] = ctxCancelPair{ctx, cancel}
@@ -345,7 +350,7 @@ func (t *cleTest) createAndRunFakeController(name string, namespace string, targ
 		t.t.Error(err)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.ctx)
 	t.mu.Lock()
 	t.ctxList[name+"/"+namespace] = ctxCancelPair{ctx, cancel}
 	t.mu.Unlock()
@@ -415,7 +420,7 @@ func (t *cleTest) pollForLease(ctx context.Context, name, namespace string, hold
 		return lease.Spec.HolderIdentity != nil && *lease.Spec.HolderIdentity == *holder, nil
 	})
 	if err != nil {
-		t.t.Fatalf("timeout awiting for Lease %s %s err: %v", name, namespace, err)
+		t.t.Fatalf("timeout waiting for Lease %s %s err: %v", name, namespace, err)
 	}
 }
 
@@ -427,28 +432,23 @@ func (t *cleTest) cancelController(name, namespace string) {
 }
 
 func (t *cleTest) cleanup() {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	for _, c := range t.ctxList {
-		c.cancel()
-	}
 	err := t.clientset.CoordinationV1().Leases("kube-system").Delete(context.TODO(), "leader-election-controller", metav1.DeleteOptions{})
 	if err != nil && !apierrors.IsNotFound(err) {
 		t.t.Error(err)
 	}
 }
 
-func setupCLE(config *rest.Config, ctx context.Context, cancel func(), t *testing.T) *cleTest {
+func setupCLE(config *rest.Config, ctx context.Context, t *testing.T) *cleTest {
 	clientset, err := kubernetes.NewForConfig(config)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	a := ctxCancelPair{ctx, cancel}
 	return &cleTest{
 		config:    config,
 		clientset: clientset,
-		ctxList:   map[string]ctxCancelPair{"main": a},
+		ctx:       ctx,
+		ctxList:   map[string]ctxCancelPair{},
 		t:         t,
 	}
 }