From bc9ca41bc11587712e5cb77488f2f4e5fbf50fc2 Mon Sep 17 00:00:00 2001
From: Sergey Kuznetsov <skuznetsov@ripple.com>
Date: Thu, 5 Sep 2024 16:33:04 +0100
Subject: [PATCH] test: Add test for WsConnection for ping response (#1619)

---
 tests/common/util/TestWsServer.cpp            | 22 +++++++++++
 tests/common/util/TestWsServer.hpp            | 10 +++++
 .../unit/util/requests/WsConnectionTests.cpp  | 37 +++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/tests/common/util/TestWsServer.cpp b/tests/common/util/TestWsServer.cpp
index 66c8b175..1d7a0e3d 100644
--- a/tests/common/util/TestWsServer.cpp
+++ b/tests/common/util/TestWsServer.cpp
@@ -75,6 +75,14 @@ TestWsConnection::send(std::string const& message, boost::asio::yield_context yi
     return std::nullopt;
 }
 
+void
+TestWsConnection::sendPing(boost::beast::websocket::ping_data const& data, boost::asio::yield_context yield)
+{
+    boost::beast::error_code errorCode;
+    ws_.async_ping(data, yield[errorCode]);
+    [&]() { ASSERT_FALSE(errorCode) << errorCode.message(); }();
+}
+
 std::optional<std::string>
 TestWsConnection::receive(boost::asio::yield_context yield)
 {
@@ -105,6 +113,20 @@ TestWsConnection::headers() const
     return headers_;
 }
 
+void
+TestWsConnection::setControlFrameCallback(
+    std::function<void(boost::beast::websocket::frame_type, std::string_view)> callback
+)
+{
+    ws_.control_callback(std::move(callback));
+}
+
+void
+TestWsConnection::resetControlFrameCallback()
+{
+    ws_.control_callback();
+}
+
 TestWsServer::TestWsServer(asio::io_context& context, std::string const& host) : acceptor_(context)
 {
     auto endpoint = asio::ip::tcp::endpoint(boost::asio::ip::make_address(host), 0);
diff --git a/tests/common/util/TestWsServer.hpp b/tests/common/util/TestWsServer.hpp
index 96cf0885..4ce2d079 100644
--- a/tests/common/util/TestWsServer.hpp
+++ b/tests/common/util/TestWsServer.hpp
@@ -25,6 +25,7 @@
 #include <boost/asio/ip/tcp.hpp>
 #include <boost/asio/spawn.hpp>
 #include <boost/beast/core/tcp_stream.hpp>
+#include <boost/beast/websocket/rfc6455.hpp>
 #include <boost/beast/websocket/stream.hpp>
 
 #include <expected>
@@ -54,6 +55,9 @@ public:
     std::optional<std::string>
     send(std::string const& message, boost::asio::yield_context yield);
 
+    void
+    sendPing(boost::beast::websocket::ping_data const& data, boost::asio::yield_context yield);
+
     // returns nullopt if the connection is closed
     std::optional<std::string>
     receive(boost::asio::yield_context yield);
@@ -63,6 +67,12 @@ public:
 
     std::vector<util::requests::HttpHeader> const&
     headers() const;
+
+    void
+    setControlFrameCallback(std::function<void(boost::beast::websocket::frame_type, std::string_view)> callback);
+
+    void
+    resetControlFrameCallback();
 };
 using TestWsConnectionPtr = std::unique_ptr<TestWsConnection>;
 
diff --git a/tests/unit/util/requests/WsConnectionTests.cpp b/tests/unit/util/requests/WsConnectionTests.cpp
index 9f226507..da02807d 100644
--- a/tests/unit/util/requests/WsConnectionTests.cpp
+++ b/tests/unit/util/requests/WsConnectionTests.cpp
@@ -25,6 +25,8 @@
 #include <boost/asio/error.hpp>
 #include <boost/asio/spawn.hpp>
 #include <boost/beast/http/field.hpp>
+#include <boost/beast/websocket/stream.hpp>
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <chrono>
@@ -33,6 +35,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <string_view>
 #include <thread>
 #include <vector>
 
@@ -318,6 +321,40 @@ TEST_F(WsConnectionTests, MultipleConnections)
     }
 }
 
+TEST_F(WsConnectionTests, RespondsToPing)
+{
+    asio::spawn(ctx, [&](asio::yield_context yield) {
+        auto serverConnection = unwrap(server.acceptConnection(yield));
+
+        testing::StrictMock<testing::MockFunction<void(boost::beast::websocket::frame_type, std::string_view)>>
+            controlFrameCallback;
+        serverConnection.setControlFrameCallback(controlFrameCallback.AsStdFunction());
+        EXPECT_CALL(controlFrameCallback, Call(boost::beast::websocket::frame_type::pong, testing::_)).WillOnce([&]() {
+            serverConnection.resetControlFrameCallback();
+            asio::spawn(ctx, [&](asio::yield_context yield) {
+                auto maybeError = serverConnection.send("got pong", yield);
+                ASSERT_FALSE(maybeError.has_value()) << *maybeError;
+            });
+        });
+
+        serverConnection.sendPing({}, yield);
+        auto message = serverConnection.receive(yield);
+        ASSERT_TRUE(message.has_value());
+        EXPECT_EQ(message, "hello") << message.value();
+    });
+
+    runSpawn([&](asio::yield_context yield) {
+        auto connection = builder.plainConnect(yield);
+        ASSERT_TRUE(connection.has_value()) << connection.error().message();
+        auto expectedMessage = connection->operator*().read(yield);
+        ASSERT_TRUE(expectedMessage) << expectedMessage.error().message();
+        EXPECT_EQ(expectedMessage.value(), "got pong");
+
+        auto error = connection->operator*().write("hello", yield);
+        ASSERT_FALSE(error) << error->message();
+    });
+}
+
 enum class WsConnectionErrorTestsBundle : int { Read = 1, Write = 2 };
 
 struct WsConnectionErrorTests : WsConnectionTestsBase, testing::WithParamInterface<WsConnectionErrorTestsBundle> {};