#include <balancer/kernel/process/master_main.h>

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

#include <util/system/thread.h>


Y_UNIT_TEST_SUITE(MainTask) {
    TString config =
    " instance = {"
    " workers = 3;"
    " thread_mode = true;"
    "};";

    Y_UNIT_TEST(AddTree) {
        NConfig::TGlobals globals;
        NSrvKernel::TMainOptions options;
        options.AllowEmptyAdminAddrs = true;

        struct TMyTree
            : public NSrvKernel::ITree
        {
            TAtomic Inited = 0;
            TAtomic Disposed = 0;
        public:
            void Init(NSrvKernel::IWorkerCtl*) override {
                AtomicIncrement(Inited);
            }
            void Dispose(NSrvKernel::IWorkerCtl*) override {
                AtomicIncrement(Disposed);
            };
        } tree;

        NSrvKernel::NProcessCore::TMainTask mainTask(config, globals, nullptr, options, nullptr);
        mainTask.AddExternalTree(tree);

        TThread masterThread([&mainTask]() {
            mainTask.Execute();
        });
        masterThread.Start();

        TInstant start = Now();
        const TDuration timeout = TDuration::Seconds(5);

        while (AtomicGet(tree.Inited) < 3 && (Now() - start) < timeout) {
            Sleep(TDuration::MilliSeconds(10));
        }

        mainTask.RemoveExternalTree(tree);

        while (AtomicGet(tree.Disposed) < 3 && (Now() - start) < timeout) {
            Sleep(TDuration::MilliSeconds(10));
        }

        mainTask.StopMaster({
            .CoolDown = TDuration(),
            .Timeout = TDuration(),
            .CloseTimeout = TDuration::Max(),
        });

        UNIT_ASSERT_VALUES_EQUAL(tree.Inited, 3);
        UNIT_ASSERT_VALUES_EQUAL(tree.Disposed, 3);
    }
}
