#include <balancer/kernel/custom_io/ut/test_common.h>
#include <balancer/kernel/custom_io/concat.h>
#include <balancer/kernel/helpers/errors.h>
#include <library/cpp/testing/unittest/registar.h>

Y_UNIT_TEST_SUITE(TCustomIoConcatTest) {
    using namespace NSrvKernel;
    const TStringBuf DATA1 = "test";
    const TStringBuf DATA2 = "data";
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);

    Y_UNIT_TEST(TestConcatInput) {
        bool isRecv = true;
        auto mock1 = MakeInputMock([&isRecv](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (isRecv) {
                lst.PushNonOwning(DATA1);
            }
            return TError{};
        });
        auto mock2 = MakeInputMock([&isRecv](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (isRecv) {
                lst.PushNonOwning(DATA2);
            }
            return TError{};
        });

        TConcatInput in(&mock1, &mock2);

        TChunkList chunk;
        auto err1 = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err1);
        UNIT_ASSERT_EQUAL(chunk, DATA1);

        isRecv = false;

        auto err2 = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err2);
        UNIT_ASSERT(chunk.Empty());

        isRecv = true;

        auto err3 = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err3);
        UNIT_ASSERT_EQUAL(chunk, DATA2);

        isRecv = false;

        for (size_t i = 0; i < 5; ++i) {
            auto err = in.Recv(chunk, DLINE);
            UNIT_ASSERT(!err);
            UNIT_ASSERT(chunk.Empty());
        }
    }

    Y_UNIT_TEST(TestConcatInputWithError) {
        bool isRecv = true;
        auto mock1  = MakeInputMock([&isRecv](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (isRecv) {
                lst.PushNonOwning(DATA1);
            }
            return TError{};
        });
        auto mock2 = MakeInputMock([](TChunkList&, TInstant deadline){
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return Y_MAKE_ERROR(TSystemError{});
        });

        TConcatInput in(&mock1, &mock2);

        TChunkList chunk;
        auto err1 = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err1);
        UNIT_ASSERT_EQUAL(chunk, DATA1);

        isRecv = false;

        for (size_t i = 0; i < 5; ++i) {
            auto err = in.Recv(chunk, DLINE);
            UNIT_ASSERT(err.GetAs<TSystemError>());
            UNIT_ASSERT(chunk.Empty());
        }
    }
};
