clio/tests/common/util/TestWsServer.hpp

#pragma once

#include "util/requests/Types.hpp"

#include <boost/asio/io_context.hpp>
#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>
#include <functional>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <vector>

class TestWsConnection {
    boost::beast::websocket::stream<boost::beast::tcp_stream> ws_;
    std::vector<util::requests::HttpHeader> headers_;

public:
    using SendCallback = std::function<void()>;
    using ReceiveCallback = std::function<void(std::string)>;

    TestWsConnection(
        boost::beast::websocket::stream<boost::beast::tcp_stream> wsStream,
        std::vector<util::requests::HttpHeader> headers
    );

    TestWsConnection(TestWsConnection&& other);

    // returns error message if error occurs
    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);

    std::optional<std::string>
    close(boost::asio::yield_context yield);

    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>;

class TestWsServer {
    boost::asio::ip::tcp::acceptor acceptor_;

public:
    TestWsServer(boost::asio::io_context& context, std::string const& host);

    std::string
    port() const;

    std::expected<TestWsConnection, util::requests::RequestError>
    acceptConnection(boost::asio::yield_context yield);

    void
    acceptConnectionAndDropIt(boost::asio::yield_context yield);

    boost::asio::ip::tcp::socket
    acceptConnectionWithoutHandshake(boost::asio::yield_context yield);
};