Necessary changes in simple_switch_grpc tests

To support P4Runtime controller arbitration.
p4lang · Sep 9, 2017 · 595bb79 · 595bb79
1 parent 93fca47
commit 595bb79
Show file tree

Hide file tree

Showing 6 changed files with 94 additions and 2 deletions.
diff --git a/targets/simple_switch_grpc/tests/base_test.cpp b/targets/simple_switch_grpc/tests/base_test.cpp
@@ -20,6 +20,7 @@
 
 #include <grpc++/grpc++.h>
 
+#include <google/rpc/code.pb.h>
 #include <p4/tmp/p4config.grpc.pb.h>
 
 #include <fstream>
@@ -44,6 +45,29 @@ SimpleSwitchGrpcBaseTest::SimpleSwitchGrpcBaseTest(
   p4info = parse_p4info(p4info_proto_txt_path);
 }
 
+void
+SimpleSwitchGrpcBaseTest::SetUp() {
+  stream = p4runtime_stub->StreamChannel(&stream_context);
+  p4::StreamMessageRequest request;
+  auto arbitration = request.mutable_arbitration();
+  arbitration->set_device_id(device_id);
+  set_election_id(arbitration->mutable_election_id());
+  stream->Write(request);
+  p4::StreamMessageResponse response;
+  stream->Read(&response);
+  ASSERT_EQ(response.update_case(), p4::StreamMessageResponse::kArbitration);
+  ASSERT_EQ(response.arbitration().status().code(), ::google::rpc::Code::OK);
+}
+
+void
+SimpleSwitchGrpcBaseTest::TearDown() {
+  stream->WritesDone();
+  p4::StreamMessageResponse response;
+  while (stream->Read(&response)) { }
+  auto status = stream->Finish();
+  EXPECT_TRUE(status.ok());
+}
+
 void
 SimpleSwitchGrpcBaseTest::update_json(const char *json_path) {
   p4::SetForwardingPipelineConfigRequest request;
@@ -67,6 +91,20 @@ SimpleSwitchGrpcBaseTest::update_json(const char *json_path) {
   ASSERT_TRUE(status.ok());
 }
 
+void
+SimpleSwitchGrpcBaseTest::set_election_id(p4::Uint128 *election_id) const {
+  election_id->set_high(0);
+  election_id->set_low(1);
+}
+
+grpc::Status
+SimpleSwitchGrpcBaseTest::Write(ClientContext *context,
+                                p4::WriteRequest &request,
+                                p4::WriteResponse *response) {
+  set_election_id(request.mutable_election_id());
+  return p4runtime_stub->Write(context, request, response);
+}
+
 }  // namespace testing
 
 }  // namespace sswitch_grpc
diff --git a/targets/simple_switch_grpc/tests/base_test.h b/targets/simple_switch_grpc/tests/base_test.h
@@ -53,10 +53,25 @@ class SimpleSwitchGrpcBaseTest : public ::testing::Test {
  protected:
   explicit SimpleSwitchGrpcBaseTest(const char *p4info_proto_txt_path);
 
+  void SetUp() override;
+
+  void TearDown() override;
+
   void update_json(const char *json_path);
 
+  void set_election_id(p4::Uint128 *election_id) const;
+
+  // calls p4runtime_stub->Write, with the appropriate election_id
+  grpc::Status Write(ClientContext *context,
+                     p4::WriteRequest &request,
+                     p4::WriteResponse *response);
+
   std::shared_ptr<grpc::Channel> p4runtime_channel;
   std::unique_ptr<p4::P4Runtime::Stub> p4runtime_stub;
+  using ReaderWriter = ::grpc::ClientReaderWriter<p4::StreamMessageRequest,
+                                                  p4::StreamMessageResponse>;
+  ClientContext stream_context;
+  std::unique_ptr<ReaderWriter> stream{nullptr};
   p4::config::P4Info p4info{};
 };
 

diff --git a/targets/simple_switch_grpc/tests/example.cpp b/targets/simple_switch_grpc/tests/example.cpp
@@ -20,6 +20,7 @@
 
 #include <grpc++/grpc++.h>
 
+#include <google/rpc/code.pb.h>
 #include <p4/p4runtime.grpc.pb.h>
 #include <p4/tmp/p4config.grpc.pb.h>
 
