mirror of
https://github.com/zeromicro/go-zero.git
synced 2025-02-03 00:38:40 +08:00
feat: support WithStreamClientInterceptor for zrpc clients (#1907)
* feat: support WithStreamClientInterceptor for zrpc clients * fix: data race
This commit is contained in:
parent
95282edb78
commit
e80a64fa67
@ -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…
Reference in New Issue
Block a user