func TestThriftTracingPropagation(t *testing.T) {
suite := &PropagationTestSuite{
Encoding: EncodingInfo{Format: tchannel.Thrift, HeadersSupported: true},
Register: func(t *testing.T, ch *tchannel.Channel) TracingCall {
opts := &thrift.ClientOptions{HostPort: ch.PeerInfo().HostPort}
thriftClient := thrift.NewClient(ch, ch.PeerInfo().ServiceName, opts)
handler := &ThriftHandler{
TraceHandler: TraceHandler{Ch: ch},
t: t,
thriftClient: gen.NewTChanSimpleServiceClient(thriftClient),
}
// Register Thrift handler
server := thrift.NewServer(ch)
server.Register(gen.NewTChanSimpleServiceServer(handler))
return handler.firstCall
},
TestCases: map[TracerType][]PropagationTestCase{
Noop: {
{ForwardCount: 2, TracingDisabled: true, ExpectedBaggage: "", ExpectedSpanCount: 0},
{ForwardCount: 2, TracingDisabled: false, ExpectedBaggage: "", ExpectedSpanCount: 0},
},
Mock: {
{ForwardCount: 2, TracingDisabled: true, ExpectedBaggage: BaggageValue, ExpectedSpanCount: 0},
{ForwardCount: 2, TracingDisabled: false, ExpectedBaggage: BaggageValue, ExpectedSpanCount: 6},
},
Jaeger: {
{ForwardCount: 2, TracingDisabled: true, ExpectedBaggage: BaggageValue, ExpectedSpanCount: 0},
{ForwardCount: 2, TracingDisabled: false, ExpectedBaggage: BaggageValue, ExpectedSpanCount: 6},
},
},
}
suite.Run(t)
}
func TestRegisterPostResponseCB(t *testing.T) {
withSetup(t, func(ctx Context, args testArgs) {
arg := &gen.Data{
B1: true,
S2: "str",
I3: 102,
}
ret := &gen.Data{
B1: false,
S2: "return-str",
I3: 105,
}
called := make(chan struct{})
cb := func(method string, response thrift.TStruct) {
assert.Equal(t, "Call", method)
res, ok := response.(*gen.SimpleServiceCallResult)
if assert.True(t, ok, "response type should be Result struct") {
assert.Equal(t, ret, res.GetSuccess(), "result should be returned value")
}
close(called)
}
args.server.Register(gen.NewTChanSimpleServiceServer(args.s1), OptPostResponse(cb))
args.s1.On("Call", ctxArg(), arg).Return(ret, nil)
res, err := args.c1.Call(ctx, arg)
require.NoError(t, err, "Call failed")
assert.Equal(t, res, ret, "Call return value wrong")
select {
case <-time.After(time.Second):
t.Errorf("post-response callback not called")
case <-called:
}
})
}
func TestRegisterPostResponseCB(t *testing.T) {
withSetup(t, func(ctx Context, args testArgs) {
arg := &gen.Data{
B1: true,
S2: "str",
I3: 102,
}
ret := &gen.Data{
B1: false,
S2: "return-str",
I3: 105,
}
var called bool
cb := func(method string, response thrift.TStruct) {
assert.Equal(t, "Call", method)
res, ok := response.(*gen.SimpleServiceCallResult)
if assert.True(t, ok, "response type should be Result struct") {
assert.Equal(t, ret, res.GetSuccess(), "result should be returned value")
}
called = true
}
args.server.Register(gen.NewTChanSimpleServiceServer(args.s1), OptPostResponse(cb))
args.s1.On("Call", ctxArg(), arg).Return(ret, nil)
res, err := args.c1.Call(ctx, arg)
assert.NoError(t, err, "Call failed")
assert.Equal(t, res, ret, "Call return value wrong")
assert.True(t, called, "post response callback not called")
})
}
func setupServer(t *testing.T, h *mocks.TChanSimpleService, sh *mocks.TChanSecondService) (*tchannel.Channel, *Server) {
ch := testutils.NewServer(t, nil)
server := NewServer(ch)
server.Register(gen.NewTChanSimpleServiceServer(h))
server.Register(gen.NewTChanSecondServiceServer(sh))
return ch, server
}
func setupServer(h *mocks.TChanSimpleService, sh *mocks.TChanSecondService) (*tchannel.Channel, *Server, error) {
ch, err := testutils.NewServerChannel(nil)
if err != nil {
return nil, nil, err
}
server := NewServer(ch)
server.Register(gen.NewTChanSimpleServiceServer(h))
server.Register(gen.NewTChanSecondServiceServer(sh))
return ch, server, nil
}
func (b *Behavior) registerThrift(ch *tchannel.Channel) {
handler := &thriftHandler{b: b, ch: ch}
server := thrift.NewServer(ch)
server.Register(gen.NewTChanSimpleServiceServer(handler))
b.thriftCall = handler.callDownstream
}