yab/rpc_encoding_flag_test.go at dev · yarpc/yab · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
package main

import (
	"context"
	"net"
	"testing"

	"github.com/opentracing/opentracing-go"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/yarpc/yab/encoding"
	yabtransport "github.com/yarpc/yab/transport"

	apitransport "go.uber.org/yarpc/api/transport"
	yarpcgrpc "go.uber.org/yarpc/transport/grpc"
)

type unaryHandlerFunc func(context.Context, *apitransport.Request, apitransport.ResponseWriter) error

func (f unaryHandlerFunc) Handle(ctx context.Context, req *apitransport.Request, resw apitransport.ResponseWriter) error {
	return f(ctx, req, resw)
}

type encodingCaptureRouter struct {
	expectedService   string
	expectedProcedure string
	capturedEncoding  chan string
}

func (r *encodingCaptureRouter) Procedures() []apitransport.Procedure {
	return nil
}

func (r *encodingCaptureRouter) Choose(_ context.Context, req *apitransport.Request) (apitransport.HandlerSpec, error) {
	if req.Service == r.expectedService && req.Procedure == r.expectedProcedure {
		select {
		case r.capturedEncoding <- string(req.Encoding):
		default:
		}
		return apitransport.NewUnaryHandlerSpec(unaryHandlerFunc(func(_ context.Context, _ *apitransport.Request, resw apitransport.ResponseWriter) error {
			_, _ = resw.Write([]byte("ok"))
			return nil
		})), nil
	}
	return apitransport.HandlerSpec{}, apitransport.UnrecognizedProcedureError(req)
}

func TestRPCEncodingFlagOverridesGRPCHeader(t *testing.T) {
	const want = "dev-override"

	gt := yarpcgrpc.NewTransport(yarpcgrpc.Tracer(opentracing.NoopTracer{}))
	require.NoError(t, gt.Start())
	defer func() { _ = gt.Stop() }()

	lis, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(t, err)
	defer func() { _ = lis.Close() }()

	router := &encodingCaptureRouter{
		expectedService:   "svc",
		expectedProcedure: "svc::echo",
		capturedEncoding:  make(chan string, 1),
	}

	inbound := gt.NewInbound(lis)
	inbound.SetRouter(router)
	require.NoError(t, inbound.Start())
	defer func() { _ = inbound.Stop() }()

	opts := TransportOptions{
		ServiceName: "svc",
		CallerName:  "caller",
		Peers:       []string{lis.Addr().String()},
		RPCEncoding: want,
	}
	tp, err := getTransport(opts, resolvedProtocolEncoding{protocol: yabtransport.GRPC, enc: encoding.Protobuf}, opentracing.NoopTracer{})
	require.NoError(t, err)

	_, err = tp.Call(context.Background(), &yabtransport.Request{TargetService: "svc", Method: "svc::echo", Body: []byte("hello")})
	require.NoError(t, err)

	select {
	case got := <-router.capturedEncoding:
		assert.Equal(t, want, got)
	default:
		t.Fatal("did not capture inbound encoding")
	}
}