func TestInboundConnection_CallOptions(t *testing.T) {
ctx, cancel := NewContext(time.Second)
defer cancel()
testutils.WithTestServer(t, nil, func(server *testutils.TestServer) {
server.RegisterFunc("test", func(ctx context.Context, args *raw.Args) (*raw.Res, error) {
assert.Equal(t, "client", CurrentCall(ctx).CallerName(), "Expected caller name to be passed through")
return &raw.Res{}, nil
})
backendName := server.ServiceName()
proxyCh := server.NewServer(&testutils.ChannelOpts{ServiceName: "proxy"})
defer proxyCh.Close()
subCh := proxyCh.GetSubChannel(backendName)
subCh.SetHandler(HandlerFunc(func(ctx context.Context, inbound *InboundCall) {
outbound, err := proxyCh.BeginCall(ctx, server.HostPort(), backendName, inbound.MethodString(), inbound.CallOptions())
require.NoError(t, err, "Create outbound call failed")
arg2, arg3, _, err := raw.WriteArgs(outbound, []byte("hello"), []byte("world"))
require.NoError(t, err, "Write outbound call failed")
require.NoError(t, raw.WriteResponse(inbound.Response(), &raw.Res{
Arg2: arg2,
Arg3: arg3,
}), "Write response failed")
}))
clientCh := server.NewClient(&testutils.ChannelOpts{
ServiceName: "client",
})
defer clientCh.Close()
_, _, _, err := raw.Call(ctx, clientCh, proxyCh.PeerInfo().HostPort, backendName, "test", nil, nil)
require.NoError(t, err, "Call through proxy failed")
})
}
func TestSetHandler(t *testing.T) {
// Generate a Handler that expects only the given methods to be called.
genHandler := func(methods ...string) Handler {
allowedMethods := make(map[string]struct{}, len(methods))
for _, m := range methods {
allowedMethods[m] = struct{}{}
}
return HandlerFunc(func(ctx context.Context, call *InboundCall) {
method := call.MethodString()
assert.Contains(t, allowedMethods, method, "unexpected call to %q", method)
err := raw.WriteResponse(call.Response(), &raw.Res{Arg3: []byte(method)})
require.NoError(t, err)
})
}
ch := testutils.NewServer(t, testutils.NewOpts().
AddLogFilter("Couldn't find handler", 1, "serviceName", "svc2", "method", "bar"))
defer ch.Close()
// Catch-all handler for the main channel that accepts foo, bar, and baz,
// and a single registered handler for a different subchannel.
ch.GetSubChannel("svc1").SetHandler(genHandler("foo", "bar", "baz"))
ch.GetSubChannel("svc2").Register(genHandler("foo"), "foo")
client := testutils.NewClient(t, nil)
client.Peers().Add(ch.PeerInfo().HostPort)
defer client.Close()
tests := []struct {
Service string
Method string
ShouldFail bool
}{
{"svc1", "foo", false},
{"svc1", "bar", false},
{"svc1", "baz", false},
{"svc2", "foo", false},
{"svc2", "bar", true},
}
for _, tt := range tests {
c := client.GetSubChannel(tt.Service)
ctx, _ := NewContext(time.Second)
_, data, _, err := raw.CallSC(ctx, c, tt.Method, nil, []byte("irrelevant"))
if tt.ShouldFail {
require.Error(t, err)
} else {
require.NoError(t, err)
assert.Equal(t, tt.Method, string(data))
}
}
st := ch.IntrospectState(nil)
assert.Equal(t, "overriden", st.SubChannels["svc1"].Handler.Type.String())
assert.Nil(t, st.SubChannels["svc1"].Handler.Methods)
assert.Equal(t, "methods", st.SubChannels["svc2"].Handler.Type.String())
assert.Equal(t, []string{"foo"}, st.SubChannels["svc2"].Handler.Methods)
}