From d41d8ac1da262dda6ba9d79081dae748a913ee6a Mon Sep 17 00:00:00 2001
From: Benjamin Wang <wachao@vmware.com>
Date: Fri, 10 Feb 2023 20:15:01 +0800
Subject: [PATCH] test:enhance the test case TestV3WatchProgressOnMemberRestart

Signed-off-by: Benjamin Wang <wachao@vmware.com>
---
 tests/integration/v3_watch_test.go | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/tests/integration/v3_watch_test.go b/tests/integration/v3_watch_test.go
index 5c86d8e56..f21a1976d 100644
--- a/tests/integration/v3_watch_test.go
+++ b/tests/integration/v3_watch_test.go
@@ -895,6 +895,7 @@ func TestV3WatchProgressOnMemberRestart(t *testing.T) {
 	defer cancel()
 
 	errC := make(chan error, 1)
+	watchReady := make(chan struct{}, 1)
 	doneC := make(chan struct{}, 1)
 	progressNotifyC := make(chan struct{}, 1)
 	go func() {
@@ -906,13 +907,14 @@ func TestV3WatchProgressOnMemberRestart(t *testing.T) {
 		)
 
 		wch := client.Watch(ctx, "foo", clientv3.WithProgressNotify())
+		watchReady <- struct{}{}
 		for wr := range wch {
 			if wr.Err() != nil {
 				errC <- fmt.Errorf("watch error: %w", wr.Err())
 				return
 			}
 
-			if wr.IsProgressNotify() {
+			if len(wr.Events) == 0 {
 				// We need to make sure at least one progress notification
 				// is received after receiving the normal watch response
 				// and before restarting the member.
@@ -939,6 +941,11 @@ func TestV3WatchProgressOnMemberRestart(t *testing.T) {
 		}
 	}()
 
+	// waiting for the watcher ready
+	t.Log("Waiting for the watcher to be ready.")
+	<-watchReady
+	time.Sleep(time.Second)
+
 	// write the key before the member restarts
 	t.Log("Writing key 'foo'")
 	_, err := client.Put(ctx, "foo", "bar1")
@@ -947,14 +954,18 @@ func TestV3WatchProgressOnMemberRestart(t *testing.T) {
 	// make sure at least one progress notification is received
 	// before restarting the member
 	t.Log("Waiting for the progress notification")
-	<-progressNotifyC
+	select {
+	case <-progressNotifyC:
+	case <-time.After(5 * time.Second):
+		t.Log("Do not receive the progress notification in 5 seconds, move forward anyway.")
+	}
 
 	// restart the member
 	t.Log("Restarting the member")
 	clus.Members[0].Stop(t)
 	t.Log("The member stopped")
 	clus.Members[0].Restart(t)
-	clus.WaitLeader(t)
+	clus.Members[0].WaitOK(t)
 	t.Log("The member restarted")
 
 	// write the same key again after the member restarted
@@ -968,7 +979,7 @@ func TestV3WatchProgressOnMemberRestart(t *testing.T) {
 		t.Fatal(err)
 	case <-doneC:
 		t.Log("Done")
-	case <-time.After(20 * time.Second):
+	case <-time.After(15 * time.Second):
 		t.Fatal("Timed out waiting for the response")
 	}
 }