From 6c62b1a2ea2a946d894c288e09b1e1ae012217a9 Mon Sep 17 00:00:00 2001
From: Harshavardhana <harsha@minio.io>
Date: Mon, 4 May 2020 21:57:42 -0700
Subject: [PATCH] fix broken retry tests

---
 cmd/retry_test.go       | 33 +++++++++++++++++++--------------
 pkg/dsync/retry_test.go |  4 +++-
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/cmd/retry_test.go b/cmd/retry_test.go
index cd9a66047..4dc5c5534 100644
--- a/cmd/retry_test.go
+++ b/cmd/retry_test.go
@@ -1,5 +1,5 @@
 /*
- * MinIO Cloud Storage, (C) 2016 MinIO, Inc.
+ * Minio Cloud Storage, (C) 2016-2020 Minio, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,25 +17,26 @@
 package cmd
 
 import (
+	"context"
 	"testing"
 	"time"
 )
 
 // Tests for retry timer.
 func TestRetryTimerSimple(t *testing.T) {
-	doneCh := make(chan struct{})
-	attemptCh := newRetryTimerSimple(doneCh)
+	rctx, cancel := context.WithCancel(context.Background())
+	attemptCh := newRetryTimerSimple(rctx)
 	i := <-attemptCh
 	if i != 0 {
-		close(doneCh)
+		cancel()
 		t.Fatalf("Invalid attempt counter returned should be 0, found %d instead", i)
 	}
 	i = <-attemptCh
 	if i <= 0 {
-		close(doneCh)
+		cancel()
 		t.Fatalf("Invalid attempt counter returned should be greater than 0, found %d instead", i)
 	}
-	close(doneCh)
+	cancel()
 	_, ok := <-attemptCh
 	if ok {
 		t.Fatal("Attempt counter should be closed")
@@ -44,20 +45,23 @@ func TestRetryTimerSimple(t *testing.T) {
 
 // Test retry time with no jitter.
 func TestRetryTimerWithNoJitter(t *testing.T) {
-	doneCh := make(chan struct{})
+	rctx, cancel := context.WithCancel(context.Background())
+
 	// No jitter
-	attemptCh := newRetryTimerWithJitter(time.Millisecond, 5*time.Millisecond, NoJitter, doneCh)
+	attemptCh := newRetryTimerWithJitter(rctx, time.Millisecond, 5*time.Millisecond, NoJitter)
 	i := <-attemptCh
 	if i != 0 {
-		close(doneCh)
+		cancel()
 		t.Fatalf("Invalid attempt counter returned should be 0, found %d instead", i)
 	}
 	// Loop through the maximum possible attempt.
 	for i = range attemptCh {
 		if i == 30 {
-			close(doneCh)
+			break
 		}
 	}
+
+	cancel()
 	_, ok := <-attemptCh
 	if ok {
 		t.Fatal("Attempt counter should be closed")
@@ -66,15 +70,16 @@ func TestRetryTimerWithNoJitter(t *testing.T) {
 
 // Test retry time with Jitter greater than MaxJitter.
 func TestRetryTimerWithJitter(t *testing.T) {
-	doneCh := make(chan struct{})
+	rctx, cancel := context.WithCancel(context.Background())
+
 	// Jitter will be set back to 1.0
-	attemptCh := newRetryTimerWithJitter(defaultRetryUnit, defaultRetryCap, 2.0, doneCh)
+	attemptCh := newRetryTimerWithJitter(rctx, time.Second, 30*time.Second, 2.0)
 	i := <-attemptCh
 	if i != 0 {
-		close(doneCh)
+		cancel()
 		t.Fatalf("Invalid attempt counter returned should be 0, found %d instead", i)
 	}
-	close(doneCh)
+	cancel()
 	_, ok := <-attemptCh
 	if ok {
 		t.Fatal("Attempt counter should be closed")
diff --git a/pkg/dsync/retry_test.go b/pkg/dsync/retry_test.go
index 8fa5fc96d..01387645c 100644
--- a/pkg/dsync/retry_test.go
+++ b/pkg/dsync/retry_test.go
@@ -57,9 +57,11 @@ func TestRetryTimerWithNoJitter(t *testing.T) {
 	// Loop through the maximum possible attempt.
 	for i = range attemptCh {
 		if i == 30 {
-			cancel()
+			break
 		}
 	}
+	cancel()
+
 	_, ok := <-attemptCh
 	if ok {
 		t.Fatal("Attempt counter should be closed")