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

Y_UNIT_TEST_SUITE(TCustomIoChunkedTransferTest) {
    using namespace NSrvKernel;
    const TStringBuf DATA = "4\r\ntest\r\n";
    const TStringBuf CANONDATA = "test";
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);
    const TVector<TStringBuf> BAD_DATA{"4\rtest\r\n", "=.test", "1\r\ntest\r\n"};

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

        TChunkedTransferDecoder decoder(&mock);
        TChunkList chunk;
        auto err = decoder.Recv(chunk, DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT_EQUAL(chunk, CANONDATA);
    }

    Y_UNIT_TEST(TestChunkedTransferDecoderWithTHttpError) {
        size_t ind = 0;
        auto mock = MakeInputMock([&ind](TChunkList& lst, TInstant deadline) {
            UNIT_ASSERT(lst.Empty());
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            lst.PushNonOwning(BAD_DATA[ind]);
            return TError{};
        });

        TChunkedTransferDecoder decoder(&mock);
        TChunkList chunk;
        for (;ind < BAD_DATA.size(); ++ind) {
            auto error = decoder.Recv(chunk, DLINE);
            UNIT_ASSERT(error);
            UNIT_ASSERT_UNEQUAL(error.GetAs<THttpError>(), nullptr);
            UNIT_ASSERT_EQUAL(error.GetAs<THttpError>()->Code(), 400);
        }
    }

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

        TChunkedTransferDecoder decoder(&mock);
        TChunkList chunk;
        auto err = decoder.Recv(chunk, DLINE);
        UNIT_ASSERT(err);
        UNIT_ASSERT(err.GetAs<TSystemError>());
    }

    Y_UNIT_TEST(TestChunkedTransferEncoder) {
        auto mock = MakeOutputMock([](TChunkList lst) {
            UNIT_ASSERT_EQUAL(lst, DATA);
            return TError{};
        });
        TChunkedTransferEncoder encoder(&mock);
        TChunkList chunk(NewChunkNonOwning(CANONDATA));
        auto err = encoder.Send(std::move(chunk), DLINE);
        UNIT_ASSERT(!err);
        UNIT_ASSERT(chunk.Empty());
    }

    Y_UNIT_TEST(TestChunkedTransferEncoderWithError) {
        auto mock = MakeOutputMock([](TChunkList lst) {
            UNIT_ASSERT_EQUAL(lst, DATA);
            return Y_MAKE_ERROR(TSystemError{});
        });
        TChunkedTransferEncoder encoder(&mock);
        TChunkList chunk(NewChunkNonOwning(CANONDATA));
        auto err = encoder.Send(std::move(chunk), DLINE);
        UNIT_ASSERT(err);
        UNIT_ASSERT(err.GetAs<TSystemError>());
    }

};
