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 }