Mojo UDP: add a client-side wrapper which handles the client<->server buffering

mojo/services/network/udp_socket_unittest.cc

Index: mojo/services/network/udp_socket_unittest.cc
diff --git a/mojo/services/network/udp_socket_unittest.cc b/mojo/services/network/udp_socket_unittest.cc
index 46c443a7ee7806d1f1cf1ae874ad3dd7c0f3580f..5aae6145278ade7745b34648a4c5d64daa85a85d 100644
--- a/mojo/services/network/udp_socket_unittest.cc
+++ b/mojo/services/network/udp_socket_unittest.cc
@@ -6,6 +6,7 @@
 #include "base/macros.h"
 #include "base/memory/scoped_ptr.h"
 #include "mojo/public/cpp/bindings/callback.h"
+#include "mojo/services/public/cpp/network/udp_socket_wrapper.h"
 #include "mojo/services/public/interfaces/network/network_service.mojom.h"
 #include "mojo/services/public/interfaces/network/udp_socket.mojom.h"
 #include "mojo/shell/shell_test_helper.h"
@@ -38,24 +39,6 @@ Array<uint8_t> CreateTestMessage(uint8_t initial, size_t size) {
   return array.Pass();
 }
 
-bool AreEqualArrays(const Array<uint8_t>& array_1,
-                    const Array<uint8_t>& array_2) {
-  if (array_1.is_null() != array_2.is_null())
-    return false;
-  else if (array_1.is_null())
-    return true;
-
-  if (array_1.size() != array_2.size())
-    return false;
-
-  for (size_t i = 0; i < array_1.size(); ++i) {
-    if (array_1[i] != array_2[i])
-      return false;
-  }
-
-  return true;
-}
-
 template <typename CallbackType>
 class TestCallbackBase {
  public:
@@ -80,7 +63,7 @@ class TestCallbackBase {
  protected:
   struct StateBase : public CallbackType::Runnable {
     StateBase() : test_callback_(nullptr) {}
-    virtual ~StateBase() {}
+    ~StateBase() override {}
 
     void set_test_callback(TestCallbackBase* test_callback) {
       test_callback_ = test_callback;
@@ -131,9 +114,9 @@ class TestCallback : public TestCallbackBase<Callback<void(NetworkErrorPtr)>> {
 
  private:
   struct State: public StateBase {
-    virtual ~State() {}
+    ~State() override {}
 
-    virtual void Run(NetworkErrorPtr result) const override {
+    void Run(NetworkErrorPtr result) const override {
       if (test_callback_) {
         TestCallback* callback = static_cast<TestCallback*>(test_callback_);
         callback->result_ = result.Pass();
@@ -158,10 +141,9 @@ class TestCallbackWithAddress
 
  private:
   struct State : public StateBase {
-    virtual ~State() {}
+    ~State() override {}
 
-    virtual void Run(NetworkErrorPtr result,
-                     NetAddressPtr net_address) const override {
+    void Run(NetworkErrorPtr result, NetAddressPtr net_address) const override {
       if (test_callback_) {
         TestCallbackWithAddress* callback =
             static_cast<TestCallbackWithAddress*>(test_callback_);
@@ -188,9 +170,9 @@ class TestCallbackWithUint32
 
  private:
   struct State : public StateBase {
-    virtual ~State() {}
+    ~State() override {}
 
-    virtual void Run(uint32_t result) const override {
+    void Run(uint32_t result) const override {
       if (test_callback_) {
         TestCallbackWithUint32* callback =
             static_cast<TestCallbackWithUint32*>(test_callback_);
@@ -203,12 +185,48 @@ class TestCallbackWithUint32
   uint32_t result_;
 };
 
+class TestReceiveCallback
+    : public TestCallbackBase<
+        Callback<void(NetworkErrorPtr, NetAddressPtr, Array<uint8_t>)>> {
+ public:
+  TestReceiveCallback() {
+    Initialize(new State());
+  }
+  ~TestReceiveCallback() {}
+
+  const NetworkErrorPtr& result() const { return result_; }
+  const NetAddressPtr& src_addr() const { return src_addr_; }
+  const Array<uint8_t>& data() const { return data_; }
+
+ private:
+  struct State : public StateBase {
+    ~State() override {}
+
+    void Run(NetworkErrorPtr result,
+             NetAddressPtr src_addr,
+             Array<uint8_t> data) const override {
+      if (test_callback_) {
+        TestReceiveCallback* callback =
+            static_cast<TestReceiveCallback*>(test_callback_);
+        callback->result_ = result.Pass();
+        callback->src_addr_ = src_addr.Pass();
+        callback->data_ = data.Pass();
+      }
+      NotifyRun();
+    }
+  };
+
+  NetworkErrorPtr result_;
+  NetAddressPtr src_addr_;
+  Array<uint8_t> data_;
+};
+
 class UDPSocketTest : public testing::Test {
  public:
   UDPSocketTest() {}
-  virtual ~UDPSocketTest() {}
+  ~UDPSocketTest() override {}
 
-  virtual void SetUp() override {
+  void SetUp() override {
     test_helper_.Init();
 
     test_helper_.application_manager()->ConnectToService(
@@ -230,16 +248,16 @@ class UDPSocketTest : public testing::Test {
 
     UDPSocketClientImpl() : run_loop_(nullptr), expected_receive_count_(0) {}
 
-    virtual ~UDPSocketClientImpl() {
+    ~UDPSocketClientImpl() override {
       while (!results_.empty()) {
         delete results_.front();
         results_.pop();
       }
     }
 
-    virtual void OnReceived(NetworkErrorPtr result,
-                            NetAddressPtr src_addr,
-                            Array<uint8_t> data) override {
+    void OnReceived(NetworkErrorPtr result,
+                    NetAddressPtr src_addr,
+                    Array<uint8_t> data) override {
       ReceiveResult* entry = new ReceiveResult();
       entry->result = result.Pass();
       entry->addr = src_addr.Pass();
@@ -356,6 +374,8 @@ TEST_F(UDPSocketTest, TestReadWrite) {
   ASSERT_EQ(net::OK, callback2.result()->code);
   ASSERT_NE(0u, callback2.net_address()->ipv4->port);
 
+  NetAddressPtr client_addr = callback2.net_address().Clone();
+
   const size_t kDatagramCount = 6;
   const size_t kDatagramSize = 255;
   udp_socket_->ReceiveMore(kDatagramCount);
@@ -376,9 +396,61 @@ TEST_F(UDPSocketTest, TestReadWrite) {
     GetReceiveResults()->pop();
 
     EXPECT_EQ(static_cast<int>(kDatagramSize), result->result->code);
-    EXPECT_TRUE(AreEqualArrays(
-        CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize),
-        result->data));
+    EXPECT_TRUE(result->addr.Equals(client_addr));
+    EXPECT_TRUE(result->data.Equals(
+        CreateTestMessage(static_cast<uint8_t>(i), kDatagramSize)));
+  }
+}
+
+TEST_F(UDPSocketTest, TestUDPSocketWrapper) {
+  UDPSocketWrapper udp_socket(udp_socket_.Pass(), 4, 4);
+
+  TestCallbackWithAddress callback1;
+  udp_socket.Bind(GetLocalHostWithAnyPort(), callback1.callback());
+  callback1.WaitForResult();
+  ASSERT_EQ(net::OK, callback1.result()->code);
+  ASSERT_NE(0u, callback1.net_address()->ipv4->port);
+
+  NetAddressPtr server_addr = callback1.net_address().Clone();
+
+  UDPSocketPtr raw_client_socket;
+  network_service_->CreateUDPSocket(GetProxy(&raw_client_socket));
+  UDPSocketWrapper client_socket(raw_client_socket.Pass(), 4, 4);
+
+  TestCallbackWithAddress callback2;
+  client_socket.Bind(GetLocalHostWithAnyPort(), callback2.callback());
+  callback2.WaitForResult();
+  ASSERT_EQ(net::OK, callback2.result()->code);
+  ASSERT_NE(0u, callback2.net_address()->ipv4->port);
+
+  NetAddressPtr client_addr = callback2.net_address().Clone();
+
+  const size_t kDatagramCount = 16;
+  const size_t kDatagramSize = 255;
+
+  for (size_t i = 1; i < kDatagramCount; ++i) {
+    scoped_ptr<TestCallback[]> send_callbacks(new TestCallback[i]);
+    scoped_ptr<TestReceiveCallback[]> receive_callbacks(
+        new TestReceiveCallback[i]);
+
+    for (size_t j = 0; j < i; ++j) {
+      client_socket.SendTo(
+          server_addr.Clone(),
+          CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize),
+          send_callbacks[j].callback());
+
+      udp_socket.ReceiveFrom(receive_callbacks[j].callback());
+    }
+
+    receive_callbacks[i - 1].WaitForResult();
+
+    for (size_t j = 0; j < i; ++j) {
+      EXPECT_EQ(static_cast<int>(kDatagramSize),
+                receive_callbacks[j].result()->code);
+      EXPECT_TRUE(receive_callbacks[j].src_addr().Equals(client_addr));
+      EXPECT_TRUE(receive_callbacks[j].data().Equals(
+          CreateTestMessage(static_cast<uint8_t>(j), kDatagramSize)));
+    }
   }
 }