From 202cd0806afc2caec127e978336d13d2c0c780cd Mon Sep 17 00:00:00 2001
From: hlts2 <hiroto.funakoshi.hiroto@gmail.com>
Date: Thu, 19 Sep 2024 19:31:37 +0900
Subject: [PATCH] fix: add mock function

Signed-off-by: hlts2 <hiroto.funakoshi.hiroto@gmail.com>
---
 internal/test/mock/grpc/grpc_client_mock.go | 14 ++++++++++----
 pkg/gateway/mirror/service/mirror_test.go   |  6 ++++++
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/internal/test/mock/grpc/grpc_client_mock.go b/internal/test/mock/grpc/grpc_client_mock.go
index e29fc50c5d..a8c7d43f4d 100644
--- 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
index edf41c8191..b96211d887 100644
--- 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) {},
 							}
 						},
 					},