rtsp_tensorrt/tests/test_input_manager.cpp

#include <gtest/gtest.h>
#include "pipeline/input/input_manager.hpp"
#include <thread>
#include <chrono>
#include <filesystem>
#include <opencv2/videoio.hpp>
#include <opencv2/imgproc.hpp>
#include <iostream>

using namespace pipeline;

class InputManagerTest : public ::testing::Test {
protected:
    void SetUp() override {
        try {
            // 创建测试视频文件
            std::filesystem::path current_path = std::filesystem::current_path();
            test_video_path_ = (current_path / "test_input_manager.avi").string();
            std::cout << "Creating test video at: " << test_video_path_ << std::endl;
            createTestVideo();
            std::cout << "Test video created successfully" << std::endl;

            // 配置测试参数
            video_config_.buffer_size = 30;  // 增加缓冲区大小
            video_config_.target_fps = 10.0f;  // 降低目标帧率
            video_config_.frame_timeout_ms = 1000;  // 增加超时时间
            video_config_.loop_playback = true;  // 循环播放
            video_config_.async_reading = true;  // 异步读取
        } catch (const std::exception& e) {
            std::cerr << "Exception in SetUp: " << e.what() << std::endl;
            throw;
        }
    }

    void TearDown() override {
        try {
            // 确保资源被正确释放
            std::cout << "Cleaning up resources..." << std::endl;
            manager_.clear();
            
            // 删除测试视频文件
            if (std::filesystem::exists(test_video_path_)) {
                std::cout << "Removing test video file..." << std::endl;
                std::filesystem::remove(test_video_path_);
            }
            std::cout << "Cleanup completed" << std::endl;
        } catch (const std::exception& e) {
            std::cerr << "Exception in TearDown: " << e.what() << std::endl;
        }
    }

    void createTestVideo() {
        // 创建一个简单的测试视频文件
        cv::VideoWriter writer(test_video_path_, 
                             cv::VideoWriter::fourcc('M', 'J', 'P', 'G'),  // 使用 MJPG 编码器
                             30, cv::Size(640, 480));
        
        if (!writer.isOpened()) {
            throw std::runtime_error("Failed to create test video file: " + test_video_path_);
        }
        
        std::cout << "Generating test video frames..." << std::endl;
        // 生成 100 帧的测试视频，增加帧数
        for (int i = 0; i < 100; ++i) {
            cv::Mat frame(480, 640, CV_8UC3, cv::Scalar(0, 0, 0));
            cv::rectangle(frame, cv::Point(100, 100), cv::Point(200, 200), 
                         cv::Scalar(0, 255, 0), 2);
            cv::putText(frame, "Test Frame " + std::to_string(i), 
                       cv::Point(50, 50), cv::FONT_HERSHEY_SIMPLEX, 
                       1.0, cv::Scalar(255, 255, 255), 2);
            writer.write(frame);
        }
        writer.release();
        
        // 验证文件是否创建成功
        if (!std::filesystem::exists(test_video_path_)) {
            throw std::runtime_error("Test video file was not created: " + test_video_path_);
        }
        std::cout << "Test video generation completed" << std::endl;
    }

