kafka_fixture.hpp

#pragma once
 
#include <atomic>
#include <chrono>
#include <deque>
#include <vector>
 
#include <userver/kafka/headers.hpp>
#include <userver/kafka/impl/broker_secrets.hpp>
#include <userver/kafka/impl/configuration.hpp>
#include <userver/kafka/impl/consumer.hpp>
#include <userver/kafka/producer.hpp>
#include <userver/utest/utest.hpp>
#include <userver/utils/box.hpp>
#include <userver/utils/span.hpp>
USERVER_NAMESPACE_BEGIN
 
namespace kafka::utest {
 
/// @ingroup userver_utest
///
/// @brief Message owning data wrapper for unit tests.
struct Message {
    std::string topic;
    std::string key;
    std::string payload;
    std::optional<std::uint32_t> partition{};
    std::vector<OwningHeader> headers{};
};
 
bool operator==(const Message& lhs, const Message& rhs);
 
/// @ingroup userver_utest
///
/// @brief Helper for Kafka unit testing.
///
/// KafkaCluster is useful for
/// - inline producer, consumer and their configuration creation;
/// - generation of unique topic names;
/// - sending and receiving batched of messages in sync form;
///
/// KafkaCluster expects that topic names used in test are already exist in
/// Kafka broker or auto-created on produce to it.
///
/// It is recommended to create all topics before tests started, because
/// topic creation leads to disk IO operation and maybe too slow in tests'
/// runtime.
///
/// If one run tests with Kafka testsuite plugin, use
/// `TESTSUITE_KAFKA_CUSTOM_TOPICS` environment variable to create topics before
/// tests started.
///
/// @note KafkaCluster determines the broker's connection string
/// based on `TESTSUITE_KAFKA_SERVER_HOST` (by default `localhost`) and
/// `TESTSUITE_KAFKA_SERVER_PORT` (by default `9099`) environment variables.
///
/// @see https://yandex.github.io/yandex-taxi-testsuite/kafka
class KafkaCluster : public ::testing::Test {
public:
    /// Kafka broker has some cold start issues on its. To stabilize tests
    /// KafkaCluster patches producer's default delivery timeout.
    /// To use custom delivery timeout in test, pass `configuration` argument to
    /// MakeProducer.
    static constexpr const std::chrono::milliseconds kDefaultTestProducerTimeout{
        USERVER_NAMESPACE::utest::kMaxTestWaitTime / 2};
 
    KafkaCluster();
 
    ~KafkaCluster() override;
 
    std::string GenerateTopic(std::uint32_t partition_cnt = 1);
 
    std::vector<std::string> GenerateTopics(std::size_t count);
 
    impl::Configuration MakeProducerConfiguration(
        const std::string& name,
        impl::ProducerConfiguration configuration = {},
        impl::Secret secrets = {}
    );
 
    impl::Configuration MakeConsumerConfiguration(
        const std::string& name,
        impl::ConsumerConfiguration configuration = {},
        impl::Secret secrets = {}
    );
 
    Producer MakeProducer(const std::string& name, impl::ProducerConfiguration configuration = {});
 
    std::deque<Producer> MakeProducers(
        std::size_t count,
        std::function<std::string(std::size_t)> nameGenerator,
        impl::ProducerConfiguration configuration = {}
    );
 
    /// @brief Creates temporary producer and send messages span.
    void SendMessages(utils::span<const Message> messages);
 
    impl::Consumer MakeConsumer(
        const std::string& name,
        const std::vector<std::string>& topics,
        impl::ConsumerConfiguration configuration = {},
        impl::ConsumerExecutionParams params = {}
    );
 
    /// @brief Starts consumer, waits until `expected_message_count` messages are
    /// consumed, calls `user_callback` if set, stops consumer.
    std::vector<Message> ReceiveMessages(
        impl::Consumer& consumer,
        std::size_t expected_messages_count,
        bool commit_after_receive = true,
        std::optional<std::function<void(MessageBatchView)>> user_callback = {}
    );
 
    /// @brief The same as previous, but working with consumer_scope passed to function.
    std::vector<Message> ReceiveMessages(
        ConsumerScope& consumer,
        std::size_t expected_messages_count,
        bool commit_after_receive = true,
        std::optional<std::function<void(MessageBatchView)>> user_callback = {}
    );
 
private:
    impl::Secret AddBootstrapServers(impl::Secret secrets) const;
    std::string InitBootstrapServers();
 
private:
    static std::atomic<std::size_t> kTopicsCount;
    class MockCluster;
    utils::Box<MockCluster> mock_;
    const std::string bootstrap_servers_;
};
}  // namespace kafka::utest
 
USERVER_NAMESPACE_END