func TestClientHostPort(t *testing.T) {
ctx, cancel := NewContext(time.Second)
defer cancel()
s1ch := testutils.NewServer(t, nil)
s2ch := testutils.NewServer(t, nil)
defer s1ch.Close()
defer s2ch.Close()
s1ch.Peers().Add(s2ch.PeerInfo().HostPort)
s2ch.Peers().Add(s1ch.PeerInfo().HostPort)
mock1, mock2 := new(mocks.TChanSecondService), new(mocks.TChanSecondService)
NewServer(s1ch).Register(gen.NewTChanSecondServiceServer(mock1))
NewServer(s2ch).Register(gen.NewTChanSecondServiceServer(mock2))
// When we call using a normal client, it can only call the other server (only peer).
c1 := gen.NewTChanSecondServiceClient(NewClient(s1ch, s2ch.PeerInfo().ServiceName, nil))
mock2.On("Echo", ctxArg(), "call1").Return("call1", nil)
res, err := c1.Echo(ctx, "call1")
assert.NoError(t, err, "call1 failed")
assert.Equal(t, "call1", res)
// When we call using a client that specifies host:port, it should call that server.
c2 := gen.NewTChanSecondServiceClient(NewClient(s1ch, s1ch.PeerInfo().ServiceName, &ClientOptions{
HostPort: s1ch.PeerInfo().HostPort,
}))
mock1.On("Echo", ctxArg(), "call2").Return("call2", nil)
res, err = c2.Echo(ctx, "call2")
assert.NoError(t, err, "call2 failed")
assert.Equal(t, "call2", res)
}
func setupBenchServer() ([]string, error) {
ch, err := testutils.NewServerChannel(testutils.NewOpts().
SetServiceName(benchServerName).
SetFramePool(tchannel.NewSyncFramePool()))
if err != nil {
return nil, err
}
fmt.Println(benchServerName, "started on", ch.PeerInfo().HostPort)
server := thrift.NewServer(ch)
server.Register(gen.NewTChanSecondServiceServer(benchSecondHandler{}))
if !*useHyperbahn {
return []string{ch.PeerInfo().HostPort}, nil
}
// Set up a Hyperbahn client and advertise it.
nodes := strings.Split(*hyperbahnNodes, ",")
config := hyperbahn.Configuration{InitialNodes: nodes}
hc, err := hyperbahn.NewClient(ch, config, nil)
if err := hc.Advertise(); err != nil {
return nil, err
}
return nodes, nil
}
func TestRequestSubChannel(t *testing.T) {
ctx, cancel := NewContext(time.Second)
defer cancel()
tchan, err := tchannel.NewChannel("svc1", nil)
require.NoError(t, err, "server NewChannel failed")
require.NoError(t, tchan.ListenAndServe(":0"), "Listen failed")
defer tchan.Close()
clientCh, err := tchannel.NewChannel("client", nil)
require.NoError(t, err, "client NewChannel failed")
defer clientCh.Close()
clientCh.Peers().Add(tchan.PeerInfo().HostPort)
tests := []tchannel.Registrar{tchan, tchan.GetSubChannel("svc2"), tchan.GetSubChannel("svc3")}
for _, ch := range tests {
mockHandler := new(mocks.TChanSecondService)
server := NewServer(ch)
server.Register(gen.NewTChanSecondServiceServer(mockHandler))
client := NewClient(clientCh, ch.ServiceName(), nil)
secondClient := gen.NewTChanSecondServiceClient(client)
echoArg := ch.ServiceName()
echoRes := echoArg + "-echo"
mockHandler.On("Echo", ctxArg(), echoArg).Return(echoRes, nil)
res, err := secondClient.Echo(ctx, echoArg)
assert.NoError(t, err, "Echo failed")
assert.Equal(t, echoRes, res)
}
}
// NewServer returns a new Server that can recieve Thrift calls or raw calls.
func NewServer(optFns ...Option) Server {
opts := getOptions(optFns)
if opts.external {
return newExternalServer(opts)
}
ch, err := tchannel.NewChannel(opts.svcName, &tchannel.ChannelOptions{
Logger: tchannel.NewLevelLogger(tchannel.NewLogger(os.Stderr), tchannel.LogLevelWarn),
})
if err != nil {
panic("failed to create channel: " + err.Error())
}
if err := ch.ListenAndServe("127.0.0.1:0"); err != nil {
panic("failed to listen on port 0: " + err.Error())
}
s := &internalServer{
ch: ch,
opts: opts,
}
tServer := thrift.NewServer(ch)
tServer.Register(gen.NewTChanSecondServiceServer(handler{calls: &s.thriftCalls}))
ch.Register(raw.Wrap(rawHandler{calls: &s.rawCalls}), "echo")
if len(opts.advertiseHosts) > 0 {
if err := s.Advertise(opts.advertiseHosts); err != nil {
panic("failed to advertise: " + err.Error())
}
}
return s
}
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 TestInvalidThriftBytes(t *testing.T) {
ctx, cancel := NewContext(time.Second)
defer cancel()
ch := testutils.NewClient(t, nil)
sCh := testutils.NewServer(t, nil)
defer sCh.Close()
svr := NewServer(sCh)
svr.Register(gen.NewTChanSecondServiceServer(new(mocks.TChanSecondService)))
tests := []struct {
name string
arg3 []byte
}{
{
name: "missing bytes",
arg3: serializeStruct(t, &gen.SecondServiceEchoArgs{Arg: "Hello world"})[:5],
},
{
name: "wrong struct",
arg3: serializeStruct(t, &gen.Data{B1: true}),
},
}
for _, tt := range tests {
sPeer := sCh.PeerInfo()
call, err := ch.BeginCall(ctx, sPeer.HostPort, sPeer.ServiceName, "SecondService::Echo", &tchannel.CallOptions{
Format: tchannel.Thrift,
})
require.NoError(t, err, "BeginCall failed")
require.NoError(t, tchannel.NewArgWriter(call.Arg2Writer()).Write([]byte{0, 0}), "Write arg2 failed")
writer, err := call.Arg3Writer()
require.NoError(t, err, "Arg3Writer failed")
_, err = writer.Write(tt.arg3)
require.NoError(t, err, "Write arg3 failed")
require.NoError(t, writer.Close(), "Close failed")
response := call.Response()
_, _, err = raw.ReadArgsV2(response)
assert.Error(t, err, "%v: Expected error", tt.name)
assert.Equal(t, tchannel.ErrCodeBadRequest, tchannel.GetSystemErrorCode(err),
"%v: Expected bad request, got %v", tt.name, err)
}
}