    // 辅助函数：等待源启动
    bool waitForSourceStart(const std::string& name, int timeout_ms = 500) {
        auto start = std::chrono::steady_clock::now();
        SourceStatus status;
        while (true) {
            if (manager_.getSourceStatus(name, status) && status.is_connected) {
                return true;
            }
            
            if (std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now() - start).count() > timeout_ms) {
                std::cout << "Timeout waiting for source to start: " << name << std::endl;
                return false;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }

    // 辅助函数：测量帧率
    float measureFrameRate(const std::string& source_name, int duration_ms = 2000) {  // 增加测量时间
        std::vector<cv::Mat> frames;
        int frame_count = 0;
        
        // 等待第一帧
        std::cout << "Waiting for first frame..." << std::endl;
        auto wait_start = std::chrono::steady_clock::now();
        bool got_first_frame = false;
        while (!got_first_frame) {
            if (manager_.getNextBatch(frames, 1000)) {  // 增加超时时间
                frames.clear();
                got_first_frame = true;
                std::cout << "Got first frame, starting measurement..." << std::endl;
                break;
            }
            if (std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now() - wait_start).count() > 10000) {  // 增加等待时间
                std::cout << "Timeout waiting for first frame" << std::endl;
                return 0.0f;
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
        
        if (!got_first_frame) {
            return 0.0f;
        }
        
        // 开始测量帧率
        auto start = std::chrono::steady_clock::now();
        while (std::chrono::duration_cast<std::chrono::milliseconds>(
               std::chrono::steady_clock::now() - start).count() < duration_ms) {
            if (manager_.getNextBatch(frames, 50)) {
                frame_count++;
                frames.clear();
            }
            std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
        
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now() - start).count();
        float fps = static_cast<float>(frame_count * 1000) / elapsed;
        std::cout << "Measurement completed: " << frame_count << " frames in " 
                  << elapsed << "ms (" << fps << " fps)" << std::endl;
        return fps;
    }

    VideoReader::Config video_config_;
    std::string test_video_path_;
    InputManager manager_{30};  // 增加队列容量
};

// 测试添加源
TEST_F(InputManagerTest, AddSource) {
    std::cout << "Starting AddSource test..." << std::endl;
    
    // 添加有效源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    EXPECT_EQ(manager_.getSourceCount(), 1);
    
    // 添加重复源
    EXPECT_FALSE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_EQ(manager_.getSourceCount(), 1);
    
    // 添加另一个源
    EXPECT_TRUE(manager_.addVideoSource("camera2", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera2"));
    EXPECT_EQ(manager_.getSourceCount(), 2);
    
    std::cout << "AddSource test completed" << std::endl;
}

// 测试移除源
TEST_F(InputManagerTest, RemoveSource) {
    std::cout << "Starting RemoveSource test..." << std::endl;
    
    // 添加源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    EXPECT_EQ(manager_.getSourceCount(), 1);
    
    // 移除源
    EXPECT_TRUE(manager_.removeSource("camera1"));
    EXPECT_EQ(manager_.getSourceCount(), 0);
    
    // 移除不存在的源
    EXPECT_FALSE(manager_.removeSource("camera1"));
    
    std::cout << "RemoveSource test completed" << std::endl;
}

// 测试获取帧
TEST_F(InputManagerTest, GetNextBatch) {
    std::cout << "Starting GetNextBatch test..." << std::endl;
    
    std::vector<cv::Mat> frames;
    
    // 添加源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    
    // 等待一小段时间让帧积累
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    
    // 获取帧
    auto start = std::chrono::steady_clock::now();
    bool got_frame = false;
    while (std::chrono::duration_cast<std::chrono::milliseconds>(
           std::chrono::steady_clock::now() - start).count() < 500) {
        if (manager_.getNextBatch(frames, 50)) {
            got_frame = true;
            break;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }
    
    EXPECT_TRUE(got_frame) << "Failed to get frame within timeout";
    EXPECT_FALSE(frames.empty());
    
    // 检查帧属性
    for (const auto& frame : frames) {
        EXPECT_FALSE(frame.empty());
        EXPECT_EQ(frame.cols, 640);
        EXPECT_EQ(frame.rows, 480);
        EXPECT_EQ(frame.channels(), 3);
    }
    
    std::cout << "GetNextBatch test completed" << std::endl;
}

// 测试源状态
TEST_F(InputManagerTest, SourceStatus) {
    std::cout << "Starting SourceStatus test..." << std::endl;
    
    SourceStatus status;
    
    // 添加源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    
    // 获取状态
    EXPECT_TRUE(manager_.getSourceStatus("camera1", status));
    EXPECT_TRUE(status.is_connected);
    EXPECT_GE(status.frame_count, 0);
    
    // 获取不存在的源状态
    EXPECT_FALSE(manager_.getSourceStatus("invalid_camera", status));
    
    std::cout << "SourceStatus test completed" << std::endl;
}

// 测试清空
TEST_F(InputManagerTest, Clear) {
    std::cout << "Starting Clear test..." << std::endl;
    
    // 添加多个源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    EXPECT_TRUE(manager_.addVideoSource("camera2", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera2"));
    EXPECT_EQ(manager_.getSourceCount(), 2);
    
    // 清空
    manager_.clear();
    EXPECT_EQ(manager_.getSourceCount(), 0);
    
    // 检查是否可以重新添加源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    EXPECT_EQ(manager_.getSourceCount(), 1);
    
    std::cout << "Clear test completed" << std::endl;
}

// 测试获取源名称
TEST_F(InputManagerTest, GetSourceNames) {
    std::cout << "Starting GetSourceNames test..." << std::endl;
    
    // 添加源
    EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera1"));
    EXPECT_TRUE(manager_.addVideoSource("camera2", video_config_, test_video_path_));
    EXPECT_TRUE(waitForSourceStart("camera2"));
    
    // 获取名称列表
    auto names = manager_.getSourceNames();
    EXPECT_EQ(names.size(), 2);
    EXPECT_TRUE(std::find(names.begin(), names.end(), "camera1") != names.end());
    EXPECT_TRUE(std::find(names.begin(), names.end(), "camera2") != names.end());
    
    std::cout << "GetSourceNames test completed" << std::endl;
}

// 测试帧率控制
TEST_F(InputManagerTest, FrameRateControl) {
    std::cout << "Starting FrameRateControl test..." << std::endl;
    
    // 设置不同的目标帧率
    const std::vector<float> target_fps = {5.0f, 10.0f, 15.0f};
    const float tolerance = 0.3f;  // 允许 30% 的误差
    
    for (float fps : target_fps) {
        std::cout << "Testing target FPS: " << fps << std::endl;
        
        // 确保清理之前的状态
        manager_.clear();
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        
        // 重新配置参数
        video_config_.target_fps = fps;
        video_config_.buffer_size = 10;  // 增加缓冲区大小
        video_config_.frame_timeout_ms = 100;  // 增加超时时间
        
        // 添加源
        EXPECT_TRUE(manager_.addVideoSource("camera1", video_config_, test_video_path_));
        EXPECT_TRUE(waitForSourceStart("camera1"));
        
        // 等待帧率稳定
        std::this_thread::sleep_for(std::chrono::milliseconds(500));
        
        // 测量实际帧率
        float measured_fps = measureFrameRate("camera1", 2000);  // 增加测量时间
        std::cout << "Measured FPS: " << measured_fps << std::endl;
        
        // 验证帧率是否在允许范围内
        EXPECT_NEAR(measured_fps, fps, fps * tolerance)
            << "Target FPS: " << fps << ", Measured FPS: " << measured_fps;
        
        // 清理
        manager_.clear();
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    
    std::cout << "FrameRateControl test completed" << std::endl;
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}