#include <maps/wikimap/mapspro/services/mrc/eye/lib/detect_road_marking/impl/yt_spec.h>

#include <library/cpp/testing/unittest/registar.h>
#include <mapreduce/yt/tests/yt_unittest_lib/yt_unittest_lib.h>

#include <util/generic/size_literals.h>

namespace maps::mrc::eye::tests {

Y_UNIT_TEST_SUITE(yt_spec_tests)
{

Y_UNIT_TEST(cpu_yt_spec)
{
    YtWorkerConfig config;
    config.useGpu = false;
    config.concurrency = 10;

    NYT::TNode spec = makeDetectWorkerSpec(config);

    UNIT_ASSERT_EQUAL(
        spec["resource_limits"]["user_slots"].AsUint64(),
        config.concurrency
    );

    UNIT_ASSERT_EQUAL(spec["mapper"]["cpu_limit"].AsInt64(), 1);
    UNIT_ASSERT_EQUAL(spec["mapper"]["memory_limit"].AsUint64(), 8_GB);
    UNIT_ASSERT_EQUAL(spec["mapper"]["gpu_limit"].IsUndefined(), true);
}

Y_UNIT_TEST(gpu_yt_spec)
{
    YtWorkerConfig config;
    config.useGpu = true;
    config.concurrency = 10;

    NYT::TNode spec = makeDetectWorkerSpec(config);

    UNIT_ASSERT_EQUAL(
        spec["resource_limits"]["user_slots"].AsUint64(),
        config.concurrency
    );

    UNIT_ASSERT_EQUAL(spec["mapper"]["gpu_limit"].AsInt64(), 1);
    UNIT_ASSERT_EQUAL(spec["mapper"]["memory_limit"].AsUint64(), 4_GB);
    UNIT_ASSERT_EQUAL(spec["mapper"]["cpu_limit"].IsUndefined(), true);
}

} // Y_UNIT_TEST_SUITE

} // namespace maps::mrc::eye::tests
