dns_server_mock.hpp

#pragma once
 
#include <exception>
#include <functional>
#include <memory>
#include <string>
#include <variant>
#include <vector>
 
#include <userver/engine/io/sockaddr.hpp>
#include <userver/engine/task/task.hpp>
 
USERVER_NAMESPACE_BEGIN
 
namespace internal::net {
struct UdpListener;
}  // namespace internal::net
 
namespace utest {
 
class DnsServerMock final {
 public:
  enum class RecordType {
    kInvalid = 0,
    kA = 1,
    kAAAA = 28,
    kCname = 5,
  };
 
  struct DnsQuery {
    RecordType type{RecordType::kInvalid};
    std::string name;
  };
 
  struct DnsAnswer {
    using AnswerData =
        std::variant<std::monostate, engine::io::Sockaddr, std::string>;
 
    RecordType type{RecordType::kInvalid};
    AnswerData data;
    int ttl{0};
  };
 
  using DnsAnswerVector = std::vector<DnsAnswer>;
 
  struct NoAnswer : std::exception {};
 
  // throwing an exception will cause SERVFAIL
  using DnsHandler = std::function<DnsAnswerVector(const DnsQuery&)>;
 
  explicit DnsServerMock(DnsHandler);
 
  std::string GetServerAddress() const;
 
 private:
  void ProcessRequests();
 
  std::shared_ptr<internal::net::UdpListener> listener_;
  DnsHandler handler_;
  engine::Task receiver_task_;
};
 
}  // namespace utest
 
USERVER_NAMESPACE_END