fix: add mock function

Signed-off-by: hlts2 <[email protected]>
vdaas · Sep 19, 2024 · 202cd08 · 202cd08
1 parent 41f84ae
commit 202cd08
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/internal/test/mock/grpc/grpc_client_mock.go b/internal/test/mock/grpc/grpc_client_mock.go
@@ -30,10 +30,11 @@ type GRPCClientMock struct {
 			addr string,
 			conn *grpc.ClientConn,
 			copts ...grpc.CallOption) error) error
-	ConnectFunc        func(ctx context.Context, addr string, dopts ...grpc.DialOption) (pool.Conn, error)
-	DisconnectFunc     func(ctx context.Context, addr string) error
-	IsConnectedFunc    func(ctx context.Context, addr string) bool
-	ConnectedAddrsFunc func() []string
+	ConnectFunc                  func(ctx context.Context, addr string, dopts ...grpc.DialOption) (pool.Conn, error)
+	DisconnectFunc               func(ctx context.Context, addr string) error
+	IsConnectedFunc              func(ctx context.Context, addr string) bool
+	ConnectedAddrsFunc           func() []string
+	SetDisableResolveDNSAddrFunc func(addr string, disabled bool)
 }
 
 // OrderedRangeConcurrent calls the OrderedRangeConcurrentFunc object.
@@ -70,3 +71,8 @@ func (gc *GRPCClientMock) Disconnect(ctx context.Context, addr string) error {
 func (gc *GRPCClientMock) IsConnected(ctx context.Context, addr string) bool {
 	return gc.IsConnectedFunc(ctx, addr)
 }
+
+// SetDisableResolveDNSAddr calls the SetDisableResolveDNSAddr object.
+func (gc *GRPCClientMock) SetDisableResolveDNSAddr(addr string, disabled bool) {
+	gc.SetDisableResolveDNSAddrFunc(addr, disabled)
+}
diff --git a/pkg/gateway/mirror/service/mirror_test.go b/pkg/gateway/mirror/service/mirror_test.go
@@ -87,6 +87,7 @@ func Test_mirr_Connect(t *testing.T) {
 								ConnectFunc: func(_ context.Context, _ string, _ ...grpc.DialOption) (conn pool.Conn, err error) {
 									return conn, err
 								},
+								SetDisableResolveDNSAddrFunc: func(addr string, disabled bool) {},
 							}
 						},
 					},
@@ -118,6 +119,7 @@ func Test_mirr_Connect(t *testing.T) {
 								ConnectFunc: func(_ context.Context, _ string, _ ...grpc.DialOption) (pool.Conn, error) {
 									return nil, errors.New("missing port in address")
 								},
+								SetDisableResolveDNSAddrFunc: func(addr string, disabled bool) {},
 							}
 						},
 					},
@@ -221,6 +223,7 @@ func Test_mirr_Disconnect(t *testing.T) {
 								DisconnectFunc: func(_ context.Context, _ string) error {
 									return nil
 								},
+								SetDisableResolveDNSAddrFunc: func(addr string, disabled bool) {},
 							}
 						},
 					},
@@ -252,6 +255,7 @@ func Test_mirr_Disconnect(t *testing.T) {
 								DisconnectFunc: func(_ context.Context, _ string) error {
 									return errors.New("missing port in address")
 								},
+								SetDisableResolveDNSAddrFunc: func(addr string, disabled bool) {},
 							}
 						},
 					},
@@ -373,6 +377,7 @@ func Test_mirr_MirrorTargets(t *testing.T) {
 								IsConnectedFunc: func(_ context.Context, addr string) bool {
 									return connected[addr]
 								},
+								SetDisableResolveDNSAddrFunc: func(addr string, disabled bool) {},
 							}
 						},
 					},
@@ -498,6 +503,7 @@ func Test_mirr_connectedOtherMirrorAddrs(t *testing.T) {
 								IsConnectedFunc: func(_ context.Context, addr string) bool {
 									return connected[addr]
 								},
+								SetDisableResolveDNSAddrFunc: func(addr string, disabled bool) {},
 							}
 						},
 					},