func TestReadByte0Byte(t *testing.T) {
chunkWriter, chunkReader := testreader.ChunkReader()
reader := getProtocolReader(chunkReader)
chunkWriter <- []byte{}
chunkWriter <- []byte{}
chunkWriter <- []byte{}
chunkWriter <- []byte("abc")
close(chunkWriter)
b, err := reader.transport.ReadByte()
assert.NoError(t, err, "ReadByte should ignore 0 byte reads")
assert.EqualValues(t, 'a', b)
b, err = reader.transport.ReadByte()
assert.NoError(t, err, "ReadByte failed")
assert.EqualValues(t, 'b', b)
b, err = reader.transport.ReadByte()
assert.NoError(t, err, "ReadByte failed")
assert.EqualValues(t, 'c', b)
b, err = reader.transport.ReadByte()
assert.Equal(t, io.EOF, err, "ReadByte should EOF")
}
func TestReadStructErr(t *testing.T) {
writer, reader := testreader.ChunkReader()
writer <- structTest.encoded[:10]
writer <- nil
close(writer)
s := &test.Data{}
err := ReadStruct(reader, s)
if assert.Error(t, err, "ReadStruct should fail") {
// Apache Thrift just prepends the error message, and doesn't give us access
// to the underlying error, so we can't check the underlying error exactly.
assert.Contains(t, err.Error(), testreader.ErrUser.Error(), "Underlying error missing")
}
}
func TestReaderErr(t *testing.T) {
tests := []struct {
chunks [][]byte
validation func(reader *Reader)
}{
{
chunks: [][]byte{
{0, 1},
nil,
{2, 3},
},
validation: func(reader *Reader) {
assert.Equal(t, uint16(1), reader.ReadUint16(), "Read unexpected value")
assert.Equal(t, uint16(0), reader.ReadUint16(), "Expected default value")
},
},
{
chunks: [][]byte{
{0, 4},
[]byte("test"),
nil,
{'A', 'b'},
},
validation: func(reader *Reader) {
assert.Equal(t, "test", reader.ReadLen16String(), "Read unexpected value")
assert.Equal(t, "", reader.ReadString(2), "Expected default value")
},
},
}
for _, tt := range tests {
writer, chunkReader := testreader.ChunkReader()
reader := NewReader(chunkReader)
defer reader.Release()
for _, chunk := range tt.chunks {
writer <- chunk
}
close(writer)
tt.validation(reader)
// Once there's an error, all further calls should fail.
assert.Equal(t, testreader.ErrUser, reader.Err(), "Unexpected error")
assert.Equal(t, uint16(0), reader.ReadUint16(), "Expected default value")
assert.Equal(t, "", reader.ReadString(1), "Expected default value")
assert.Equal(t, "", reader.ReadLen16String(), "Expected default value")
assert.Equal(t, testreader.ErrUser, reader.Err(), "Unexpected error")
}
}
func TestCall(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
ctx := context.Background()
caller := "caller"
service := "service"
tests := []struct {
procedure string
headers yarpc.Headers
body []byte
responseBody [][]byte
want []byte
wantErr string
wantHeaders yarpc.Headers
}{
{
procedure: "foo",
body: []byte{1, 2, 3},
responseBody: [][]byte{{4}, {5}, {6}},
want: []byte{4, 5, 6},
},
{
procedure: "bar",
body: []byte{1, 2, 3},
responseBody: [][]byte{{4}, {5}, nil, {6}},
wantErr: "error set by user",
},
{
procedure: "headers",
headers: yarpc.NewHeaders().With("x", "y"),
body: []byte{},
responseBody: [][]byte{},
want: []byte{},
wantHeaders: yarpc.NewHeaders().With("a", "b"),
},
}
for _, tt := range tests {
outbound := transporttest.NewMockUnaryOutbound(mockCtrl)
client := New(channel.MultiOutbound(caller, service,
transport.Outbounds{
Unary: outbound,
}))
writer, responseBody := testreader.ChunkReader()
for _, chunk := range tt.responseBody {
writer <- chunk
}
close(writer)
outbound.EXPECT().Call(gomock.Any(),
transporttest.NewRequestMatcher(t,
&transport.Request{
Caller: caller,
Service: service,
Procedure: tt.procedure,
Headers: transport.Headers(tt.headers),
Encoding: Encoding,
Body: bytes.NewReader(tt.body),
}),
).Return(
&transport.Response{
Body: ioutil.NopCloser(responseBody),
Headers: transport.Headers(tt.wantHeaders),
}, nil)
resBody, res, err := client.Call(
ctx,
yarpc.NewReqMeta().Procedure(tt.procedure).Headers(tt.headers),
tt.body)
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.want, resBody)
assert.Equal(t, tt.wantHeaders, res.Headers())
}
}
}
}
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)
}
}
}
}