#include <balancer/modules/balancer/module.h>
#include <library/cpp/testing/unittest/registar.h>
#include <balancer/kernel/http/parser/response_builder.h>

using namespace NSrvKernel;

class THttpOutputMock {
public:
    TError Send(TChunkList, TInstant) noexcept {
        return {};
    }

    TError SendHead(TResponse&&, bool, TInstant) noexcept {
        return {};
    }

    TError SendTrailers(THeaders&&, TInstant) noexcept {
        return {};
    }
};


Y_UNIT_TEST_SUITE(TouchTrackingHttpOutputTest) {

    Y_UNIT_TEST(TestHeadHeadSendSend) {
        THttpOutputMock outputMock;

        bool sentSomething = false;
        TMaybe<size_t> whoTouched;

        TTouchTrackingHttpOutput output1{ outputMock, sentSomething, 0, whoTouched };

        {
            auto err = output1.SendHead(BuildResponse().Version11().Code(200).ContentLength(10), false, TInstant::Max());
            UNIT_ASSERT(!err);
            UNIT_ASSERT(output1.Touched());
        }

        sentSomething = true;

        TTouchTrackingHttpOutput output2{ outputMock, sentSomething, 1, whoTouched };
        {
            auto err = output2.SendHead(BuildResponse().Version11().Code(200).ContentLength(10), false, TInstant::Max());
            UNIT_ASSERT(err);
            UNIT_ASSERT(!output2.Touched());
        }

        {
            auto err = output1.Send({}, TInstant::Max());
            UNIT_ASSERT(!err);
        }

        {
            auto err = output2.Send({}, TInstant::Max());
            UNIT_ASSERT(err);
        }
    }

    Y_UNIT_TEST(TestHeadSendHeadSend) {

        THttpOutputMock outputMock;

        bool sentSomething = false;
        TMaybe<size_t> whoTouched;

        TTouchTrackingHttpOutput output1{ outputMock, sentSomething, 0, whoTouched };

        {
            auto err = output1.SendHead(BuildResponse().Version11().Code(200).ContentLength(10), false, TInstant::Max());
            UNIT_ASSERT(!err);
            UNIT_ASSERT(output1.Touched());
        }

        sentSomething = true;

        {
            auto err = output1.Send({}, TInstant::Max());
            UNIT_ASSERT(!err);
        }

        TTouchTrackingHttpOutput output2{ outputMock, sentSomething, 1, whoTouched };
        {
            auto err = output2.SendHead(BuildResponse().Version11().Code(200).ContentLength(10), false, TInstant::Max());
            UNIT_ASSERT(err);
            UNIT_ASSERT(!output2.Touched());
        }

        {
            auto err = output2.Send({}, TInstant::Max());
            UNIT_ASSERT(err);
        }

        {
            auto err = output1.Send({}, TInstant::Max());
            UNIT_ASSERT(!err);
        }
    }

    Y_UNIT_TEST(TestSendSend) {
        THttpOutputMock outputMock;

        bool sentSomething = false;
        TMaybe<size_t> whoTouched;

        TTouchTrackingHttpOutput output1{ outputMock, sentSomething, 0, whoTouched };
        {
            auto err = output1.Send({}, TInstant::Max());
            UNIT_ASSERT(!err);
        }

        sentSomething = true;

        TTouchTrackingHttpOutput output2{ outputMock, sentSomething, 1, whoTouched };
        {
            auto err = output2.Send({}, TInstant::Max());
            UNIT_ASSERT(err);
        }
    }
}
