#include "capped_buffered_output.h"

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

#include <util/random/mersenne.h>

namespace {
    // by courtesy of util/stream/buffered_ut.cpp
    template <class TOut>
    inline void Run(TOut && out) {
        TMersenne<ui64> r;

        for (size_t i = 0; i < 1000; ++i) {
            const size_t c = r.GenRand() % 10000;
            TString s;

            for (size_t j = 0; j < c; ++j) {
                s.append('A' + (r.GenRand() % 10));
            }

            out.Write(s.data(), s.size());
        }

        out.Finish();
    }

    class TWriteCapValidator : public IOutputStream {
    public:
        TWriteCapValidator(size_t cap)
            : Cap(cap)
            , IncompleteWrite(false)
        {
        }

    protected:
        void DoWrite(const void* /*buf*/, size_t len) override {
            UNIT_ASSERT(len <= Cap);
            if (len < Cap) {
                UNIT_ASSERT(IncompleteWrite == false); // allowed only once (last chunk)
                IncompleteWrite = true;
            }
        }

    private:
        size_t Cap;
        bool IncompleteWrite;
    };

    class TStringOutputWithSuffix : public TStringOutput {
    public:
        TStringOutputWithSuffix(TString& outstring, TString suffix)
            : TStringOutput(outstring)
            , Suffix(std::move(suffix))
        {
        }

    protected:
        void DoFinish() override {
            DoWrite(Suffix.data(), Suffix.size());
            TStringOutput::DoFinish();
        }

    private:
        TString Suffix;
    };
}

Y_UNIT_TEST_SUITE(CappedBufferedOutputSuite) {
    Y_UNIT_TEST(TestEqual) {
        TString s1;
        TString s2;

        Run(TStringOutput(s1));
        Run(TCappedBufferedOutput<TStringOutput>(8000, s2));

        UNIT_ASSERT_VALUES_EQUAL(s1, s2);
    }

    Y_UNIT_TEST(TestCap) {
        Run(TCappedBufferedOutput<TWriteCapValidator>(123, 123));
    }

    Y_UNIT_TEST(TestSuffix) {
        TString s1;
        TString s2;
        const TString suffix = ".suffix";

        Run(TStringOutputWithSuffix(s1, suffix));
        Run(TCappedBufferedOutput<TStringOutputWithSuffix>(8000, s2, suffix));

        UNIT_ASSERT_VALUES_EQUAL(s1, s2);
    }
}
