rtsp_tensorrt/tests/test_pipeline.cpp

#include <gtest/gtest.h>
#include <filesystem>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "pipeline/common/pipeline.hpp"
#include "pipeline/common/yaml_config_parser.hpp"

namespace fs = std::filesystem;
namespace pipeline {

class PipelineTest : public ::testing::Test {
protected:
    void SetUp() override {
        // 创建测试目录
        test_dir_ = "/tmp/test_pipeline";
        if (fs::exists(test_dir_)) {
            fs::remove_all(test_dir_);
        }
        fs::create_directories(test_dir_);
        
        // 创建日志目录
        log_dir_ = test_dir_ + "/logs";
        fs::create_directories(log_dir_);
        
        // 创建测试视频
        video_path_ = test_dir_ + "/test_video.mp4";
        createTestVideo(video_path_);
        
        // 创建配置文件
        config_path_ = test_dir_ + "/config.yaml";
        createTestConfig(config_path_);
    }

    void TearDown() override {
        // 只在测试成功时删除文件
        if (!::testing::Test::HasFailure() && fs::exists(test_dir_)) {
            fs::remove_all(test_dir_);
        }
    }
    
    void createTestVideo(const std::string& path) {
        cv::Mat frame(480, 640, CV_8UC3, cv::Scalar(0, 0, 0));
        cv::VideoWriter writer(path, cv::VideoWriter::fourcc('a', 'v', 'c', '1'), 30, frame.size());
        for (int i = 0; i < 30; ++i) {
            cv::putText(frame, "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();
    }
    
    void createTestConfig(const std::string& path) {
        std::ofstream config_file(path);
        
        // 输入配置
        config_file << "input:\n";
        config_file << "  max_batch_size: 1\n";
        config_file << "  sources:\n";
        config_file << "    - type: video\n";
        config_file << "      name: test_video\n";
        config_file << "      url: " << video_path_ << "\n";
        config_file << "      buffer_size: 30\n";
        config_file << "      outputs: [output_video]\n";
        config_file << "\n";
        
        // 推理配置
        config_file << "inference:\n";
        config_file << "  model:\n";
        config_file << "    onnx_path: /app/models/yolov8n.onnx\n";
        config_file << "    engine_path: /app/models/yolov8n.engine\n";
        config_file << "    input_shape: [3, 640, 640]\n";
        config_file << "    precision: FP16\n";
        config_file << "  threshold:\n";
        config_file << "    conf: 0.5\n";
        config_file << "    nms: 0.45\n";
        config_file << "  gpu_id: 0\n";
        config_file << "\n";
        
        // 渲染配置
        config_file << "render:\n";
        config_file << "  enable: true\n";
        config_file << "  window:\n";
        config_file << "    name: Detection Results\n";
        config_file << "    width: 1280\n";
        config_file << "    height: 720\n";
        config_file << "    fullscreen: false\n";
        config_file << "  default_style:\n";
        config_file << "    box_color: [0, 255, 0]\n";
        config_file << "    text_color: [255, 255, 255]\n";
        config_file << "    transparency: 0.0\n";
        config_file << "    box_thickness: 2\n";
        config_file << "    font_scale: 0.5\n";
        config_file << "    font_thickness: 1\n";
        config_file << "  class_styles:\n";
        config_file << "    person:\n";
        config_file << "      box_color: [0, 255, 0]\n";
        config_file << "      text_color: [255, 255, 255]\n";
        config_file << "      transparency: 0.0\n";
        config_file << "      box_thickness: 2\n";
        config_file << "      font_scale: 0.5\n";
        config_file << "      font_thickness: 1\n";
        config_file << "  metrics:\n";
        config_file << "    show_fps: true\n";
        config_file << "    show_inference_time: true\n";
        config_file << "    show_gpu_usage: true\n";
        config_file << "    update_interval_ms: 1000\n";
        config_file << "\n";
        
        // 输出配置
        config_file << "output:\n";
        config_file << "  targets:\n";
        config_file << "    - type: video\n";
        config_file << "      name: output_video\n";
        config_file << "      path: " << test_dir_ << "/output.mp4\n";
        config_file << "      fps: 30\n";
        config_file << "      codec: h264\n";
        config_file << "      bitrate: 4000000\n";
        config_file << "\n";
        
        // 日志配置
        config_file << "log:\n";
        config_file << "  level: info\n";
        config_file << "  save_path: " << log_dir_ << "\n\n";
        
        config_file.close();
    }

    std::string test_dir_;
    std::string video_path_;
    std::string config_path_;
    std::string log_dir_;
};

// 测试配置加载
TEST_F(PipelineTest, ConfigLoading) {
    YamlConfigParser parser;
    ASSERT_TRUE(parser.parse(config_path_));
    
    const auto& config = parser.getConfig();
    
    // 验证输入配置
    EXPECT_EQ(config.input.max_batch_size, 1);
    EXPECT_EQ(config.input.sources.size(), 1);
    EXPECT_EQ(config.input.sources[0].type, "video");
    EXPECT_EQ(config.input.sources[0].name, "test_video");
    EXPECT_EQ(config.input.sources[0].url, video_path_);
    
    // 验证推理配置
    EXPECT_EQ(config.inference.engine_path, "/app/models/yolov8n.engine");
    EXPECT_EQ(config.inference.input_shape, std::vector<int>({3, 640, 640}));
    EXPECT_EQ(config.inference.precision, "FP16");
    
    // 验证输出配置
    EXPECT_EQ(config.output.targets.size(), 1);
    EXPECT_EQ(config.output.targets[0].type, "video");
    EXPECT_EQ(config.output.targets[0].name, "output_video");
}

// 测试Pipeline初始化
TEST_F(PipelineTest, Initialization) {
    Pipeline pipeline(config_path_, true);  // 使用测试模式
    ASSERT_TRUE(pipeline.init());
    ASSERT_TRUE(pipeline.start());
    EXPECT_TRUE(pipeline.isRunning());
    pipeline.stop();
    pipeline.wait();
}

// 测试Pipeline运行
TEST_F(PipelineTest, Running) {
    Pipeline pipeline(config_path_, true);  // 使用测试模式
    ASSERT_TRUE(pipeline.init());
    
    // 运行一段时间
    ASSERT_TRUE(pipeline.start());
    std::this_thread::sleep_for(std::chrono::seconds(2));
    pipeline.stop();
    pipeline.wait();
    
    // 验证输出文件是否生成
    std::string output_path = test_dir_ + "/output.mp4";
    EXPECT_TRUE(fs::exists(output_path));
    EXPECT_GT(fs::file_size(output_path), 0);
}

// 测试错误处理
TEST_F(PipelineTest, ErrorHandling) {
    // 测试无效的配置文件
    Pipeline pipeline_invalid("invalid_config.yaml", true);
    EXPECT_FALSE(pipeline_invalid.init());
    
    // 测试无效的引擎路径
    std::string invalid_config_path = test_dir_ + "/invalid_config.yaml";
    std::ofstream config_file(invalid_config_path);
    config_file << "inference:\n";
    config_file << "  model:\n";
    config_file << "    onnx_path: /app/models/yolov8n.onnx\n";
    config_file << "    engine_path: invalid_engine_path\n";
    config_file << "    input_shape: [3, 640, 640]\n";
    config_file.close();
    
    Pipeline pipeline_invalid_engine(invalid_config_path, true);
    EXPECT_FALSE(pipeline_invalid_engine.init());
}

// 测试资源管理
TEST_F(PipelineTest, ResourceManagement) {
    {
        Pipeline pipeline(config_path_, true);
        ASSERT_TRUE(pipeline.init());
        ASSERT_TRUE(pipeline.start());
        std::this_thread::sleep_for(std::chrono::seconds(1));
        // 不调用stop,让析构函数处理
    }
    
    // 验证资源是否正确释放
    EXPECT_TRUE(fs::exists(test_dir_ + "/output.mp4"));
}

} // namespace pipeline