Repository: yarpc/yab
Pull request #317 (2021-01-09 17:45:00)

gRPC streams support - 4
## Aim
Support gRPC streams in Yab

## Changes
* Add stream request dispatcher for streaming RPC methods
* Supports interactive IO for streaming methods, i.e. for bi-directional streams you could input the stream message in STDIN and see the stream response
* Concurrently send and receive stream messages for Bidirectional streaming RPC

Next Stack: #319

## How to test?
Download [grpc-stream-test-server.zip](https://github.com/yarpc/yab/files/5826809/grpc-stream-test-server.zip), follow the `readme.md` to run the test server and dispatch streaming yab queries

### Client Stream
```
yab  -F testprotoset --service "test-server" --method uber.yarpc.encoding.protobuf.Test/ClientStream -p grpc://localhost:55555 -r '{"value": "hi"}{"value": "world"}'

{
    "value" "2" // count of the number of requests sent
}
```
STDIN mode
```
yab  -F testprotoset --service "test-server" --method uber.yarpc.encoding.protobuf.Test/ClientStream -p grpc://localhost:55555 -r '-'

{"value": "req-msg-1"}
{"value": "req-msg-2"}

{
    "value" "2" // count of the number of requests sent
}
```

### Server Stream
```
yab  -F testprotoset --service "test-server" --method uber.yarpc.encoding.protobuf.Test/ServerStream -p grpc://localhost:55555 -r '{"value": "hi"}'

// Responds with n messages, where n = len(req.value)
{
  "value": "hi"
}

{
  "value": "hi"
}
```

### Bidirectional Stream
```
yab -F testprotoset --service "test-server" --method uber.yarpc.encoding.protobuf.Test/Stream -p grpc://localhost:55555 -r '{"value": "h"}{"value": "i"}{"value": " world"}'
{
  "value": "h"
}

{
  "value": "hi"
}

{
  "value": "hi world"
}
```

STDIN mode

```
yab -F testprotoset --service "test-server" --method uber.yarpc.encoding.protobuf.Test/Stream -p grpc://localhost:55555 -r '-'
{"value": "req-1"}
{
  "value": "req-1"
}

{"value": "req-2"}
{
  "value": "req-1req-2"
}

```



codecov[bot]: # [Codecov](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=h1) Report
> Merging [#317](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=desc) (ba84158) into [grpc-streams-3](https://codecov.io/gh/yarpc/yab/commit/2166cf9e7b0349e5c5ba233d87982c5db8a63385?el=desc) (2166cf9) will **decrease** coverage by `0.65%`.
> The diff coverage is `60.57%`.

[![Impacted file tree graph](https://codecov.io/gh/yarpc/yab/pull/317/graphs/tree.svg?width=650&height=150&src=pr&token=yEgrqUTCuh)](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=tree)

```diff
@@                Coverage Diff                 @@
##           grpc-streams-3     #317      +/-   ##
==================================================
- Coverage           79.43%   78.78%   -0.66%     
==================================================
  Files                  60       60              
  Lines                2966     3068     +102     
==================================================
+ Hits                 2356     2417      +61     
- Misses                507      541      +34     
- Partials              103      110       +7     
```


| [Impacted Files](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [testdata/protobuf/simple/simple.pb.go](https://codecov.io/gh/yarpc/yab/pull/317/diff?src=pr&el=tree#diff-dGVzdGRhdGEvcHJvdG9idWYvc2ltcGxlL3NpbXBsZS5wYi5nbw==) | `25.00% <23.68%> (-0.76%)` | :arrow_down: |
| [main.go](https://codecov.io/gh/yarpc/yab/pull/317/diff?src=pr&el=tree#diff-bWFpbi5nbw==) | `89.65% <81.81%> (-2.39%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=footer). Last update [2166cf9...ba84158](https://codecov.io/gh/yarpc/yab/pull/317?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

jronak: @Dogild  Addressed the comments, going to update the PR summary with the Yab streaming request example once the sample server (lottery/proberFx) is ready to accept streaming RPCs.

Will include the example in the commit summary when I squash and merge

Patch:
diff --git a/handler.go b/handler.go
new file mode 100644
index 0000000..4ca4557
--- /dev/null
+++ b/handler.go
@@ -0,0 +1,344 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"sync"
+
+	"github.com/uber/tchannel-go"
+	"github.com/yarpc/yab/encoding"
+	"github.com/yarpc/yab/transport"
+	yarpctransport "go.uber.org/yarpc/api/transport"
+	"go.uber.org/zap"
+)
+
+// streamRequestSupplierFn is a function that returns the stream message body
+// if supplier has no more messages it must return io.EOF
+type streamRequestSupplierFn func() (requestBody []byte, err error)
+
+// streamResponseHandlerFn is a function that receives the stream response body
+type streamResponseHandlerFn func(responseBody []byte) error
+
+type requestHandler struct {
+	out        output
+	logger     *zap.Logger
+	opts       Options
+	transport  transport.Transport
+	resolved   resolvedProtocolEncoding
+	serializer encoding.Serializer
+	body       io.Reader
+	headers    map[string]string
+}
+
+func (r requestHandler) handle() {
+	if r.isStreamingMethod() {
+		r.handleStreamRequest()
+		return
+	}
+
+	r.handleUnaryRequest()
+}
+
+// handleUnaryRequest launches initial unary request and unary benchmark
+func (r requestHandler) handleUnaryRequest() {
+	reqInput, err := ioutil.ReadAll(r.body)
+	if err != nil {
+		r.out.Fatalf("Failed while reading the input: %v\n", err)
+	}
+
+	req, err := r.serializer.Request(reqInput)
+	if err != nil {
+		r.out.Fatalf("Failed while serializing the input: %v\n", err)
+	}
+
+	req, err = prepareRequest(req, r.headers, r.opts)
+	if err != nil {
+		r.out.Fatalf("Failed while preparing the request: %v\n", err)
+	}
+
+	// Decides if warm requests must be dispatched before benchmark.
+	if r.shouldMakeInitialRequest() {
+		makeInitialRequest(r.out, r.transport, r.serializer, req)
+	}
+
+	runBenchmark(r.out, r.logger, r.opts, r.resolved, benchmarkMethod{
+		serializer: r.serializer,
+		req:        req,
+	})
+}
+
+// handleStreamRequest launches initial stream request and stream benchmark
+func (r requestHandler) handleStreamRequest() {
+	streamSerializer, ok := r.serializer.(encoding.StreamSerializer)
+	if !ok {
+		r.out.Fatalf("Serializer does not support streaming: %v\n", r.serializer.Encoding())
+	}
+
+	streamReq, streamMsgReader, err := streamSerializer.StreamRequest(r.body)
+	if err != nil {
+		r.out.Fatalf("Failed to create streaming request: %v\n", err)
+	}
+
+	streamReq.Request, err = prepareRequest(streamReq.Request, r.headers, r.opts)
+	if err != nil {
+		r.out.Fatalf("Failed while preparing the request: %v\n", err)
+	}
+
+	if r.shouldMakeInitialRequest() {
+		if err = makeStreamRequest(r.transport, streamReq, r.serializer, streamRequestSupplier(streamMsgReader), r.responseHandler); err != nil {
+			r.out.Fatalf("%v\n", err)
+		}
+	}
+
+	// TODO: handle streaming benchmark.
+}
+
+// shouldMakeInitialRequest returns true if initial request must be made
+func (r requestHandler) shouldMakeInitialRequest() bool {
+	return !(r.opts.BOpts.enabled() && r.opts.BOpts.WarmupRequests == 0)
+}
+
+// responseHandler validates the response bytes and prints indented JSON body
+func (r requestHandler) responseHandler(body []byte) error {
+	res, err := r.serializer.Response(&transport.Response{Body: body})
+	if err != nil {
+		return fmt.Errorf("Failed while serializing stream response: %v", err)
+	}
+
+	bs, err := json.MarshalIndent(res, "", "  ")
+	if err != nil {
+		return fmt.Errorf("Failed to convert map to JSON: %v\nMap: %+v", err, res)
+	}
+
+	r.out.Printf("%s\n\n", bs)
+	return nil
+}
+
+// isStreamingMethod returns true if RPC is streaming type
+func (r requestHandler) isStreamingMethod() bool {
+	return r.serializer.MethodType() != encoding.Unary
+}
+
+// streamRequestSupplier returns a stream request supplier function which uses
+// provided stream request reader
+func streamRequestSupplier(streamMsgReader encoding.StreamRequestReader) streamRequestSupplierFn {
+	nextBodyFn := func() ([]byte, error) {
+		msg, err := streamMsgReader.NextBody()
+		if err == io.EOF {
+			return nil, err
+		}
+		if err != nil {
+			return nil, fmt.Errorf("Failed while reading stream input: %v", err)
+		}
+
+		return msg, nil
+	}
+	return nextBodyFn
+}
+
+// makeStreamRequest opens a stream rpc from the given transport and stream request
+// it then delegates to handler based on rpc type to handle request and response of the stream
+// nextBodyFn is called to get the next stream message body
+// responseHandlerFn is called with the response of the stream
+func makeStreamRequest(t transport.Transport, streamReq *transport.StreamRequest, serializer encoding.Serializer,
+	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+
+	streamTransport, ok := t.(transport.StreamTransport)
+	if !ok {
+		return fmt.Errorf("Transport does not support stream calls: %q", t.Protocol())
+	}
+
+	// Uses tchannel context to remain compatible with tchannel transport
+	// although it does not support streaming, this needs to be removed later.
+	ctx, cancel := tchannel.NewContext(streamReq.Request.Timeout)
+	defer cancel()
+	ctx = makeContextWithTrace(ctx, t, streamReq.Request, 0)
+
+	stream, err := streamTransport.CallStream(ctx, streamReq)
+	if err != nil {
+		return fmt.Errorf("Failed while making stream call: %v", err)
+	}
+
+	switch serializer.MethodType() {
+	case encoding.BidirectionalStream:
+		return makeBidiStream(ctx, stream, nextBodyFn, responseHandlerFn)
+	case encoding.ClientStream:
+		return makeClientStream(ctx, stream, nextBodyFn, responseHandlerFn)
+	default:
+		return makeServerStream(ctx, stream, nextBodyFn, responseHandlerFn)
+	}
+}
+
+// makeServerStream starts server-side streaming rpc
+func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream,
+	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+
+	req, err := nextBodyFn()
+	// Use nil body if no initial request input is empty, since request
+	// is mandatory in server streaming rpc.
+	if err != nil && err != io.EOF {
+		return err
+	}
+	if err == nil {
+		// Verify there is no second request.
+		if _, err = nextBodyFn(); err == nil {
+			return fmt.Errorf("Request data contains more than 1 message for server-streaming RPC")
+		} else if err != io.EOF {
+			return err
+		}
+	}
+
+	if err = sendStreamMessage(ctx, stream, req); err != nil {
+		return err
+	}
+
+	for err == nil {
+		var resBody []byte
+		if resBody, err = receiveStreamMessage(ctx, stream); err != nil {
+			break
+		}
+
+		err = responseHandlerFn(resBody)
+	}
+
+	if err == io.EOF {
+		return nil
+	}
+	return err
+}
+
+// makeClientStream starts client-side streaming rpc
+func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream,
+	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+
+	var err error
+	for err == nil {
+		var reqBody []byte
+		reqBody, err = nextBodyFn()
+		if err != nil {
+			break
+		}
+		err = sendStreamMessage(ctx, stream, reqBody)
+	}
+
+	if err == io.EOF {
+		err = closeSendStream(ctx, stream)
+	}
+	if err != nil {
+		return err
+	}
+
+	res, err := receiveStreamMessage(ctx, stream)
+	if err != nil {
+		return err
+	}
+	return responseHandlerFn(res)
+}
+
+// makeBidiStream starts bi-directional streaming rpc
+func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream,
+	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+
+	var wg sync.WaitGroup
+	var sendErr error
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	wg.Add(1)
+	// Start go routine to concurrently send stream messages.
+	go func() {
+		defer wg.Done()
+
+		var err error
+		for err == nil {
+			var reqBody []byte
+			reqBody, err = nextBodyFn()
+			if err == io.EOF {
+				err = closeSendStream(ctx, stream)
+				break
+			}
+			if err != nil {
+				// Cancel the context to unblock the routine waiting on receiving
+				// stream messages.
+				cancel()
+				break
+			}
+
+			err = sendStreamMessage(ctx, stream, reqBody)
+		}
+
+		if err != nil {
+			sendErr = err
+		}
+	}()
+
+	var receiveErr error
+	for receiveErr == nil {
+		var resBody []byte
+		resBody, receiveErr = receiveStreamMessage(ctx, stream)
+		if receiveErr != nil {
+			break
+		}
+
+		receiveErr = responseHandlerFn(resBody)
+	}
+
+	cancel()
+	wg.Wait()
+
+	if sendErr != nil && sendErr != io.EOF {
+		return sendErr
+	}
+
+	if receiveErr != nil && receiveErr != io.EOF {
+		return receiveErr
+	}
+
+	return nil
+}
+
+// sendStreamMessage sends the stream message using message body provided
+func sendStreamMessage(ctx context.Context, stream *yarpctransport.ClientStream, msgBody []byte) error {
+	// TODO: print the stream message being sent on STDOUT to inform user about
+	// request message dispatch.
+	req := &yarpctransport.StreamMessage{Body: ioutil.NopCloser(bytes.NewReader(msgBody))}
+	if err := stream.SendMessage(ctx, req); err != nil {
+		if err == io.EOF {
+			return err
+		}
+		return fmt.Errorf("Failed while sending stream request: %v", err)
+	}
+	return nil
+}
+
+// receiveStreamMessage receives and returns the message from the given stream
+func receiveStreamMessage(ctx context.Context, stream *yarpctransport.ClientStream) ([]byte, error) {
+	msg, err := stream.ReceiveMessage(ctx)
+	if err == io.EOF {
+		return nil, err
+	}
+	if err != nil {
+		return nil, fmt.Errorf("Failed while receiving stream response: %v", err)
+	}
+
+	bytes, err := ioutil.ReadAll(msg.Body)
+	if err != nil {
+		return nil, fmt.Errorf("Failed while reading stream response: %v", err)
+	}
+	return bytes, err
+}
+
+// closeSendStream closes the stream from the client side while
+// stream can continue to receive messages from server
+func closeSendStream(ctx context.Context, stream *yarpctransport.ClientStream) error {
+	// YARPC stream.Close method internally invokes closeSend on gRPC clientStream.
+	if err := stream.Close(ctx); err != nil {
+		return fmt.Errorf("Failed to close send stream: %v", err)
+	}
+	return nil
+}
diff --git a/main.go b/main.go
index fc86234..6df616c 100644
--- a/main.go
+++ b/main.go
@@ -21,11 +21,11 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"log"
 	"os"
 	"regexp"
@@ -340,30 +340,17 @@ func runWithOptions(opts Options, out output, logger *zap.Logger) {
 		out.Fatalf("Failed while parsing options: %v\n", err)
 	}
 
-	reqInput, err := ioutil.ReadAll(reqReader)
-	if err != nil {
-		out.Fatalf("Failed while reading body input: %v\n", err)
-	}
-
-	// req is the transport.Request that will be used to make a call.
-	req, err := serializer.Request(reqInput)
-	if err != nil {
-		out.Fatalf("Failed while parsing request input: %v\n", err)
-	}
-	req, err = prepareRequest(req, headers, opts)
-	if err != nil {
-		out.Fatalf("Failed while preparing the request: %v\n", err)
-	}
-
-	// Only make the request if the user hasn't specified 0 warmup.
-	if !(opts.BOpts.enabled() && opts.BOpts.WarmupRequests == 0) {
-		makeInitialRequest(out, transport, serializer, req)
-	}
-
-	runBenchmark(out, logger, opts, resolved, benchmarkMethod{
+	handler := requestHandler{
+		out:        out,
+		logger:     logger,
+		opts:       opts,
+		transport:  transport,
+		resolved:   resolved,
 		serializer: serializer,
-		req:        req,
-	})
+		body:       reqReader,
+		headers:    headers,
+	}
+	handler.handle()
 }
 
 func createJaegerTracer(opts Options, out output) (opentracing.Tracer, io.Closer) {
@@ -466,6 +453,11 @@ func makeRequestWithTracePriority(t transport.Transport, request *transport.Requ
 	ctx, cancel := tchannel.NewContext(request.Timeout)
 	defer cancel()
 
+	ctx = makeContextWithTrace(ctx, t, request, trace)
+	return t.Call(ctx, request)
+}
+
+func makeContextWithTrace(ctx context.Context, t transport.Transport, request *transport.Request, trace uint16) context.Context {
 	if tracer := t.Tracer(); tracer != nil {
 		span := tracer.StartSpan(request.Method)
 		opentracing_ext.SamplingPriority.Set(span, trace)
@@ -474,8 +466,7 @@ func makeRequestWithTracePriority(t transport.Transport, request *transport.Requ
 		}
 		ctx = opentracing.ContextWithSpan(ctx, span)
 	}
-
-	return t.Call(ctx, request)
+	return ctx
 }
 
 func makeInitialRequest(out output, transport transport.Transport, serializer encoding.Serializer, req *transport.Request) {


Test patch:
diff --git a/integration_test.go b/integration_test.go
index 823c63b..eabdbf5 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -25,9 +25,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"testing"
 	"time"
 
@@ -359,40 +361,376 @@ func setupYARPCServer(t *testing.T, inbound ytransport.Inbound, opts ...ythrift.
 	return dispatcher
 }
 
-type simpleService struct{}
+type simpleService struct {
+	expectedInput []simple.Foo // messages expected from the client
+
+	returnOutput []simple.Foo // messages to be streamed back to client
+	returnError  error
+}
 
 func (s *simpleService) Baz(c context.Context, in *simple.Foo) (*simple.Foo, error) {
 	if in.Test > 0 {
 		return in, nil
 	}
+
 	return nil, fmt.Errorf("negative input")
 }
 
-func (*simpleService) BidiStream(simple.Bar_BidiStreamServer) error {
-	return nil
+func (s *simpleService) ClientStream(stream simple.Bar_ClientStreamServer) error {
+	idx := 0
+	for idx < len(s.expectedInput) {
+		msg, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+
+		if !reflect.DeepEqual(*msg, s.expectedInput[idx]) {
+			return fmt.Errorf("unexpected input found at index: %d", idx)
+		}
+
+		idx++
+	}
+
+	if len(s.returnOutput) > 0 {
+		return stream.SendAndClose(&s.returnOutput[0])
+	}
+
+	return s.returnError
 }
 
-func (*simpleService) ClientStream(simple.Bar_ClientStreamServer) error {
-	return nil
+func (s *simpleService) ServerStream(req *simple.Foo, stream simple.Bar_ServerStreamServer) error {
+	if len(s.expectedInput) > 0 && !reflect.DeepEqual(*req, s.expectedInput[0]) {
+		return fmt.Errorf("unexpected input found")
+	}
+
+	for _, req := range s.returnOutput {
+		if err := stream.Send(&req); err != nil {
+			return err
+		}
+	}
+
+	return s.returnError
 }
 
-func (*simpleService) ServerStream(*simple.Foo, simple.Bar_ServerStreamServer) error {
-	return nil
+func (s *simpleService) BidiStream(stream simple.Bar_BidiStreamServer) error {
+	idx := 0
+	for idx < len(s.expectedInput) {
+		msg, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+
+		if !reflect.DeepEqual(*msg, s.expectedInput[idx]) {
+			return fmt.Errorf("unexpected input found at index: %d", idx)
+		}
+
+		idx++
+	}
+
+	for _, req := range s.returnOutput {
+		if err := stream.Send(&req); err != nil {
+			return err
+		}
+	}
+
+	return s.returnError
 }
 
-func setupGRPCServer(t *testing.T) (net.Addr, *grpc.Server) {
+func setupGRPCServer(t *testing.T, svc *simpleService) (net.Addr, *grpc.Server) {
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 
 	s := grpc.NewServer()
-	simple.RegisterBarServer(s, &simpleService{})
+	simple.RegisterBarServer(s, svc)
 	reflection.Register(s)
 	go s.Serve(ln)
 	return ln.Addr(), s
 }
 
+func TestGRPCStream(t *testing.T) {
+	tests := []struct {
+		desc    string
+		opts    Options
+		wantRes string
+		wantErr string
+
+		returnError   error
+		returnOutput  []simple.Foo
+		expectedInput []simple.Foo
+	}{
+		{
+			desc: "client streaming",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ClientStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{"test":1}{"test":2}{"test":-1} {"test":10}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes: `{
+  "test": 4
+}`,
+			expectedInput: []simple.Foo{{Test: 1}, {Test: 2}, {Test: -1}, {Test: 10}},
+			returnOutput:  []simple.Foo{{Test: 4}},
+		},
+		{
+			desc: "client streaming with YAML input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ClientStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON: `test: 1
+---
+test: 2
+---
+test: -1
+---
+test: 10
+---`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes: `{
+  "test": 4
+}`,
+			expectedInput: []simple.Foo{{Test: 1}, {Test: 2}, {Test: -1}, {Test: 10}},
+			returnOutput:  []simple.Foo{{Test: 4}},
+		},
+		{
+			desc: "client streaming with empty input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ClientStream",
+					Timeout:           timeMillisFlag(time.Second),
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes:      `{}`,
+			returnOutput: []simple.Foo{{}},
+		},
+		{
+			desc: "sever streaming with multiple input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ServerStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{}{}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantErr: "Request data contains more than 1 message for server-streaming RPC\n",
+		},
+		{
+			desc: "bidi streaming with immidiate error",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantErr:       "Failed while receiving stream response: code:unknown message:test error\n",
+			expectedInput: []simple.Foo{{}},
+			returnError:   errors.New("test error"),
+		},
+		{
+			desc: "bidi streaming with EOF",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+		},
+		{
+			desc: "error on stream call due to timeout",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Nanosecond),
+					RequestJSON:       `{}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantErr: "Failed while making stream call: code:unavailable message:roundrobin peer list timed out waiting for peer: context deadline exceeded\n",
+		},
+		{
+			desc: "server streaming",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ServerStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{"test":2}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes: `{
+  "test": 1
+}
+
+{
+  "test": 2
+}`,
+			expectedInput: []simple.Foo{{Test: 2}},
+			returnOutput:  []simple.Foo{{Test: 1}, {Test: 2}},
+		},
+		{
+			desc: "server streaming with YAML input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ServerStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `test: 2`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes: `{
+  "test": 1
+}
+
+{
+  "test": 2
+}`,
+			expectedInput: []simple.Foo{{Test: 2}},
+			returnOutput:  []simple.Foo{{Test: 1}, {Test: 2}},
+		},
+		{
+			desc: "bidirectional streaming",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{"test":250}{"test":1}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes: `{
+  "test": 350
+}
+
+{
+  "test": 101
+}`,
+			expectedInput: []simple.Foo{{Test: 250}, {Test: 1}},
+			returnOutput:  []simple.Foo{{Test: 350}, {Test: 101}},
+		},
+		{
+			desc: "bidirectional streaming with YAML input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON: `test: 250
+---
+test: 1
+---`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantRes: `{
+  "test": 350
+}
+
+{
+  "test": 101
+}`,
+			expectedInput: []simple.Foo{{Test: 250}, {Test: 1}},
+			returnOutput:  []simple.Foo{{Test: 350}, {Test: 101}},
+		},
+		{
+			desc: "client streaming with invalid input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ClientStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantErr: "Failed while reading stream input: unexpected EOF\n",
+		},
+		{
+			desc: "bidirectional streaming with invalid second input",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second),
+					RequestJSON:       `{"test":250}{"test_err": 1}`,
+				},
+				TOpts: TransportOptions{
+					ServiceName: "foo",
+				},
+			},
+			wantErr: "Failed while reading stream input: could not parse given request body as message of type \"Foo\": Message type Foo has no known field named test_err\n",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			addr, server := setupGRPCServer(t, &simpleService{
+				returnOutput:  tt.returnOutput,
+				expectedInput: tt.expectedInput,
+				returnError:   tt.returnError,
+			})
+			defer server.Stop()
+
+			tt.opts.TOpts.Peers = []string{"grpc://" + addr.String()}
+
+			gotOut, gotErr := runTestWithOpts(tt.opts)
+			assert.Equal(t, tt.wantErr, gotErr)
+			assert.Contains(t, gotOut, tt.wantRes)
+		})
+	}
+}
+
 func TestGRPCReflectionSource(t *testing.T) {
-	addr, server := setupGRPCServer(t)
+	addr, server := setupGRPCServer(t, &simpleService{})
 	defer server.GracefulStop()
 
 	tests := []struct {
@@ -411,7 +749,7 @@ func TestGRPCReflectionSource(t *testing.T) {
 				},
 				TOpts: TransportOptions{
 					ServiceName: "foo",
-					Peers:       []string{"grpc://" + addr.String()},
+					Peers:       []string{addr.String()},
 				},
 			},
 			wantRes: `"test": 1`,
diff --git a/main_test.go b/main_test.go
index 260de56..2583d3c 100644
--- a/main_test.go
+++ b/main_test.go
@@ -110,7 +110,7 @@ func TestRunWithOptions(t *testing.T) {
 					Peers:       []string{"1.1.1.1:1"},
 				},
 			},
-			errMsg: "while parsing request input",
+			errMsg: "Failed while serializing the input: yaml: line 1: did not find expected ',' or '}'",
 		},
 		{
 			desc: "Invalid host:port, fail to make request",
diff --git a/testdata/yarpc/integration/fooserver/server.go b/testdata/yarpc/integration/fooserver/server.go
index fe6de7d..0d16bb7 100644
--- a/testdata/yarpc/integration/fooserver/server.go
+++ b/testdata/yarpc/integration/fooserver/server.go
@@ -5,10 +5,11 @@ package fooserver
 
 import (
 	"context"
+
+	"github.com/yarpc/yab/testdata/yarpc/integration"
 	"go.uber.org/thriftrw/wire"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/encoding/thrift"
-	"github.com/yarpc/yab/testdata/yarpc/integration"
 )
 
 // Interface is the server-side interface for the Foo service.