#include <balancer/kernel/thread/time_manager.h>
#include <balancer/kernel/thread/time_arbiter.h>

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

using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TTimeManagerTest) {
    class TTimeArbiterStupidStub final : public IThreadTimeArbiter {
    public:
        void SleepUntil(TInstant tp) noexcept override {
            if (tp > CurrentTime) {
                CurrentTime = tp;
            }
        }

        TInstant Now() const noexcept override {
            return CurrentTime;
        }

    private:
        TInstant CurrentTime = ::Now();
    };

    Y_UNIT_TEST(SmokingTest) {
        TInstant now = TThreadTimeManager::Instance().Now();
        UNIT_ASSERT(now <= Now());
    }

    Y_UNIT_TEST(StubTest) {
        TThreadTimeManager::Instance().EnableTesting();
        {
            // RAII that will put original time arbiter at block completion
            TThreadTimeManager::TMockGuard guard{MakeHolder<TTimeArbiterStupidStub>()};

            TInstant now = TThreadTimeManager::Instance().Now();
            UNIT_ASSERT(now <= Now());

            // Sleep test
            TDuration dur = TDuration::Seconds(5);
            TThreadTimeManager::Instance().Sleep(dur);
            UNIT_ASSERT_EQUAL(TThreadTimeManager::Instance().Now(), now + dur);

            // SleepUntil test
            TInstant tp = TThreadTimeManager::Instance().Now() + dur * 2;
            TThreadTimeManager::Instance().SleepUntil(tp);
            UNIT_ASSERT_EQUAL(TThreadTimeManager::Instance().Now(), tp);

            // We wouldn't sleep until time point in the past
            TThreadTimeManager::Instance().SleepUntil(tp - dur);
            UNIT_ASSERT_EQUAL(TThreadTimeManager::Instance().Now(), tp);

            // ToDeadline test
            TInstant deadline = TThreadTimeManager::Instance().ToDeadline(dur);
            UNIT_ASSERT_EQUAL(deadline, TThreadTimeManager::Instance().Now() + dur);
        }
    }
}
