test: Deflake TestV3LeasePrmote integration test

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")
 	}
 }

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
+			}
+		}
+	}
+}