Alicho/tests/process_manager/test_process_manager_manager.cpp

#include <gtest/gtest.h>
#include <thread>
#include <chrono>
#include <atomic>
#include <condition_variable>
#include <mutex>
#include <filesystem>
#include <fstream>

#include "process_manager.h"
#include "process_types.h"

using namespace alicho;

// ============================================================================
// 测试辅助工具
// ============================================================================

// 回调计数器，用于跟踪回调函数的调用
struct callback_counter {
	std::atomic<int> state_change_count{0};
	std::atomic<int> resource_update_count{0};
	std::atomic<int> error_count{0};

	std::mutex              mutex;
	std::condition_variable cv;

	process_state last_state{process_state::UNKNOWN};
	process_error last_error{process_error::SUCCESS};
	std::string   last_error_message;

	void reset() {
		state_change_count    = 0;
		resource_update_count = 0;
		error_count           = 0;
		last_state            = process_state::UNKNOWN;
		last_error            = process_error::SUCCESS;
		last_error_message.clear();
	}

	bool wait_for_state_change(int                       expected_count,
	                           std::chrono::milliseconds timeout = std::chrono::milliseconds(5000)) {
		std::unique_lock<std::mutex> lock(mutex);
		return cv.wait_for(lock, timeout, [this, expected_count]() {
			return state_change_count >= expected_count;
		});
	}

	bool wait_for_error(int expected_count, std::chrono::milliseconds timeout = std::chrono::milliseconds(5000)) {
		std::unique_lock<std::mutex> lock(mutex);
		return cv.wait_for(lock, timeout, [this, expected_count]() {
			return error_count >= expected_count;
		});
	}
};

// 全局回调计数器
static callback_counter g_callback_counter;

// 回调函数
void state_change_callback(const process_info& info, process_state new_state) {
	std::lock_guard<std::mutex> lock(g_callback_counter.mutex);
	g_callback_counter.state_change_count++;
	g_callback_counter.last_state = new_state;
	g_callback_counter.cv.notify_all();
}

void resource_callback(const process_info& info, const process_resource_usage& usage) {
	std::lock_guard<std::mutex> lock(g_callback_counter.mutex);
	g_callback_counter.resource_update_count++;
	g_callback_counter.cv.notify_all();
}

void error_callback(const process_info& info, process_error error, const std::string& message) {
	std::lock_guard<std::mutex> lock(g_callback_counter.mutex);
	g_callback_counter.error_count++;
	g_callback_counter.last_error         = error;
	g_callback_counter.last_error_message = message;
	g_callback_counter.cv.notify_all();
}

// 创建一个简单的测试可执行文件
std::string create_test_executable() {
	std::string exe_path;

	#ifdef _WIN32
	exe_path = "test_process.bat";
	std::ofstream bat_file(exe_path);
	bat_file << "@echo off\n";
	bat_file << "ping 127.0.0.1 -n 6 > nul\n"; // 运行约5秒，给监控器足够时间
	bat_file << "exit /b 0\n";
	bat_file.close();
	#else
	exe_path = "./test_process.sh";
	std::ofstream script_file(exe_path);
	script_file << "#!/bin/bash\n";
	script_file << "sleep 5\n"; // 运行5秒然后退出
	script_file << "exit 0\n";
	script_file.close();

	// 设置执行权限
	std::filesystem::permissions(exe_path, std::filesystem::perms::owner_exec, std::filesystem::perm_options::add);
	#endif

	return exe_path;
}

// 创建一个立即失败的测试可执行文件
std::string create_failing_executable() {
	std::string exe_path;

	#ifdef _WIN32
	exe_path = "test_fail_process.bat";
	std::ofstream bat_file(exe_path);
	bat_file << "@echo off\n";
	bat_file << "exit /b 1\n"; // 立即以错误码退出
	bat_file.close();
	#else
	exe_path = "./test_fail_process.sh";
	std::ofstream script_file(exe_path);
	script_file << "#!/bin/bash\n";
	script_file << "exit 1\n"; // 立即以错误码退出
	script_file.close();

	// 设置执行权限
	std::filesystem::permissions(exe_path, std::filesystem::perms::owner_exec, std::filesystem::perm_options::add);
	#endif

	return exe_path;
}