@@ -57,6 +58,27 @@ test() {
 
   auto p4info = parse_p4info(test_proto_txt);
 
+  auto set_election_id = [](p4::Uint128 *election_id) {
+    election_id->set_high(0);
+    election_id->set_low(1);
+  };
+
+  // initial handshake: open bidirectional stream and advertise election
+  // id. This stream needs to stay open for the lifetime of the controller.
+  ClientContext stream_context;
+  auto stream = pi_stub_->StreamChannel(&stream_context);
+  {
+    p4::StreamMessageRequest request;
+    auto arbitration = request.mutable_arbitration();
+    arbitration->set_device_id(dev_id);
+    set_election_id(arbitration->mutable_election_id());
+    stream->Write(request);
+    p4::StreamMessageResponse response;
+    stream->Read(&response);
+    assert(response.update_case() == p4::StreamMessageResponse::kArbitration);
+    assert(response.arbitration().status().code() == ::google::rpc::Code::OK);
+  }
+
   {
     p4::SetForwardingPipelineConfigRequest request;
     request.set_action(
@@ -110,6 +132,7 @@ test() {
   // add entry
   {
     p4::WriteRequest request;
+    set_election_id(request.mutable_election_id());
     request.set_device_id(dev_id);
     auto update = request.add_updates();
     update->set_type(p4::Update_Type_INSERT);
@@ -147,6 +170,7 @@ test() {
   // remove entry
   {
     p4::WriteRequest request;
+    set_election_id(request.mutable_election_id());
     request.set_device_id(dev_id);
     auto update = request.add_updates();
     update->set_type(p4::Update_Type_DELETE);
@@ -164,6 +188,14 @@ test() {
     assert(rep.entities().size() == 0);
   }
 
+  {
+    stream->WritesDone();
+    p4::StreamMessageResponse response;
+    while (stream->Read(&response)) { }
+    auto status = stream->Finish();
+    assert(status.ok());
+  }
+
   return 0;
 }
 

diff --git a/targets/simple_switch_grpc/tests/test_basic.cpp b/targets/simple_switch_grpc/tests/test_basic.cpp
@@ -48,6 +48,7 @@ class SimpleSwitchGrpcTest_Basic : public SimpleSwitchGrpcBaseTest {
       : SimpleSwitchGrpcBaseTest(simple_router_proto) { }
 
   void SetUp() override {
+    SimpleSwitchGrpcBaseTest::SetUp();
     update_json(simple_router_json);
   }
 };
@@ -90,7 +91,7 @@ TEST_F(SimpleSwitchGrpcTest_Basic, Entries) {
     update->set_allocated_entity(&entity);
     ClientContext context;
     p4::WriteResponse rep;
-    auto status = p4runtime_stub->Write(&context, request, &rep);
+    auto status = Write(&context, request, &rep);
     EXPECT_TRUE(status.ok());
     update->release_entity();
   }
@@ -127,7 +128,7 @@ TEST_F(SimpleSwitchGrpcTest_Basic, Entries) {
     update->set_allocated_entity(&entity);
     ClientContext context;
     p4::WriteResponse rep;
-    auto status = p4runtime_stub->Write(&context, request, &rep);
+    auto status = Write(&context, request, &rep);
     EXPECT_TRUE(status.ok());
     update->release_entity();
   }

diff --git a/targets/simple_switch_grpc/tests/test_grpc_dp.cpp b/targets/simple_switch_grpc/tests/test_grpc_dp.cpp
@@ -49,6 +49,7 @@ class SimpleSwitchGrpcTest_GrpcDataplane : public SimpleSwitchGrpcBaseTest {
             dataplane_channel)) { }
 
   void SetUp() override {
+    SimpleSwitchGrpcBaseTest::SetUp();
     update_json(loopback_json);
   }
 

diff --git a/targets/simple_switch_grpc/tests/test_packet_io.cpp b/targets/simple_switch_grpc/tests/test_packet_io.cpp
@@ -44,6 +44,7 @@ class SimpleSwitchGrpcTest_PacketIO : public SimpleSwitchGrpcBaseTest {
       : SimpleSwitchGrpcBaseTest(loopback_proto) { }
 
   void SetUp() override {
+    SimpleSwitchGrpcBaseTest::SetUp();
     update_json(loopback_json);
   }
 };
@@ -55,8 +56,12 @@ TEST_F(SimpleSwitchGrpcTest_PacketIO, SendAndReceiveCPUPacket) {
   ClientContext context;
   auto stream = p4runtime_stub->StreamChannel(&context);
 
+  // create a new master stream
   auto arbitration = request.mutable_arbitration();
   arbitration->set_device_id(device_id);
+  auto election_id = arbitration->mutable_election_id();
+  election_id->set_high(1);
+  election_id->set_low(0);
   stream->Write(request);
 
   auto packet_out = request.mutable_packet();