1 files changed, 137 insertions, 0 deletions

diff --git a/test/cpp/end2end/mock_test.cc b/test/cpp/end2end/mock_test.cc
index 16c04032ab..b7968ce230 100644
--- a/test/cpp/end2end/mock_test.cc
+++ b/test/cpp/end2end/mock_test.cc
@@ -90,6 +90,55 @@ class FakeClient {
     EXPECT_TRUE(s.ok());
   }
 
+  void DoRequestStream() {
+    EchoRequest request;
+    EchoResponse response;
+
+    ClientContext context;
+    grpc::string msg("hello");
+    grpc::string exp(msg);
+
+    std::unique_ptr<ClientWriterInterface<EchoRequest>>
+        cstream = stub_->RequestStream(&context, &response);
+
+    request.set_message(msg);
+    EXPECT_TRUE(cstream->Write(request));
+
+    msg = ", world";
+    request.set_message(msg);
+    exp.append(msg);
+    EXPECT_TRUE(cstream->Write(request));
+
+    cstream->WritesDone();
+    Status s = cstream->Finish();
+
+    EXPECT_EQ(exp, response.message());
+    EXPECT_TRUE(s.ok());
+  }
+
+  void DoResponseStream() {
+    EchoRequest request;
+    EchoResponse response;
+    request.set_message("hello world");
+
+    ClientContext context;
+    std::unique_ptr<ClientReaderInterface<EchoResponse>>
+        cstream = stub_->ResponseStream(&context, request);
+
+    grpc::string exp = "";
+    EXPECT_TRUE(cstream->Read(&response));
+    exp.append(response.message() + " ");
+
+    EXPECT_TRUE(cstream->Read(&response));
+    exp.append(response.message());
+
+    EXPECT_FALSE(cstream->Read(&response));
+    EXPECT_EQ(request.message(), exp);
+
+    Status s = cstream->Finish();
+    EXPECT_TRUE(s.ok());
+  }
+
   void DoBidiStream() {
     EchoRequest request;
     EchoResponse response;
@@ -135,6 +184,31 @@ class TestServiceImpl : public EchoTestService::Service {
     return Status::OK;
   }
 
+  Status RequestStream(ServerContext* context,
+                       ServerReader<EchoRequest>* reader,
+                       EchoResponse* response) {
+    EchoRequest request;
+    grpc::string resp("");
+    while (reader->Read(&request)) {
+      gpr_log(GPR_INFO, "recv msg %s", request.message().c_str());
+      resp.append(request.message());
+    }
+    response->set_message(resp);
+    return Status::OK;
+  }
+
+  Status ResponseStream(ServerContext* context,
+                        const EchoRequest* request,
+                        ServerWriter<EchoResponse>* writer) {
+    EchoResponse response;
+    vector<grpc::string> tokens = split(request->message());
+    for (grpc::string token : tokens) {
+      response.set_message(token);
+      writer->Write(response);
+    }
+    return Status::OK;
+  }
+
   Status BidiStream(
       ServerContext* context,
       ServerReaderWriter<EchoResponse, EchoRequest>* stream) override {
@@ -147,6 +221,26 @@ class TestServiceImpl : public EchoTestService::Service {
     }
     return Status::OK;
   }
+ private:
+  const vector<grpc::string> split(const grpc::string& input) {
+    grpc::string buff("");
+    vector<grpc::string> result;
+
+    for (auto n:input) {
+      if (n != ' ') {
+        buff += n;
+        continue;
+      }
+      if (buff == "")
+        continue;
+      result.push_back(buff);
+      buff = "";
+    }
+    if (buff != "")
+      result.push_back(buff);
+
+    return result;
+  }
 };
 
 class MockTest : public ::testing::Test {
@@ -191,6 +285,49 @@ TEST_F(MockTest, SimpleRpc) {
   client.DoEcho();
 }
 
+TEST_F(MockTest, ClientStream) {
+  ResetStub();
+  FakeClient client(stub_.get());
+  client.DoRequestStream();
+
+  MockEchoTestServiceStub stub;
+  auto w = new MockClientWriter<EchoRequest>();
+  EchoResponse resp;
+  resp.set_message("hello, world");
+
+  EXPECT_CALL(*w, Write(_, _)).Times(2).WillRepeatedly(Return(true));
+  EXPECT_CALL(*w, WritesDone());
+  EXPECT_CALL(*w, Finish()).WillOnce(Return(Status::OK));
+
+  EXPECT_CALL(stub, RequestStreamRaw(_, _)).WillOnce(DoAll(SetArgPointee<1>(resp), Return(w)));
+  client.ResetStub(&stub);
+  client.DoRequestStream();
+}
+
+TEST_F(MockTest, ServerStream) {
+  ResetStub();
+  FakeClient client(stub_.get());
+  client.DoResponseStream();
+
+  MockEchoTestServiceStub stub;
+  auto r = new MockClientReader<EchoResponse>();
+  EchoResponse resp1;
+  resp1.set_message("hello");
+  EchoResponse resp2;
+  resp2.set_message("world");
+
+  EXPECT_CALL(*r, Read(_)).
+      WillOnce(DoAll(SetArgPointee<0>(resp1), Return(true))).
+      WillOnce(DoAll(SetArgPointee<0>(resp2), Return(true))).
+      WillOnce(Return(false));
+  EXPECT_CALL(*r, Finish()).WillOnce(Return(Status::OK));
+
+  EXPECT_CALL(stub, ResponseStreamRaw(_, _)).WillOnce(Return(r));
+
+  client.ResetStub(&stub);
+  client.DoResponseStream();
+}
+
 ACTION_P(copy, msg) {
   arg0->set_message(msg->message());
 }