ANSCustomModels/tests/HelmetDetection/HelmetDetectionTest.cpp

#include "TestCommon.h"
#include "ANSCustomCodeHelmetDetection.h"

// ===========================================================================
// Unit Tests — no model files required
// ===========================================================================

class HelmetUnitTest : public ::testing::Test {
protected:
    ANSCustomHMD detector;
};

TEST_F(HelmetUnitTest, EmptyFrameReturnsNoDetections) {
    cv::Mat empty;
    auto results = detector.RunInference(empty);
    EXPECT_TRUE(results.empty());
}

TEST_F(HelmetUnitTest, TinyFrameReturnsNoDetections) {
    cv::Mat tiny = TestUtils::CreateTestFrame(5, 5);
    auto results = detector.RunInference(tiny);
    EXPECT_TRUE(results.empty());
}

TEST_F(HelmetUnitTest, UninitializedDetectorReturnsNoDetections) {
    cv::Mat frame = TestUtils::CreateTestFrame(640, 480);
    auto results = detector.RunInference(frame);
    EXPECT_TRUE(results.empty());
}

TEST_F(HelmetUnitTest, RunInferenceWithCameraId) {
    cv::Mat frame = TestUtils::CreateTestFrame(640, 480);
    auto results = detector.RunInference(frame, "test_cam_01");
    EXPECT_TRUE(results.empty());
}

TEST_F(HelmetUnitTest, ConfigureParametersReturnsValidConfig) {
    CustomParams params;
    bool result = detector.ConfigureParameters(params);
    EXPECT_TRUE(result);
}

TEST_F(HelmetUnitTest, DestroySucceeds) {
    EXPECT_TRUE(detector.Destroy());
}

TEST_F(HelmetUnitTest, DestroyCanBeCalledMultipleTimes) {
    EXPECT_TRUE(detector.Destroy());
    EXPECT_TRUE(detector.Destroy());
}

TEST_F(HelmetUnitTest, InitializeWithInvalidDirectoryFails) {
    std::string labelMap;
    bool result = detector.Initialize("C:\\NonExistent\\Path\\Model", 0.5f, labelMap);
    EXPECT_FALSE(result);
}

TEST_F(HelmetUnitTest, OptimizeBeforeInitializeReturnsFalse) {
    EXPECT_FALSE(detector.OptimizeModel(true));
}

// ===========================================================================
// Integration Tests — require model files on disk
// ===========================================================================

class HelmetIntegrationTest : public ::testing::Test {
protected:
    ANSCustomHMD detector;
    std::string labelMap;
    std::vector<std::string> classes;

    void SetUp() override {
        if (!TestConfig::ModelExists(TestConfig::HELMET_MODEL_DIR)) {
            GTEST_SKIP() << "Helmet model not found at: " << TestConfig::HELMET_MODEL_DIR;
        }
        bool ok = detector.Initialize(TestConfig::HELMET_MODEL_DIR, 0.6f, labelMap);
        ASSERT_TRUE(ok) << "Failed to initialize Helmet detector";
        classes = TestUtils::ParseLabelMap(labelMap);
    }

    void TearDown() override {
        detector.Destroy();
    }
};

TEST_F(HelmetIntegrationTest, InitializeProducesLabelMap) {
    EXPECT_FALSE(labelMap.empty());
    EXPECT_FALSE(classes.empty());
}

TEST_F(HelmetIntegrationTest, InferenceOnSolidFrameReturnsNoDetections) {
    cv::Mat frame = TestUtils::CreateTestFrame(1920, 1080);
    auto results = detector.RunInference(frame, "test_cam");
    EXPECT_TRUE(results.empty()) << "Solid gray frame should not trigger helmet detection";
}

TEST_F(HelmetIntegrationTest, InferenceOnSmallFrame) {
    cv::Mat frame = TestUtils::CreateTestFrame(320, 240);
    auto results = detector.RunInference(frame, "test_cam");
    SUCCEED();
}

TEST_F(HelmetIntegrationTest, InferenceOnLargeFrame) {
    cv::Mat frame = TestUtils::CreateTestFrame(3840, 2160);
    auto results = detector.RunInference(frame, "test_cam");
    SUCCEED();
}

TEST_F(HelmetIntegrationTest, DetectionResultFieldsAreValid) {
    if (!TestConfig::VideoExists(TestConfig::HELMET_VIDEO)) {
        GTEST_SKIP() << "Helmet test video not found";
    }

    cv::VideoCapture cap(TestConfig::HELMET_VIDEO);
    ASSERT_TRUE(cap.isOpened());

    bool detectionFound = false;
    for (int i = 0; i < 300 && !detectionFound; i++) {
        cv::Mat frame;
        if (!cap.read(frame)) break;

        auto results = detector.RunInference(frame, "test_cam");
        for (const auto& obj : results) {
            detectionFound = true;
            EXPECT_GE(obj.confidence, 0.0f);
            EXPECT_LE(obj.confidence, 1.0f);
            EXPECT_GE(obj.box.width, 0);
            EXPECT_GE(obj.box.height, 0);
            EXPECT_GE(obj.classId, 0);
        }
    }
    cap.release();
}

TEST_F(HelmetIntegrationTest, PerformanceBenchmark) {
    if (!TestConfig::VideoExists(TestConfig::HELMET_VIDEO)) {
        GTEST_SKIP() << "Helmet test video not found";
    }

    auto [totalDetections, avgMs] = TestUtils::RunVideoFrames(detector, TestConfig::HELMET_VIDEO, 100);
    ASSERT_GE(totalDetections, 0) << "Video could not be opened";

    std::cout << "[Helmet] 100 frames: avg=" << avgMs << "ms/frame, "
              << "detections=" << totalDetections << std::endl;

    EXPECT_LT(avgMs, 200.0) << "Average inference time exceeds 200ms";
}

TEST_F(HelmetIntegrationTest, ThreadSafetyConcurrentInference) {
    cv::Mat frame1 = TestUtils::CreateTestFrame(640, 480, cv::Scalar(100, 100, 100));
    cv::Mat frame2 = TestUtils::CreateTestFrame(640, 480, cv::Scalar(200, 200, 200));

    std::vector<CustomObject> results1, results2;

    std::thread t1([&]() { results1 = detector.RunInference(frame1, "cam_1"); });
    std::thread t2([&]() { results2 = detector.RunInference(frame2, "cam_2"); });

    t1.join();
    t2.join();

    SUCCEED();
}