diff --git a/zrpc/client.go b/zrpc/client.go index 3c1ff363..81d370f3 100644 --- a/zrpc/client.go +++ b/zrpc/client.go @@ -11,8 +11,9 @@ import ( ) var ( - WithDialOption = internal.WithDialOption - WithTimeout = internal.WithTimeout + WithDialOption = internal.WithDialOption + WithTimeout = internal.WithTimeout + WithUnaryClientInterceptor = internal.WithUnaryClientInterceptor ) type ( diff --git a/zrpc/client_test.go b/zrpc/client_test.go index b048f32c..90baa997 100644 --- a/zrpc/client_test.go +++ b/zrpc/client_test.go @@ -60,14 +60,26 @@ func TestDepositServer_Deposit(t *testing.T) { }, } - directClient := MustNewClient(RpcClientConf{ - Endpoints: []string{"foo"}, - App: "foo", - Token: "bar", - Timeout: 1000, - }, WithDialOption(grpc.WithInsecure()), WithDialOption(grpc.WithContextDialer(dialer()))) + directClient := MustNewClient( + RpcClientConf{ + Endpoints: []string{"foo"}, + App: "foo", + Token: "bar", + Timeout: 1000, + }, + WithDialOption(grpc.WithInsecure()), + WithDialOption(grpc.WithContextDialer(dialer())), + WithUnaryClientInterceptor(func(ctx context.Context, method string, req, reply interface{}, + cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + return invoker(ctx, method, req, reply, cc, opts...) + }), + ) targetClient, err := NewClientWithTarget("foo", WithDialOption(grpc.WithInsecure()), - WithDialOption(grpc.WithContextDialer(dialer()))) + WithDialOption(grpc.WithContextDialer(dialer())), WithUnaryClientInterceptor( + func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, + invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + return invoker(ctx, method, req, reply, cc, opts...) + })) assert.Nil(t, err) clients := []Client{ directClient,