// ============================================================================
// 基础功能测试
// ============================================================================

class ProcessManagerTest : public ::testing::Test {
protected:
	void SetUp() override {
		// 重置回调计数器
		g_callback_counter.reset();

		// 设置回调函数
		auto& manager = process_manager::instance();
		manager.set_state_change_callback(state_change_callback);
		manager.set_resource_callback(resource_callback);
		manager.set_error_callback(error_callback);

		// 创建测试可执行文件
		test_executable_    = create_test_executable();
		failing_executable_ = create_failing_executable();
	}

	void TearDown() override {
		// 终止所有进程
		auto& manager = process_manager::instance();
		manager.terminate_all_processes(true);

		// 清理测试文件
		try {
			if (std::filesystem::exists(test_executable_)) {
				std::filesystem::remove(test_executable_);
			}
			if (std::filesystem::exists(failing_executable_)) {
				std::filesystem::remove(failing_executable_);
			}
		}
		catch (...) {
			// 忽略清理错误
		}
	}

	std::string test_executable_;
	std::string failing_executable_;
};

/**
 * @brief 测试单例模式正常工作
 */
TEST_F(ProcessManagerTest, SingletonPattern) {
	auto& manager1 = process_manager::instance();
	auto& manager2 = process_manager::instance();

	// 应该是同一个实例
	EXPECT_EQ(&manager1, &manager2);

	// 初始状态应该是0个进程
	EXPECT_EQ(manager1.process_count(), 0);
}

/**
 * @brief 测试进程启动功能
 */
TEST_F(ProcessManagerTest, LaunchProcess) {
	auto& manager = process_manager::instance();

	// 启动一个简单的进程
	uint32_t process_id = manager.launch_process(
		"test_process",
		test_executable_,
		{},
		sandbox_type::HOST,
		true // 启用监控
	);

	EXPECT_NE(process_id, 0); // 应该返回有效的进程ID

	// 检查进程是否被添加到管理列表
	EXPECT_TRUE(manager.has_process(process_id));
	EXPECT_EQ(manager.process_count(), 1);

	// 检查进程信息
	const process_info* info = manager.get_process_info(process_id);
	ASSERT_NE(info, nullptr);
	EXPECT_EQ(info->name, "test_process");
	EXPECT_EQ(info->executable_path, test_executable_);
	EXPECT_EQ(info->type, sandbox_type::HOST);

	// 等待进程开始运行
	std::this_thread::sleep_for(std::chrono::milliseconds(100));

	// 检查进程是否在运行
	EXPECT_TRUE(manager.is_process_running(process_id));

	// 终止进程
	auto result = manager.terminate_process(process_id);
	EXPECT_EQ(result, process_error::SUCCESS);
}

/**
 * @brief 测试进程终止功能
 */
TEST_F(ProcessManagerTest, TerminateProcess) {
	auto& manager = process_manager::instance();

	// 启动进程
	uint32_t process_id = manager.launch_process(
		"test_process_terminate",
		test_executable_,
		{},
		sandbox_type::HOST,
		false // 不启用监控，简化测试
	);

	ASSERT_NE(process_id, 0);

	// 等待进程启动
	std::this_thread::sleep_for(std::chrono::milliseconds(100));

	// 正常终止进程
	auto result = manager.terminate_process(process_id, false);
	EXPECT_EQ(result, process_error::SUCCESS);

	// 检查进程是否已停止
	EXPECT_FALSE(manager.is_process_running(process_id));
}

/**
 * @brief 测试强制终止进程
 */
