#include <balancer/kernel/custom_io/ut/test_common.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/helpers/errors.h>

#include <library/cpp/coroutine/engine/impl.h>

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

Y_UNIT_TEST_SUITE(TCustomIoStreamTest) {
    using namespace NSrvKernel;
    const char DATA[] = "testtesttest";
    const char TEST[] = "test";
    constexpr TStringBuf CHUNK_DATA = "testtesttest";

    Y_UNIT_TEST(TestChunksOutputStream) {
        TChunksOutputStream st;
        for (size_t i = 0; i < 3; ++i) {
            st.Write(TEST, 4);
        }
        UNIT_ASSERT_EQUAL(st.Chunks(), DATA);
    }

    Y_UNIT_TEST(TestChunksOutputStream2) {
        TChunksOutputStream out;

        for (size_t i = 0; i < 3; ++i) {
            out.Write(TEST, 4);
        }

        char buf[CHUNK_DATA.size() + 1];
        out.Chunks().CopyDataTo(buf, CHUNK_DATA.size());
        buf[CHUNK_DATA.size()] = 0;
        UNIT_ASSERT_VALUES_EQUAL(buf, DATA);
    }

    Y_UNIT_TEST(TestChunkInputStream) {
        TChunkList chunk(NewChunkNonOwning(CHUNK_DATA));
        TChunkInputStream st(std::move(chunk));
        for (size_t i = 0; i < 3; ++i) {
            char lst[5];
            lst[4] = 0;
            st.Read(&lst, 4);
            UNIT_ASSERT_VALUES_EQUAL(TEST, lst);
        }
        TString str;
        st >> str;
        UNIT_ASSERT(str.empty());
    }

    static void DoTestYieldingOutput(TCont* cont, void* arg) {
        TChunksOutputStream* st = (TChunksOutputStream*)(arg);
        TYieldingOutput out(*st, *cont, 8, 0);

        for (size_t i = 0; i < 3; ++i) {
            out.Write(TEST, 4);
        }
    }

    Y_UNIT_TEST(TestYieldingOutput) {
        TChunksOutputStream st;

        TContExecutor e(32000);
        e.Execute(DoTestYieldingOutput, &st);
        UNIT_ASSERT_EQUAL(st.Chunks(), DATA);
    }

};
