From 19dc0cb413974014912d6f49f2cb12720d7948c9 Mon Sep 17 00:00:00 2001
From: Benjamin Wang <wachao@vmware.com>
Date: Tue, 13 Dec 2022 05:08:46 +0800
Subject: [PATCH] client: enhance the function shouldRetryWatch and added unit
 test

Signed-off-by: Benjamin Wang <wachao@vmware.com>
---
 client/v3/watch.go      | 10 ++++------
 client/v3/watch_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/client/v3/watch.go b/client/v3/watch.go
index bb0b5485e..6942a948a 100644
--- a/client/v3/watch.go
+++ b/client/v3/watch.go
@@ -18,7 +18,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"strings"
 	"sync"
 	"time"
 
@@ -588,8 +587,7 @@ func (w *watchGrpcStream) run() {
 
 			switch {
 			case pbresp.Created:
-				cancelReasonError := v3rpc.Error(errors.New(pbresp.CancelReason))
-				if shouldRetryWatch(cancelReasonError) {
+				if shouldRetryWatch(pbresp.CancelReason) {
 					var newErr error
 					if wc, newErr = w.newWatchClient(); newErr != nil {
 						w.lg.Error("failed to create a new watch client", zap.Error(newErr))
@@ -717,9 +715,9 @@ func (w *watchGrpcStream) run() {
 	}
 }
 
-func shouldRetryWatch(cancelReasonError error) bool {
-	return (strings.Compare(cancelReasonError.Error(), v3rpc.ErrGRPCInvalidAuthToken.Error()) == 0) ||
-		(strings.Compare(cancelReasonError.Error(), v3rpc.ErrGRPCAuthOldRevision.Error()) == 0)
+func shouldRetryWatch(cancelReason string) bool {
+	return (cancelReason == v3rpc.ErrGRPCInvalidAuthToken.Error()) ||
+		(cancelReason == v3rpc.ErrGRPCAuthOldRevision.Error())
 }
 
 // nextResume chooses the next resuming to register with the grpc stream. Abandoned
diff --git a/client/v3/watch_test.go b/client/v3/watch_test.go
index 2a56ca4a9..0a94f08cd 100644
--- a/client/v3/watch_test.go
+++ b/client/v3/watch_test.go
@@ -17,7 +17,10 @@ package clientv3
 import (
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"go.etcd.io/etcd/api/v3/mvccpb"
+	"go.etcd.io/etcd/api/v3/v3rpc/rpctypes"
 )
 
 func TestEvent(t *testing.T) {
@@ -53,3 +56,39 @@ func TestEvent(t *testing.T) {
 		}
 	}
 }
+
+func TestShouldRetryWatch(t *testing.T) {
+	testCases := []struct {
+		name          string
+		msg           string
+		expectedRetry bool
+	}{
+		{
+			name:          "equal to ErrGRPCInvalidAuthToken",
+			msg:           rpctypes.ErrGRPCInvalidAuthToken.Error(),
+			expectedRetry: true,
+		},
+		{
+			name:          "equal to ErrGRPCAuthOldRevision",
+			msg:           rpctypes.ErrGRPCAuthOldRevision.Error(),
+			expectedRetry: true,
+		},
+		{
+			name:          "valid grpc error but not equal to ErrGRPCInvalidAuthToken or ErrGRPCAuthOldRevision",
+			msg:           rpctypes.ErrGRPCUserEmpty.Error(),
+			expectedRetry: false,
+		},
+		{
+			name:          "invalid grpc error and not equal to ErrGRPCInvalidAuthToken or ErrGRPCAuthOldRevision",
+			msg:           "whatever error message",
+			expectedRetry: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.expectedRetry, shouldRetryWatch(tc.msg))
+		})
+	}
+}