#include <balancer/kernel/ban/banned_addresses.h>
#include <balancer/kernel/net/addr.h>

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

Y_UNIT_TEST_SUITE(TBannedAddressesTest) {
    Y_UNIT_TEST(TestBannedAddAndClear) {
        using namespace NSrvKernel;
        TSharedAllocator allocator;
        TSharedStatsManager statsManager(allocator);
        statsManager.SetWorkersCount(0, 0);
        TBannedAddresses banned(statsManager);
        TIp4Raw ip1 = {1, 1, 1, 1};
        TIp4Raw ip2 = {2, 2, 2, 2};
        TSockAddr bannedAddr(ip1, 1111);
        TSockAddr anotherPort(ip1, 2222);
        TSockAddr anotherIp(ip2, 1111);

        UNIT_ASSERT(!banned.IsBanned(bannedAddr));
        UNIT_ASSERT(!banned.IsBanned(anotherPort));
        UNIT_ASSERT(!banned.IsBanned(anotherIp));
        UNIT_ASSERT(banned.Add(bannedAddr, Now() + TDuration::Seconds(30), 1000000));
        UNIT_ASSERT(banned.IsBanned(bannedAddr));
        UNIT_ASSERT(banned.IsBanned(anotherPort));
        UNIT_ASSERT(!banned.IsBanned(anotherIp));
        banned.Clear();
        UNIT_ASSERT(!banned.IsBanned(bannedAddr));
        UNIT_ASSERT(!banned.IsBanned(anotherPort));
        UNIT_ASSERT(!banned.IsBanned(anotherIp));
    }

    Y_UNIT_TEST(TestBannedTtl) {
        using namespace NSrvKernel;
        TSharedAllocator allocator;
        TSharedStatsManager statsManager(allocator);
        statsManager.SetWorkersCount(0, 0);
        TBannedAddresses banned(statsManager);
        TIp4Raw ip1 = {1, 1, 1, 1};
        TIp4Raw ip2 = {2, 2, 2, 2};
        TIp4Raw ip3 = {1, 2, 1, 2};
        TSockAddr bannedAddr(ip1, 1111);
        TSockAddr bannedAddr2(ip2, 1111);
        TSockAddr bannedAddr3(ip3, 1111);

        UNIT_ASSERT(banned.Add(bannedAddr, Now() + TDuration::Seconds(20), 1000000));
        UNIT_ASSERT(banned.Add(bannedAddr2, Now() + TDuration::Seconds(10), 1000000));
        UNIT_ASSERT(banned.Add(bannedAddr2, Now() + TDuration::Seconds(40), 1000000));
        UNIT_ASSERT(banned.Add(bannedAddr3, Now() + TDuration::Seconds(10), 1000000));
        UNIT_ASSERT(banned.Add(bannedAddr3, Now() + TDuration::Seconds(20), 1000000));
        UNIT_ASSERT(banned.IsBanned(bannedAddr));
        UNIT_ASSERT(banned.IsBanned(bannedAddr2));
        UNIT_ASSERT(banned.IsBanned(bannedAddr3));
        banned.ProcessTtl(Now() + TDuration::Seconds(30));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr));
        UNIT_ASSERT(banned.IsBanned(bannedAddr2));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr3));
        UNIT_ASSERT(banned.Add(bannedAddr3, Now() + TDuration::Seconds(60), 1000000));
        banned.ProcessTtl(Now() + TDuration::Seconds(50));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr2));
        UNIT_ASSERT(banned.IsBanned(bannedAddr3));
        banned.ProcessTtl(Now() + TDuration::Seconds(70));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr2));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr3));
    }

    Y_UNIT_TEST(TestBannedLimit) {
        using namespace NSrvKernel;
        TSharedAllocator allocator;
        TSharedStatsManager statsManager(allocator);
        statsManager.SetWorkersCount(0, 0);
        TBannedAddresses banned(statsManager);
        TIp4Raw ip1 = {1, 1, 1, 1};
        TIp4Raw ip2 = {2, 2, 2, 2};
        TIp4Raw ip3 = {1, 2, 1, 2};
        TSockAddr bannedAddr(ip1, 1111);
        TSockAddr bannedAddr2(ip2, 1111);
        TSockAddr bannedAddr3(ip3, 1111);
        UNIT_ASSERT(banned.Add(bannedAddr, Now() + TDuration::Seconds(20), 2));
        UNIT_ASSERT(banned.Add(bannedAddr2, Now() + TDuration::Seconds(20), 2));
        UNIT_ASSERT(!banned.Add(bannedAddr3, Now() + TDuration::Seconds(20), 2));
        UNIT_ASSERT(banned.IsBanned(bannedAddr));
        UNIT_ASSERT(banned.IsBanned(bannedAddr2));
        UNIT_ASSERT(!banned.IsBanned(bannedAddr3));
    }
}
