
#include <saas/searchproxy/common/abstract.h>
#include <saas/searchproxy/configs/serviceconfig.h>

#include <search/meta/scatter/source.h>

#include <infra/yp_service_discovery/libs/sdlib/client.h>
#include <library/cpp/testing/unittest/registar.h>

namespace NProxyMeta {
    void FillIntervalsAndSources(const TSearchSources& searchSources,
        NSaas::TShardIntervals& shardIntervals,
        TVector<THolder<NScatter::ISource>>& sources, const NScatter::TSourceOptions& opts,
        const bool precharge, NYP::NServiceDiscovery::TSDClient* sdClient);
}

Y_UNIT_TEST_SUITE(TestSourcesSharding) {
    Y_UNIT_TEST(TestShardingParsing) {
        TSearchSources searchSources;
        searchSources.resize(8);
        searchSources[0].ProtoConfig_.SetServerGroup("0-8190");
        searchSources[0].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[1].ProtoConfig_.SetServerGroup("16382-24572");
        searchSources[1].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[2].ProtoConfig_.SetServerGroup("24573-32763");
        searchSources[2].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[3].ProtoConfig_.SetServerGroup("32764-40954");
        searchSources[3].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[4].ProtoConfig_.SetServerGroup("40955-49145");
        searchSources[4].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[5].ProtoConfig_.SetServerGroup("49146-57336");
        searchSources[5].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[6].ProtoConfig_.SetServerGroup("57337-65533");
        searchSources[6].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        searchSources[7].ProtoConfig_.SetServerGroup("8191-16381");
        searchSources[7].ProtoConfig_.SetSearchScript("(tcp2://localhost:7309 )@1@1");

        NScatter::TSourceOptions opts;
        opts.EnableIpV6 = true;
        opts.EnableUnresolvedHosts = true;
        NSaas::TShardIntervals shardIntervals;
        TVector<THolder<NScatter::ISource>> scatterSources;
        NProxyMeta::FillIntervalsAndSources(searchSources, shardIntervals, scatterSources, opts, false, nullptr);
        UNIT_ASSERT_EQUAL(scatterSources.size(), shardIntervals.size());
        UNIT_ASSERT(IsSorted(shardIntervals.cbegin(), shardIntervals.cend()));
        for (ui32 i = 0; i < shardIntervals.size(); ++i) {
            const ui32 num = FromString<ui32>(scatterSources[i]->Descr);
            UNIT_ASSERT_EQUAL(searchSources[num].ProtoConfig_.GetServerGroup(), shardIntervals[i].ToString());
        }
    }
}

