#pragma once

#include <balancer/kernel/balancer/policy.h>

#include <library/cpp/testing/gmock_in_unittest/gmock.h>

namespace NSrvKernel::NTesting {

class TPolicyPartMock : public IPolicy {
public:
    MOCK_METHOD(void, MarkAsRetry, (), (override));
    MOCK_METHOD(IBackend*, Next, (IAlgorithm*, bool), (override));
};

class TPolicyMock : public TPolicyPartMock {
public:
    MOCK_METHOD(void, RegisterSuccess, (), (override));
    MOCK_METHOD(void, RegisterFail, (), (override));
};

}  // namespace NSrvKernel::NTesting
