#include "test_common.h"

#include <balancer/kernel/custom_io/countio.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/http/parser/response_builder.h>

#include <library/cpp/testing/unittest/registar.h>

Y_UNIT_TEST_SUITE(TCustomIoCountTest) {
    using namespace NSrvKernel;
    const TStringBuf DATA = "test";
    constexpr int NUM = 12111;
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);
    const TResponse SAMPLE_RESPONSE = BuildResponse().Version11()
                                                     .Code(HTTP_OK)
                                                     .Header("abc", "def");

    Y_UNIT_TEST(TestCountInputEmpty) {
        TChunkList chunk;

        auto mock = MakeInputMock([](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return TError{};
        });

        TCountInput countIn(&mock);
        auto err = countIn.Recv(chunk, DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT(chunk.Empty());
        UNIT_ASSERT(!countIn.Readed());
    }

    Y_UNIT_TEST(TestCountInput) {
        TChunkList chunk{ NewChunkNonOwning(DATA) };

        auto mock = MakeInputMock([](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            lst.Push(ToString(NUM));
            return TError{};
        });

        TCountInput countIn(&mock);
        auto err = countIn.Recv(chunk, DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT_EQUAL(chunk, ToString(NUM));
        UNIT_ASSERT_EQUAL(countIn.Readed(), ToString(NUM).size());
    }

    Y_UNIT_TEST(TestCountInputWithError) {
        TChunkList chunk{ NewChunkNonOwning(DATA) };

        auto mockWithError = MakeInputMock([](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return Y_MAKE_ERROR(TSystemError{});
        });

        TCountInput countIn(&mockWithError);
        auto err = countIn.Recv(chunk, DLINE);
        UNIT_ASSERT(chunk.Empty());
        UNIT_ASSERT(!countIn.Readed());
        UNIT_ASSERT(err.GetAs<TSystemError>());
    }
};
