feat: support WithStreamClientInterceptor for zrpc clients (#1907)

* feat: support WithStreamClientInterceptor for zrpc clients

* fix: data race
master
Kevin Wan 3 years ago committed by GitHub
parent 95282edb78
commit e80a64fa67
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -462,7 +462,7 @@ func TestStructedLogWithDuration(t *testing.T) {
WithDuration(time.Second).Info(message) WithDuration(time.Second).Info(message)
var entry logEntry var entry logEntry
if err := json.Unmarshal([]byte(w.builder.String()), &entry); err != nil { if err := json.Unmarshal([]byte(w.String()), &entry); err != nil {
t.Error(err) t.Error(err)
} }
assert.Equal(t, levelInfo, entry.Level) assert.Equal(t, levelInfo, entry.Level)
@ -515,7 +515,7 @@ func TestErrorfWithWrappedError(t *testing.T) {
defer writer.Store(old) defer writer.Store(old)
Errorf("hello %w", errors.New(message)) Errorf("hello %w", errors.New(message))
assert.True(t, strings.Contains(w.builder.String(), "hello there")) assert.True(t, strings.Contains(w.String(), "hello there"))
} }
func TestMustNil(t *testing.T) { func TestMustNil(t *testing.T) {

@ -38,7 +38,7 @@ func captureOutput(f func()) string {
f() f()
SetLevel(prevLevel) SetLevel(prevLevel)
return w.builder.String() return w.String()
} }
func getContent(jsonStr string) string { func getContent(jsonStr string) string {

@ -15,6 +15,8 @@ var (
WithDialOption = internal.WithDialOption WithDialOption = internal.WithDialOption
// WithNonBlock sets the dialing to be nonblock. // WithNonBlock sets the dialing to be nonblock.
WithNonBlock = internal.WithNonBlock WithNonBlock = internal.WithNonBlock
// WithStreamClientInterceptor is an alias of internal.WithStreamClientInterceptor.
WithStreamClientInterceptor = internal.WithStreamClientInterceptor
// WithTimeout is an alias of internal.WithTimeout. // WithTimeout is an alias of internal.WithTimeout.
WithTimeout = internal.WithTimeout WithTimeout = internal.WithTimeout
// WithTransportCredentials return a func to make the gRPC calls secured with given credentials. // WithTransportCredentials return a func to make the gRPC calls secured with given credentials.

@ -131,6 +131,13 @@ func WithNonBlock() ClientOption {
} }
} }
// WithStreamClientInterceptor returns a func to customize a ClientOptions with given interceptor.
func WithStreamClientInterceptor(interceptor grpc.StreamClientInterceptor) ClientOption {
return func(options *ClientOptions) {
options.DialOptions = append(options.DialOptions, WithStreamClientInterceptors(interceptor))
}
}
// WithTimeout returns a func to customize a ClientOptions with given timeout. // WithTimeout returns a func to customize a ClientOptions with given timeout.
func WithTimeout(timeout time.Duration) ClientOption { func WithTimeout(timeout time.Duration) ClientOption {
return func(options *ClientOptions) { return func(options *ClientOptions) {

@ -31,6 +31,17 @@ func TestWithNonBlock(t *testing.T) {
assert.True(t, options.NonBlock) assert.True(t, options.NonBlock)
} }
func TestWithStreamClientInterceptor(t *testing.T) {
var options ClientOptions
opt := WithStreamClientInterceptor(func(ctx context.Context, desc *grpc.StreamDesc,
cc *grpc.ClientConn, method string, streamer grpc.Streamer,
opts ...grpc.CallOption) (grpc.ClientStream, error) {
return nil, nil
})
opt(&options)
assert.Equal(t, 1, len(options.DialOptions))
}
func TestWithTransportCredentials(t *testing.T) { func TestWithTransportCredentials(t *testing.T) {
var options ClientOptions var options ClientOptions
opt := WithTransportCredentials(nil) opt := WithTransportCredentials(nil)

@ -71,7 +71,7 @@ func (s *rpcServer) Start(register RegisterFn) error {
WithStreamServerInterceptors(streamInterceptors...)) WithStreamServerInterceptors(streamInterceptors...))
server := grpc.NewServer(options...) server := grpc.NewServer(options...)
register(server) register(server)
// we need to make sure all others are wrapped up // we need to make sure all others are wrapped up,
// so we do graceful stop at shutdown phase instead of wrap up phase // so we do graceful stop at shutdown phase instead of wrap up phase
waitForCalled := proc.AddWrapUpListener(func() { waitForCalled := proc.AddWrapUpListener(func() {
server.GracefulStop() server.GracefulStop()

Loading…
Cancel
Save