TEST_F(ProcessManagerTest, ForceTerminateProcess) {
	auto& manager = process_manager::instance();

	// 启动进程
	uint32_t process_id = manager.launch_process(
		"test_process_force_terminate",
		test_executable_,
		{},
		sandbox_type::HOST,
		false
	);

	ASSERT_NE(process_id, 0);

	// 等待进程启动
	std::this_thread::sleep_for(std::chrono::milliseconds(100));

	// 强制终止进程
	auto result = manager.terminate_process(process_id, true);
	EXPECT_EQ(result, process_error::SUCCESS);

	// 检查进程是否已停止
	EXPECT_FALSE(manager.is_process_running(process_id));
}

/**
 * @brief 测试终止所有进程
 */
TEST_F(ProcessManagerTest, TerminateAllProcesses) {
	auto& manager = process_manager::instance();

	// 启动多个进程
	std::vector<uint32_t> process_ids;
	for (int i = 0; i < 3; ++i) {
		uint32_t process_id = manager.launch_process(
			"test_process_" + std::to_string(i),
			test_executable_,
			{},
			sandbox_type::HOST,
			false
		);
		if (process_id != 0) {
			process_ids.push_back(process_id);
		}
	}

	EXPECT_GT(process_ids.size(), 0);
	EXPECT_EQ(manager.process_count(), process_ids.size());

	// 等待进程启动
	std::this_thread::sleep_for(std::chrono::milliseconds(200));

	// 终止所有进程
	auto result = manager.terminate_all_processes();
	EXPECT_EQ(result, process_error::SUCCESS);

	// 检查所有进程是否已停止
	for (uint32_t process_id : process_ids) {
		EXPECT_FALSE(manager.is_process_running(process_id));
	}
}

/**
 * @brief 测试获取进程ID列表
 */
TEST_F(ProcessManagerTest, GetAllProcessIds) {
	auto& manager = process_manager::instance();

	// 初始应该没有进程
	auto initial_ids = manager.get_all_process_ids();
	EXPECT_TRUE(initial_ids.empty());

	// 启动几个进程
	std::vector<uint32_t> launched_ids;
	for (int i = 0; i < 2; ++i) {
		uint32_t process_id = manager.launch_process(
			"test_process_list_" + std::to_string(i),
			test_executable_,
			{},
			sandbox_type::HOST,
			false
		);
		if (process_id != 0) {
			launched_ids.push_back(process_id);
		}
	}

	// 获取进程ID列表
	auto current_ids = manager.get_all_process_ids();
	EXPECT_EQ(current_ids.size(), launched_ids.size());

	// 验证ID列表是否包含所有已启动的进程
	for (uint32_t launched_id : launched_ids) {
		EXPECT_TRUE(std::find(current_ids.begin(), current_ids.end(), launched_id) != current_ids.end());
	}
}

// ============================================================================
// 错误处理测试
// ============================================================================

/**
 * @brief 测试启动不存在的可执行文件
 */
TEST_F(ProcessManagerTest, LaunchNonExistentExecutable) {
	auto& manager = process_manager::instance();

	// 尝试启动不存在的可执行文件
	uint32_t process_id = manager.launch_process(
		"non_existent_process",
		"non_existent_executable.exe",
		{},
		sandbox_type::HOST,
		false
	);

	// 应该返回0表示失败
	EXPECT_EQ(process_id, 0);

	// 进程计数应该没有增加
	EXPECT_EQ(manager.process_count(), 0);
}

/**
 * @brief 测试终止不存在的进程
 */
TEST_F(ProcessManagerTest, TerminateNonExistentProcess) {
	auto& manager = process_manager::instance();

	// 尝试终止不存在的进程
	uint32_t fake_process_id = 99999;
	auto     result          = manager.terminate_process(fake_process_id);

	EXPECT_EQ(result, process_error::PROCESS_NOT_RUNNING);
}

/**
 * @brief 测试重复终止已终止的进程
 */
