func (h *multiHopHandler) Handle(ctx context.Context, reqMeta yarpc.ReqMeta, body interface{}) (interface{}, yarpc.ResMeta, error) {
if h.phoneClient == nil {
panic("call SetClient() and SetTransport() first")
}
assertBaggageMatches(ctx, h.t, h.wantBaggage)
span := opentracing.SpanFromContext(ctx)
for key, value := range h.addBaggage {
span.SetBaggageItem(key, value)
}
ctx = opentracing.ContextWithSpan(ctx, span)
var resp js.RawMessage
phoneResMeta, err := h.phoneClient.Call(
ctx,
yarpc.NewReqMeta().Procedure("phone").Headers(reqMeta.Headers()),
&server.PhoneRequest{
Service: "ctxclient",
Procedure: h.phoneCallTo,
Transport: h.phoneCallTransport,
Body: &js.RawMessage{'{', '}'},
}, &resp)
resMeta := yarpc.NewResMeta().Headers(phoneResMeta.Headers())
return map[string]interface{}{}, resMeta, err
}
func TestHandleSuccessWithResponseHeaders(t *testing.T) {
h := func(ctx context.Context, r yarpc.ReqMeta, _ *simpleRequest) (*simpleResponse, yarpc.ResMeta, error) {
resMeta := yarpc.NewResMeta().Headers(yarpc.NewHeaders().With("foo", "bar"))
return &simpleResponse{Success: true}, resMeta, nil
}
handler := jsonHandler{
reader: structReader{reflect.TypeOf(simpleRequest{})},
handler: reflect.ValueOf(h),
}
resw := new(transporttest.FakeResponseWriter)
err := handler.Handle(context.Background(), &transport.Request{
Procedure: "simpleCall",
Encoding: "json",
Body: jsonBody(`{"name": "foo", "attributes": {"bar": 42}}`),
}, resw)
require.NoError(t, err)
assert.Equal(t, transport.NewHeaders().With("foo", "bar"), resw.Headers)
}
func resMetaFromReqMeta(reqMeta yarpc.ReqMeta) yarpc.ResMeta {
return yarpc.NewResMeta().Headers(reqMeta.Headers())
}
func (h helloHandler) Echo(ctx context.Context, reqMeta yarpc.ReqMeta, e *echo.EchoRequest) (*echo.EchoResponse, yarpc.ResMeta, error) {
return &echo.EchoResponse{Message: e.Message, Count: e.Count + 1},
yarpc.NewResMeta().Headers(reqMeta.Headers()),
nil
}
func yarpcEcho(ctx context.Context, reqMeta yarpc.ReqMeta, body []byte) ([]byte, yarpc.ResMeta, error) {
return body, yarpc.NewResMeta().Headers(reqMeta.Headers()), nil
}
func TestRawHandler(t *testing.T) {
// handler to use for test cases where the handler should not be called
handlerNotCalled :=
func(ctx context.Context, reqMeta yarpc.ReqMeta, body []byte) ([]byte, yarpc.ResMeta, error) {
t.Errorf("unexpected call handle(%v, %v)", reqMeta, body)
return nil, nil, fmt.Errorf("unexpected call handle(%v, %v)", reqMeta, body)
}
tests := []struct {
procedure string
headers transport.Headers
bodyChunks [][]byte
handler UnaryHandler
wantErr string
wantHeaders transport.Headers
wantBody []byte
}{
{
procedure: "foo",
bodyChunks: [][]byte{
{1, 2, 3},
{4, 5, 6},
},
handler: func(ctx context.Context, reqMeta yarpc.ReqMeta, body []byte) ([]byte, yarpc.ResMeta, error) {
assert.Equal(t, "foo", reqMeta.Procedure())
assert.Equal(t, []byte{1, 2, 3, 4, 5, 6}, body)
return []byte("hello"), nil, nil
},
wantBody: []byte("hello"),
},
{
procedure: "bar",
bodyChunks: [][]byte{
{1, 2, 3},
nil, // triggers a read error
{4, 5, 6},
},
handler: handlerNotCalled,
wantErr: "error set by user",
// TODO consistent error messages between languages
},
{
procedure: "baz",
bodyChunks: [][]byte{},
handler: func(ctx context.Context, reqMeta yarpc.ReqMeta, body []byte) ([]byte, yarpc.ResMeta, error) {
assert.Equal(t, []byte{}, body)
return nil, nil, fmt.Errorf("great sadness")
},
wantErr: "great sadness",
},
{
procedure: "responseHeaders",
bodyChunks: [][]byte{},
handler: func(ctx context.Context, reqMeta yarpc.ReqMeta, body []byte) ([]byte, yarpc.ResMeta, error) {
resMeta := yarpc.NewResMeta().Headers(yarpc.NewHeaders().With("hello", "world"))
return []byte{}, resMeta, nil
},
wantHeaders: transport.NewHeaders().With("hello", "world"),
},
}
for _, tt := range tests {
handler := rawUnaryHandler{tt.handler}
resw := new(transporttest.FakeResponseWriter)
writer, chunkReader := testreader.ChunkReader()
for _, chunk := range tt.bodyChunks {
writer <- chunk
}
close(writer)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
err := handler.Handle(ctx, &transport.Request{
Procedure: tt.procedure,
Headers: tt.headers,
Encoding: "raw",
Body: chunkReader,
}, resw)
if tt.wantErr != "" {
if assert.Error(t, err) {
assert.Equal(t, err.Error(), tt.wantErr)
}
} else {
if assert.NoError(t, err) {
assert.Equal(t, tt.wantHeaders, resw.Headers)
assert.Equal(t, tt.wantBody, resw.Body.Bytes(),
"body does not match for %s", tt.procedure)
}
}
}
}
// Echo endpoint for the Echo service.
func (EchoThrift) Echo(ctx context.Context, reqMeta yarpc.ReqMeta, ping *echo.Ping) (*echo.Pong, yarpc.ResMeta, error) {
return &echo.Pong{Boop: ping.Beep}, yarpc.NewResMeta().Headers(reqMeta.Headers()), nil
}
// EchoJSON implements the echo procedure.
func EchoJSON(ctx context.Context, reqMeta yarpc.ReqMeta, body map[string]interface{}) (map[string]interface{}, yarpc.ResMeta, error) {
return body, yarpc.NewResMeta().Headers(reqMeta.Headers()), nil
}
func (h *singleHopHandler) Handle(ctx context.Context, reqMeta yarpc.ReqMeta, body interface{}) (interface{}, yarpc.ResMeta, error) {
assertBaggageMatches(ctx, h.t, h.wantBaggage)
resMeta := yarpc.NewResMeta().Headers(reqMeta.Headers())
return map[string]interface{}{}, resMeta, nil
}