From a9996f87688ea0c131fb18e32d647285e75d7e01 Mon Sep 17 00:00:00 2001
From: Joe Betz <jpbetz@google.com>
Date: Thu, 26 Oct 2017 13:28:51 -0700
Subject: [PATCH] test: Deflake TestV3LeasePrmote integration test

---
 integration/v3_lease_test.go | 27 ++++++++++++++++++++-------
 pkg/testutil/testutil.go     | 27 +++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/integration/v3_lease_test.go b/integration/v3_lease_test.go
index 82f9be24f..7ec2d3c76 100644
--- a/integration/v3_lease_test.go
+++ b/integration/v3_lease_test.go
@@ -36,7 +36,9 @@ func TestV3LeasePrmote(t *testing.T) {
 	defer clus.Terminate(t)
 
 	// create lease
-	lresp, err := toGRPC(clus.RandClient()).Lease.LeaseGrant(context.TODO(), &pb.LeaseGrantRequest{TTL: 5})
+	lresp, err := toGRPC(clus.RandClient()).Lease.LeaseGrant(context.TODO(), &pb.LeaseGrantRequest{TTL: 3})
+	ttl := time.Duration(lresp.TTL) * time.Second
+	afterGrant := time.Now()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -45,10 +47,11 @@ func TestV3LeasePrmote(t *testing.T) {
 	}
 
 	// wait until the lease is going to expire.
-	time.Sleep(time.Duration(lresp.TTL-1) * time.Second)
+	time.Sleep(time.Until(afterGrant.Add(ttl - time.Second)))
 
 	// kill the current leader, all leases should be refreshed.
 	toStop := clus.waitLeader(t, clus.Members)
+	beforeStop := time.Now()
 	clus.Members[toStop].Stop(t)
 
 	var toWait []*member
@@ -60,19 +63,29 @@ func TestV3LeasePrmote(t *testing.T) {
 	clus.waitLeader(t, toWait)
 	clus.Members[toStop].Restart(t)
 	clus.waitLeader(t, clus.Members)
+	afterReelect := time.Now()
 
 	// ensure lease is refreshed by waiting for a "long" time.
 	// it was going to expire anyway.
-	time.Sleep(3 * time.Second)
+	time.Sleep(time.Until(beforeStop.Add(ttl - time.Second)))
 
 	if !leaseExist(t, clus, lresp.ID) {
 		t.Error("unexpected lease not exists")
 	}
 
-	// let lease expires. total lease = 5 seconds and we already
-	// waits for 3 seconds, so 3 seconds more is enough.
-	time.Sleep(3 * time.Second)
-	if leaseExist(t, clus, lresp.ID) {
+	// wait until the renewed lease is expected to expire.
+	time.Sleep(time.Until(afterReelect.Add(ttl)))
+
+	// wait for up to 10 seconds for lease to expire.
+	expiredCondition := func() (bool, error) {
+		return !leaseExist(t, clus, lresp.ID), nil
+	}
+	expired, err := testutil.Poll(100*time.Millisecond, 10*time.Second, expiredCondition)
+	if err != nil {
+		t.Error(err)
+	}
+
+	if !expired {
 		t.Error("unexpected lease exists")
 	}
 }
diff --git a/pkg/testutil/testutil.go b/pkg/testutil/testutil.go
index db2dd3285..0e02ddc0d 100644
--- a/pkg/testutil/testutil.go
+++ b/pkg/testutil/testutil.go
@@ -55,3 +55,30 @@ func FatalStack(t *testing.T, s string) {
 	t.Error(string(stackTrace[:n]))
 	t.Fatalf(s)
 }
+
+// ConditionFunc returns true when a condition is met.
+type ConditionFunc func() (bool, error)
+
+// Poll calls a condition function repeatedly on a polling interval until it returns true, returns an error
+// or the timeout is reached. If the condition function returns true or an error before the timeout, Poll
+// immediately returns with the true value or the error. If the timeout is exceeded, Poll returns false.
+func Poll(interval time.Duration, timeout time.Duration, condition ConditionFunc) (bool, error) {
+	timeoutCh := time.After(timeout)
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-timeoutCh:
+			return false, nil
+		case <-ticker.C:
+			success, err := condition()
+			if err != nil {
+				return false, err
+			}
+			if success {
+				return true, nil
+			}
+		}
+	}
+}