#include <balancer/kernel/io/iobase.h>
#include <library/cpp/testing/unittest/registar.h>
#include <util/generic/deque.h>

Y_UNIT_TEST_SUITE(TBalancerIoBaseTest) {
    using namespace NSrvKernel;
    const TInstant DLINE = TInstant::Now() + TDuration::Days(1);

    class TFakeInput : public IIoInput {
    public:
        TFakeInput(TDeque<TChunkList> portions)
            : Portions(std::move(portions))
        {}

        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            UNIT_ASSERT_VALUES_EQUAL(deadline, DLINE);
            if (Portions) {
                lst = std::move(Portions.front());
                Portions.pop_front();
            } else {
                UNIT_ASSERT(!Eof);
                lst = TChunkList();
                Eof = true;
            }
            return {};
        }

        TDeque<TChunkList> Portions;
        bool Eof = false;
    };

    auto MakeLongList() {
        TDeque<TChunkList> lst;
        lst.emplace_back().Push("a").Push("b");
        return lst;
    }

    auto MakeShortLists() {
        TDeque<TChunkList> lst;
        lst.emplace_back("a");
        lst.emplace_back("b");
        return lst;
    }

    Y_UNIT_TEST(TestFillBuffer) {
        size_t filled;
        TError error;
        {
            TFakeInput finp{{}};
            auto error = finp.FillBuffer(100, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 0);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "");
        }
        {
            TFakeInput finp{{}};
            auto error = finp.FillBuffer(0, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 0);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "");
        }
        {
            TFakeInput finp(MakeLongList());
            auto error = finp.FillBuffer(0, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 0);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "");
        }
        {
            TFakeInput finp(MakeLongList());
            auto error = finp.FillBuffer(1, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 2);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "ab");
        }
        {
            TFakeInput finp(MakeLongList());
            auto error = finp.FillBuffer(3, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 2);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "ab");
        }
        {
            TFakeInput finp(MakeShortLists());
            auto error = finp.FillBuffer(1, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 1);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "a");
        }
        {
            TFakeInput finp(MakeShortLists());
            auto error = finp.FillBuffer(3, DLINE).AssignTo(filled);
            UNIT_ASSERT(!error);
            UNIT_ASSERT_EQUAL(filled, 2);
            UNIT_ASSERT_VALUES_EQUAL(finp.RecvBuffered(), "ab");
        }
    }
};
