Explorar el Código

export WithUnaryClientInterceptor

kevin hace 4 años
padre
commit
1abe21aa2a
Se han modificado 2 ficheros con 22 adiciones y 9 borrados
  1. 3 2
      zrpc/client.go
  2. 19 7
      zrpc/client_test.go

+ 3 - 2
zrpc/client.go

@@ -11,8 +11,9 @@ import (
 )
 )
 
 
 var (
 var (
-	WithDialOption = internal.WithDialOption
-	WithTimeout    = internal.WithTimeout
+	WithDialOption             = internal.WithDialOption
+	WithTimeout                = internal.WithTimeout
+	WithUnaryClientInterceptor = internal.WithUnaryClientInterceptor
 )
 )
 
 
 type (
 type (

+ 19 - 7
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()),
 	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)
 	assert.Nil(t, err)
 	clients := []Client{
 	clients := []Client{
 		directClient,
 		directClient,