TEST_F(ProcessManagerTest, TerminateAlreadyTerminatedProcess) {
	auto& manager = process_manager::instance();

	// 启动进程
	uint32_t process_id = manager.launch_process(
		"test_process_double_terminate",
		test_executable_,
		{},
		sandbox_type::HOST,
		false
	);

	ASSERT_NE(process_id, 0);

	// 等待进程启动
	std::this_thread::sleep_for(std::chrono::milliseconds(100));

	// 第一次终止
	auto result1 = manager.terminate_process(process_id);
	EXPECT_EQ(result1, process_error::SUCCESS);

	// 等待进程真正停止
	std::this_thread::sleep_for(std::chrono::milliseconds(100));

	// 第二次终止同一个进程
	auto result2 = manager.terminate_process(process_id);
	EXPECT_EQ(result2, process_error::PROCESS_NOT_RUNNING);
}

/**
 * @brief 测试查询不存在进程的信息
 */
TEST_F(ProcessManagerTest, GetNonExistentProcessInfo) {
	auto& manager = process_manager::instance();

	// 查询不存在的进程信息
	uint32_t            fake_process_id = 88888;
	const process_info* info            = manager.get_process_info(fake_process_id);

	EXPECT_EQ(info, nullptr);
	EXPECT_FALSE(manager.has_process(fake_process_id));
	EXPECT_FALSE(manager.is_process_running(fake_process_id));
}

// ============================================================================
// 沙箱类型测试
// ============================================================================

/**
 * @brief 测试不同沙箱类型的进程
 */
TEST_F(ProcessManagerTest, DifferentSandboxTypes) {
	auto& manager = process_manager::instance();

	// 测试不同的沙箱类型
	std::vector<sandbox_type> sandbox_types = {
		sandbox_type::HOST,
		sandbox_type::PLUGIN,
		sandbox_type::CUSTOM
	};

	std::vector<uint32_t> process_ids;

	for (auto type : sandbox_types) {
		uint32_t process_id = manager.launch_process(
			"test_sandbox_" + sandbox_type_to_string(type),
			test_executable_,
			{},
			type,
			false
		);

		if (process_id != 0) {
			process_ids.push_back(process_id);

			// 验证沙箱类型
			const process_info* info = manager.get_process_info(process_id);
			ASSERT_NE(info, nullptr);
			EXPECT_EQ(info->type, type);
		}
	}

	EXPECT_GT(process_ids.size(), 0);
}

// ============================================================================
// 进程监控测试
// ============================================================================

/**
 * @brief 测试进程状态监控回调
 */
TEST_F(ProcessManagerTest, ProcessStateMonitoring) {
	auto& manager = process_manager::instance();

	// 启动进程并启用监控
	uint32_t process_id = manager.launch_process(
		"test_process_monitoring",
		test_executable_,
		{},
		sandbox_type::HOST,
		true // 启用监控
	);

	ASSERT_NE(process_id, 0);

	// 等待状态变化回调
	EXPECT_TRUE(g_callback_counter.wait_for_state_change(1, std::chrono::milliseconds(3000)));

	// 等待进程运行一段时间
	std::this_thread::sleep_for(std::chrono::milliseconds(500));

	// 终止进程
	auto result = manager.terminate_process(process_id);
	EXPECT_EQ(result, process_error::SUCCESS);

	// 等待更多状态变化
	std::this_thread::sleep_for(std::chrono::milliseconds(500));

	// 验证收到了回调
	EXPECT_GT(g_callback_counter.state_change_count.load(), 0);
}

/**
 * @brief 测试参数传递
 */
TEST_F(ProcessManagerTest, ProcessArguments) {
	auto& manager = process_manager::instance();

	// 传递参数
	std::vector<std::string> arguments = {"arg1", "arg2", "arg3"};

	uint32_t process_id = manager.launch_process(
		"test_process_args",
		test_executable_,
		arguments,
		sandbox_type::HOST,
		false
	);

	if (process_id != 0) {
		// 验证参数是否正确保存
		const process_info* info = manager.get_process_info(process_id);
		ASSERT_NE(info, nullptr);
		EXPECT_EQ(info->arguments, arguments);

		// 清理
		manager.terminate_process(process_id);
	}
}

// ============================================================================
// 测试主入口
// ============================================================================

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