Refactor read/write mock (#315)

bismuth01 · web-flow · commit 39ebf82b10c8 · 2025-07-04T14:00:43.000+05:00
diff --git a/test/libp2p/basic/message_read_writer_test.cpp b/test/libp2p/basic/message_read_writer_test.cpp
@@ -9,6 +9,8 @@
 #include <gtest/gtest.h>
 #include <libp2p/multi/uvarint.hpp>
 #include "mock/libp2p/connection/layer_connection_mock.hpp"
+#include "testutil/expect_read.hpp"
+#include "testutil/expect_write.hpp"
 #include "testutil/gmock_actions.hpp"
 
 using namespace libp2p;
@@ -41,17 +43,10 @@ class MessageReadWriterTest : public testing::Test {
   bool operation_completed_ = false;
 };
 
-ACTION_P(ReadPut, buf) {
-  ASSERT_GE(arg0.size(), buf.size());
-  std::copy(buf.begin(), buf.end(), arg0.begin());
-  arg2(buf.size());
-}
-
 TEST_F(MessageReadWriterTest, Read) {
-  EXPECT_CALL(*conn_mock_, read(_, 1, _))
-      .WillOnce(ReadPut(len_varint_.toBytes()));
-  EXPECT_CALL(*conn_mock_, read(_, kMsgLength, _))
-      .WillOnce(ReadPut(msg_bytes_));
+  EXPECT_CALL_READ(*conn_mock_)
+      .WILL_READ(len_varint_.toBytes())
+      .WILL_READ(msg_bytes_);
 
   msg_rw_->read([this](auto &&res) {
     ASSERT_TRUE(res);
@@ -62,18 +57,8 @@ TEST_F(MessageReadWriterTest, Read) {
   ASSERT_TRUE(operation_completed_);
 }
 
-ACTION_P2(CheckWrite, buf, varint) {
-  ASSERT_EQ(arg0.size(), buf.size());
-
-  for (auto i = 0u; i < varint.size(); ++i) {
-    ASSERT_EQ(arg0[i], varint.toBytes()[i]);
-  }
-  arg2(buf.size());
-}
-
 TEST_F(MessageReadWriterTest, Write) {
-  EXPECT_CALL(*conn_mock_, writeSome(_, kMsgLength + 1, _))
-      .WillOnce(CheckWrite(msg_with_varint_bytes_, len_varint_));
+  EXPECT_CALL_WRITE(*conn_mock_).WILL_WRITE(msg_with_varint_bytes_);
 
   msg_rw_->write(msg_bytes_, [this](auto &&res) {
     ASSERT_TRUE(res);
diff --git a/test/libp2p/connection/security_conn/plaintext_connection_test.cpp b/test/libp2p/connection/security_conn/plaintext_connection_test.cpp
@@ -11,6 +11,8 @@
 #include <qtils/test/outcome.hpp>
 #include "mock/libp2p/connection/layer_connection_mock.hpp"
 #include "mock/libp2p/crypto/key_marshaller_mock.hpp"
+#include "testutil/expect_read.hpp"
+#include "testutil/expect_write.hpp"
 #include "testutil/gmock_actions.hpp"
 
 using namespace libp2p::connection;
@@ -120,63 +122,29 @@ TEST_F(PlaintextConnectionTest, RemoteMultiaddr) {
  */
 TEST_F(PlaintextConnectionTest, Read) {
   const int size = 100;
-  EXPECT_CALL(*connection_, read(_, _, _)).WillOnce(AsioSuccess(size));
+  EXPECT_CALL_READ(*connection_).WILL_READ_SIZE(size);
   auto buf = std::make_shared<std::vector<uint8_t>>(size, 0);
   secure_connection_->read(*buf, size, [size, buf](auto &&res) {
     ASSERT_OUTCOME_SUCCESS(res);
     ASSERT_EQ(res.value(), size);
   });
 }
 
-/**
- * @given plaintext secure connection
- * @when invoking readSome method of the connection
- * @then method behaves as expected
- */
-TEST_F(PlaintextConnectionTest, ReadSome) {
-  const int size = 100;
-  const int smaller = 50;
-  EXPECT_CALL(*connection_, readSome(_, _, _))
-      .WillOnce(AsioSuccess(smaller /* less than 100 */));
-  auto buf = std::make_shared<std::vector<uint8_t>>(size, 0);
-  secure_connection_->readSome(*buf, smaller, [smaller, buf](auto &&res) {
-    ASSERT_OUTCOME_SUCCESS(res);
-    ASSERT_EQ(res.value(), smaller);
-  });
-}
-
 /**
  * @given plaintext secure connection
  * @when invoking write method of the connection
  * @then method behaves as expected
  */
 TEST_F(PlaintextConnectionTest, Write) {
   const int size = 100;
-  EXPECT_CALL(*connection_, writeSome(_, _, _)).WillOnce(AsioSuccess(size));
+  EXPECT_CALL_WRITE(*connection_).WILL_WRITE_SIZE(size);
   auto buf = std::make_shared<std::vector<uint8_t>>(size, 0);
   libp2p::writeReturnSize(secure_connection_, *buf, [size, buf](auto &&res) {
     ASSERT_OUTCOME_SUCCESS(res);
     ASSERT_EQ(res.value(), size);
   });
 }
 
-/**
- * @given plaintext secure connection
- * @when invoking writeSome method of the connection
- * @then method behaves as expected
- */
-TEST_F(PlaintextConnectionTest, WriteSome) {
-  const int size = 100;
-  const int smaller = 50;
-  EXPECT_CALL(*connection_, writeSome(_, _, _))
-      .WillOnce(AsioSuccess(smaller /* less than 100 */));
-  auto buf = std::make_shared<std::vector<uint8_t>>(size, 0);
-  secure_connection_->writeSome(*buf, smaller, [smaller, buf](auto &&res) {
-    ASSERT_OUTCOME_SUCCESS(res);
-    ASSERT_EQ(res.value(), smaller);
-  });
-}
-
 /**
  * @given plaintext secure connection
  * @when invoking isClosed method of the connection
diff --git a/test/libp2p/protocol/echo_test.cpp b/test/libp2p/protocol/echo_test.cpp
@@ -6,8 +6,11 @@
 
 #include <gtest/gtest.h>
 #include <libp2p/protocol/echo.hpp>
+#include <qtils/bytestr.hpp>
 #include <qtils/test/outcome.hpp>
 #include "mock/libp2p/connection/stream_mock.hpp"
+#include "testutil/expect_read.hpp"
+#include "testutil/expect_write.hpp"
 #include "testutil/gmock_actions.hpp"
 #include "testutil/prepare_loggers.hpp"
 
@@ -17,30 +20,8 @@ using ::testing::_;
 using ::testing::Return;
 using std::string_literals::operator""s;
 
-// set what we read
-ACTION_P(SetReadMsg, msg) {
-  std::copy(msg.begin(), msg.end(), arg0.begin());
-  arg2(msg.size());
-}
-
-// assert that we write the same as we read
-ACTION_P(WriteMsgAssertEqual, msg) {
-  std::string sub;
-
-  if (*arg0.begin() == 0) {
-    // EOF
-    return;
-  }
-
-  auto begin = arg0.begin();
-  auto end = arg0.begin();
-
-  std::advance(end, msg.size());
-  std::copy(begin, end, std::back_inserter(sub));
-  EXPECT_EQ(sub, msg);
-
-  arg2(msg.size());
-}
+const auto msg = "hello"s;
+const auto msg_bytes = qtils::str2byte(msg);
 
 /**
  * @given Stream
@@ -52,7 +33,6 @@ TEST(EchoTest, Server) {
 
   Echo echo;
   auto stream = std::make_shared<connection::StreamMock>();
-  auto msg = "hello"s;
 
   EXPECT_CALL(*stream, isClosedForRead())
       .WillOnce(Return(false))
@@ -63,8 +43,8 @@ TEST(EchoTest, Server) {
   EXPECT_CALL(*stream, close(_))
       .WillOnce(Arg0CallbackWithArg(outcome::success()));
 
-  EXPECT_CALL(*stream, readSome(_, _, _)).WillOnce(SetReadMsg(msg));
-  EXPECT_CALL(*stream, writeSome(_, _, _)).WillOnce(WriteMsgAssertEqual(msg));
+  EXPECT_CALL_READ(*stream).WILL_READ(msg_bytes);
+  EXPECT_CALL_WRITE(*stream).WILL_WRITE(msg_bytes);
 
   echo.handle(StreamAndProtocol{stream, {}});
 }
@@ -74,15 +54,14 @@ TEST(EchoTest, Server) {
  * @when client writes string "hello" to the Stream
  * @then client reads back the same string
  */
-TEST(EchoTest, DISABLED_Client) {
+TEST(EchoTest, Client) {
   Echo echo;
   auto stream = std::make_shared<connection::StreamMock>();
-  auto msg = "hello"s;
 
   EXPECT_CALL(*stream, isClosedForWrite()).WillOnce(Return(false));
 
-  EXPECT_CALL(*stream, writeSome(_, _, _)).WillOnce(WriteMsgAssertEqual(msg));
-  EXPECT_CALL(*stream, readSome(_, _, _)).WillOnce(WriteMsgAssertEqual(msg));
+  EXPECT_CALL_READ(*stream).WILL_READ(msg_bytes).WILL_READ_ERROR();
+  EXPECT_CALL_WRITE(*stream).WILL_WRITE(msg_bytes);
 
   bool executed = false;
 
diff --git a/test/libp2p/protocol/identify_delta_test.cpp b/test/libp2p/protocol/identify_delta_test.cpp
@@ -19,6 +19,7 @@
 #include "mock/libp2p/network/connection_manager_mock.hpp"
 #include "mock/libp2p/peer/peer_repository_mock.hpp"
 #include "mock/libp2p/peer/protocol_repository_mock.hpp"
+#include "testutil/expect_read.hpp"
 #include "testutil/gmock_actions.hpp"
 #include "testutil/prepare_loggers.hpp"
 
@@ -154,12 +155,9 @@ ACTION_P(ReadPut, buf) {
  */
 TEST_F(IdentifyDeltaTest, Receive) {
   // handle
-  EXPECT_CALL(*stream_, read(_, 1, _))
-      .WillOnce(ReadPut(std::span(msg_added_rm_protos_bytes_.data(), 1)));
-  EXPECT_CALL(*stream_, read(_, added_rm_proto_len_.toUInt64(), _))
-      .WillOnce(ReadPut(std::span(
-          msg_added_rm_protos_bytes_.data() + added_proto_len_.size(),
-          msg_added_rm_protos_bytes_.size() - added_proto_len_.size())));
+  EXPECT_CALL_READ(*stream_)
+      .WILL_READ(BytesOut(msg_added_rm_protos_bytes_).first(1))
+      .WILL_READ(BytesOut(msg_added_rm_protos_bytes_).subspan(1));
 
   // deltaReceived
   EXPECT_CALL(*stream_, remotePeerId())
diff --git a/test/libp2p/protocol/identify_test.cpp b/test/libp2p/protocol/identify_test.cpp
@@ -24,6 +24,8 @@
 #include "mock/libp2p/peer/key_repository_mock.hpp"
 #include "mock/libp2p/peer/peer_repository_mock.hpp"
 #include "mock/libp2p/peer/protocol_repository_mock.hpp"
+#include "testutil/expect_read.hpp"
+#include "testutil/expect_write.hpp"
 #include "testutil/prepare_loggers.hpp"
 
 using namespace libp2p;
@@ -184,10 +186,7 @@ TEST_F(IdentifyTest, Send) {
   EXPECT_CALL(host_, getLibp2pClientVersion()).WillOnce(Return(kClientVersion));
 
   // handle Identify request and check it
-  EXPECT_CALL(*stream_, writeSome(_, _, _))
-      .WillOnce(Success(
-          BytesIn(identify_pb_msg_bytes_.data(), identify_pb_msg_bytes_.size()),
-          outcome::success(identify_pb_msg_bytes_.size())));
+  EXPECT_CALL_WRITE(*stream_).WILL_WRITE(identify_pb_msg_bytes_);
 
   identify_->handle(StreamAndProtocol{stream_, {}});
 }
@@ -215,12 +214,9 @@ TEST_F(IdentifyTest, Receive) {
               newStream(kRemotePeerInfo, StreamProtocols{kIdentifyProto}, _))
       .WillOnce(InvokeArgument<2>(StreamAndProtocol{stream_, kIdentifyProto}));
 
-  EXPECT_CALL(*stream_, read(_, 1, _))
-      .WillOnce(ReadPut(std::span(identify_pb_msg_bytes_.data(), 1)));
-  EXPECT_CALL(*stream_, read(_, pb_msg_len_varint_->toUInt64(), _))
-      .WillOnce(ReadPut(std::span(
-          identify_pb_msg_bytes_.data() + pb_msg_len_varint_->size(),
-          identify_pb_msg_bytes_.size() - pb_msg_len_varint_->size())));
+  EXPECT_CALL_READ(*stream_)
+      .WILL_READ(BytesOut(identify_pb_msg_bytes_).first(1))
+      .WILL_READ(BytesOut(identify_pb_msg_bytes_).subspan(1));
 
   EXPECT_CALL(*stream_, remotePeerId())
       .Times(2)
diff --git a/test/libp2p/protocol/ping_test.cpp b/test/libp2p/protocol/ping_test.cpp
@@ -22,6 +22,8 @@
 #include "mock/libp2p/crypto/random_generator_mock.hpp"
 #include "mock/libp2p/host/host_mock.hpp"
 #include "mock/libp2p/peer/peer_repository_mock.hpp"
+#include "testutil/expect_read.hpp"
+#include "testutil/expect_write.hpp"
 
 using libp2p::basic::Scheduler;
 using libp2p::basic::SchedulerMock;
@@ -97,18 +99,9 @@ ACTION_P(ReadPut, buf) {
  * writes it back
  */
 TEST_F(PingTest, PingServer) {
-  EXPECT_CALL(*stream_, read(_, kPingMsgSize, _))
-      .WillOnce(ReadPut(buffer_))
-      .WillOnce(  // no second read
-          InvokeArgument<2>(std::error_code{}));
-
-  auto if_eq_buf = [&](BytesIn actual) {
-    auto expected = BytesIn(buffer_);
-    return std::equal(
-        actual.begin(), actual.end(), expected.begin(), expected.end());
-  };
-  EXPECT_CALL(*stream_, writeSome(Truly(if_eq_buf), kPingMsgSize, _))
-      .WillOnce(InvokeArgument<2>(buffer_.size()));
+  EXPECT_CALL_READ(*stream_).WILL_READ(buffer_).WILL_READ_ERROR();
+
+  EXPECT_CALL_WRITE(*stream_).WILL_WRITE(buffer_);
 
   EXPECT_CALL(*stream_, isClosedForWrite()).WillOnce(Return(false));
   EXPECT_CALL(*stream_, isClosedForRead())
@@ -135,17 +128,9 @@ TEST_F(PingTest, PingClient) {
   EXPECT_CALL(*rand_gen_, randomBytes(kPingMsgSize))
       .Times(2)
       .WillRepeatedly(Return(buffer_));
-  auto if_eq_buf = [&](BytesIn actual) {
-    auto expected = BytesIn(buffer_);
-    return std::equal(
-        actual.begin(), actual.end(), expected.begin(), expected.end());
-  };
-  EXPECT_CALL(*stream_, writeSome(Truly(if_eq_buf), kPingMsgSize, _))
-      .WillOnce(InvokeArgument<2>(buffer_.size()))
-      .WillOnce(  // no second write
-          InvokeArgument<2>(
-              make_error_code(boost::asio::error::invalid_argument)));
-  EXPECT_CALL(*stream_, read(_, kPingMsgSize, _)).WillOnce(ReadPut(buffer_));
+
+  EXPECT_CALL_WRITE(*stream_).WILL_WRITE(buffer_).WILL_WRITE_ERROR();
+  EXPECT_CALL_READ(*stream_).WILL_READ(buffer_);
 
   EXPECT_CALL(*stream_, isClosedForWrite())
       .Times(2)
@@ -175,12 +160,7 @@ TEST_F(PingTest, PingClientTimeoutExpired) {
       .WillOnce(InvokeArgument<2>(StreamAndProtocol{stream_, kPingProto}));
 
   EXPECT_CALL(*rand_gen_, randomBytes(kPingMsgSize)).WillOnce(Return(buffer_));
-  auto if_eq_buf = [&](BytesIn actual) {
-    auto expected = BytesIn(buffer_);
-    return std::equal(
-        actual.begin(), actual.end(), expected.begin(), expected.end());
-  };
-  EXPECT_CALL(*stream_, writeSome(Truly(if_eq_buf), kPingMsgSize, _));
+  EXPECT_CALL_WRITE(*stream_).WILL_WRITE(buffer_);
 
   EXPECT_CALL(*stream_, isClosedForWrite()).WillOnce(Return(false));
 
diff --git a/test/libp2p/security/plaintext_adaptor_test.cpp b/test/libp2p/security/plaintext_adaptor_test.cpp
@@ -44,9 +44,6 @@ class PlaintextAdaptorTest : public testing::Test {
     key_marshaller = std::make_shared<KeyMarshallerMock>();
     adaptor = std::make_shared<Plaintext>(marshaller, idmgr, key_marshaller);
 
-    ON_CALL(*conn, read(_, _, _)).WillByDefault(Arg2CallbackWithArg(5));
-    ON_CALL(*conn, writeSome(_, _, _)).WillByDefault(Arg2CallbackWithArg(5));
-
     ON_CALL(*idmgr, getKeyPair()).WillByDefault(ReturnRef(local_keypair));
     ON_CALL(*idmgr, getId()).WillByDefault(ReturnRef(local_pid));
     ON_CALL(*marshaller, marshal(_))
diff --git a/test/mock/libp2p/connection/layer_connection_mock.hpp b/test/mock/libp2p/connection/layer_connection_mock.hpp
@@ -7,20 +7,26 @@
 #pragma once
 
 #include <libp2p/connection/layer_connection.hpp>
+#include <libp2p/basic/read_return_size.hpp>
 
 #include <gmock/gmock.h>
 
 namespace libp2p::connection {
 
-  class LayerConnectionMock : public virtual LayerConnection {
+  class LayerConnectionMock : public std::enable_shared_from_this<LayerConnectionMock>, public virtual LayerConnection {
    public:
     ~LayerConnectionMock() override = default;
 
     MOCK_CONST_METHOD0(isClosed, bool(void));
 
     MOCK_METHOD0(close, outcome::result<void>());
 
-    MOCK_METHOD3(read, void(BytesOut, size_t, Reader::ReadCallbackFunc));
+    void read(BytesOut out,
+              size_t ambigous_size,
+              Reader::ReadCallbackFunc cb) override {
+      ASSERT_EQ(out.size(), ambigous_size);
+      readReturnSize(shared_from_this(), out, cb);
+    }
     MOCK_METHOD3(readSome, void(BytesOut, size_t, Reader::ReadCallbackFunc));
     MOCK_METHOD3(writeSome, void(BytesIn, size_t, Writer::WriteCallbackFunc));
     MOCK_METHOD2(deferReadCallback,
diff --git a/test/mock/libp2p/connection/stream_mock.hpp b/test/mock/libp2p/connection/stream_mock.hpp
diff --git a/test/testutil/expect_read.hpp b/test/testutil/expect_read.hpp
diff --git a/test/testutil/expect_write.hpp b/test/testutil/expect_write.hpp
diff --git a/test/testutil/libp2p/message_read_writer_helper.cpp b/test/testutil/libp2p/message_read_writer_helper.cpp
