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

Y_UNIT_TEST_SUITE(TCustomIoLimitIoTest) {
    using namespace NSrvKernel;
    const TStringBuf DATA = "test0test1test2";
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);

    Y_UNIT_TEST(TestLengthLimitInput) {
        auto mock = MakeInputMock([](TChunkList& lst, TInstant) {
            UNIT_ASSERT(lst.Empty());
            lst.PushNonOwning(DATA);
            return TError{};
        });

        TLengthLimitInput in(&mock, DATA.size());
        UNIT_ASSERT(!in.Depleted());

        TChunkList chunk;
        auto err = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT_EQUAL(chunk, DATA);
        UNIT_ASSERT(in.Depleted());

        for (size_t i = 0; i < 3; ++i) {
            in.SetLimit(5);
            auto err = in.Recv(chunk, DLINE);
            UNIT_ASSERT(!err);
            UNIT_ASSERT_EQUAL(chunk, "test" + ToString(i));
            UNIT_ASSERT(in.Depleted());
        }

    }

    Y_UNIT_TEST(TestLengthLimitInputWithError) {
        auto mock = MakeInputMock([](TChunkList&, TInstant deadline) {
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return Y_MAKE_ERROR(TSystemError{});
        });

        TLengthLimitInput in(&mock, DATA.size());
        UNIT_ASSERT(!in.Depleted());
        TChunkList chunk;
        auto err1 = in.Recv(chunk, DLINE);
        UNIT_ASSERT(err1.GetAs<TSystemError>());
        UNIT_ASSERT(chunk.Empty());
        UNIT_ASSERT(!in.Depleted());

        in.SetLimit(0);
        UNIT_ASSERT(in.Depleted());
        auto err2 = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err2);
        UNIT_ASSERT(chunk.Empty());
    }

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

        TLengthEqualInput in(&mock, DATA.size());
        TChunkList chunk;
        auto err = in.Recv(chunk, DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT_EQUAL(chunk, DATA);
    }

    Y_UNIT_TEST(TestTLengthEqualInputWithException) {
        auto mock = MakeInputMock([](TChunkList&, TInstant deadline) {
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return TError{};
        });

        TLengthEqualInput in(&mock, DATA.size());
        TChunkList chunk;
        auto error = in.Recv(chunk, DLINE);
        UNIT_ASSERT(error);
        UNIT_ASSERT_UNEQUAL(error.GetAs<THttpError>(), nullptr);
        UNIT_ASSERT_EQUAL(error.GetAs<THttpError>()->Code(), 400);
    }

    Y_UNIT_TEST(TestTLengthEqualInputWithError) {
        auto mock = MakeInputMock([](TChunkList&, TInstant deadline) {
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            return Y_MAKE_ERROR(TSystemError{});
        });

        TLengthEqualInput in(&mock, DATA.size());
        TChunkList chunk;
        auto err = in.Recv(chunk, DLINE);
        UNIT_ASSERT(err.GetAs<TSystemError>());
        UNIT_ASSERT(chunk.Empty());
    }
};
