#include "chunk.h"

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

#include <util/stream/file.h>
#include <util/system/tempfile.h>
#include <util/stream/null.h>

#define CDATA "./chunkedio"

Y_UNIT_TEST_SUITE(TestChunkedIO) {
    static const char test_data[] = "87s6cfbsudg cuisg s igasidftasiy tfrcua6s";

    TString CombString(const TString& s, size_t chunkSize) {
        TString result;
        for (size_t pos = 0; pos < s.size(); pos += 2 * chunkSize) 
            result += s.substr(pos, chunkSize); 
        return result; 
    } 
 
    void WriteTestData(IOutputStream * stream, TString * contents) {
        contents->clear(); 
        for (size_t i = 0; i < sizeof(test_data); ++i) { 
            stream->Write(test_data, i); 
            contents->append(test_data, i); 
        } 
    } 
 
    void ReadInSmallChunks(IInputStream * stream, TString * contents) {
        char buf[11]; 
        size_t read = 0; 
 
        contents->clear(); 
        do { 
            read = stream->Read(buf, sizeof(buf)); 
            contents->append(buf, read); 
        } while (read > 0); 
    } 
 
    void ReadCombed(IInputStream * stream, TString * contents, size_t chunkSize) {
        Y_ASSERT(chunkSize < 128);
        char buf[128]; 
 
        contents->clear(); 
        while (true) { 
            size_t read = stream->Load(buf, chunkSize); 
            contents->append(buf, read); 
            if (read == 0) 
                break; 
 
            size_t toSkip = chunkSize; 
            size_t skipped = 0; 
            do { 
                skipped = stream->Skip(toSkip); 
                toSkip -= skipped; 
            } while (skipped != 0 && toSkip != 0); 
        } 
    } 
 
    Y_UNIT_TEST(TestChunkedIo) {
        TTempFile tmpFile(CDATA);
        TString tmp;

        {
            TUnbufferedFileOutput fo(CDATA);
            TChunkedOutput co(&fo);
            WriteTestData(&co, &tmp); 
        } 

        { 
            TUnbufferedFileInput fi(CDATA);
            TChunkedInput ci(&fi); 
            TString r;

            ReadInSmallChunks(&ci, &r); 
 
            UNIT_ASSERT_EQUAL(r, tmp); 
        }

        {
            TUnbufferedFileInput fi(CDATA);
            TChunkedInput ci(&fi);
            TString r;

            ReadCombed(&ci, &r, 11); 

            UNIT_ASSERT_EQUAL(r, CombString(tmp, 11)); 
        }
    }

    Y_UNIT_TEST(TestBadChunk) {
        bool hasError = false;

        try {
            TString badChunk = "10\r\nqwerty";
            TMemoryInput mi(badChunk.data(), badChunk.size());
            TChunkedInput ci(&mi);
            TransferData(&ci, &Cnull);
        } catch (...) {
            hasError = true;
        }

        UNIT_ASSERT(hasError);
    }
}