momentry_core/test_shutdown_improved.py

#!/opt/homebrew/bin/python3.11
"""
改进的关机测试脚本
测试进程检测和停止逻辑
"""

import os
import sys
import time
import subprocess
import psutil
from datetime import datetime


def find_processes(keywords):
    """查找包含关键字的进程"""
    processes = []
    for proc in psutil.process_iter(["pid", "name", "cmdline"]):
        try:
            cmdline = " ".join(proc.info["cmdline"]) if proc.info["cmdline"] else ""
            name = proc.info["name"] or ""

            for keyword in keywords:
                if keyword in cmdline or keyword in name:
                    processes.append(proc)
                    break
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
    return processes


def stop_process_gracefully(pid, timeout=5):
    """优雅地停止进程"""
    try:
        proc = psutil.Process(pid)

        # 发送 SIGTERM
        proc.terminate()

        # 等待
        try:
            proc.wait(timeout=timeout)
            return True
        except psutil.TimeoutExpired:
            # 发送 SIGKILL
            proc.kill()
            proc.wait(timeout=2)
            return True
    except psutil.NoSuchProcess:
        return True
    except Exception as e:
        print(f"停止进程 {pid} 失败: {e}")
        return False


def test_service_stopping(service_name, keywords, stop_commands=None):
    """测试服务停止"""
    print(f"\n测试停止 {service_name}...")

    # 查找进程
    processes = find_processes(keywords)
    print(f"  找到 {len(processes)} 个进程")

    if processes:
        for proc in processes:
            print(f"  PID {proc.pid}: {proc.info.get('name', 'unknown')}")

    # 执行停止命令
    if stop_commands:
        for cmd in stop_commands:
            print(f"  执行命令: {cmd}")
            try:
                result = subprocess.run(
                    cmd, shell=True, capture_output=True, text=True, timeout=10
                )
                print(f"    返回码: {result.returncode}")
                if result.stdout:
                    print(f"    输出: {result.stdout[:100]}...")
            except Exception as e:
                print(f"    错误: {e}")

    # 等待
    time.sleep(3)

    # 检查剩余进程
    remaining = find_processes(keywords)
    if remaining:
        print(f"  ❌ {service_name} 仍有 {len(remaining)} 个进程在运行")
        return False
    else:
        print(f"  ✅ {service_name} 已停止")
        return True


def main():
    print("=" * 60)
    print("改进的关机机制测试")
    print(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)

    # 测试的服务列表
    test_services = [
        {
            "name": "Redis",
            "keywords": ["redis-server"],
            "stop_commands": ["redis-cli shutdown"],
        },
        {
            "name": "PostgreSQL",
            "keywords": ["postgres"],
            "stop_commands": [
                "pg_ctl -D /Users/accusys/momentry/var/postgresql stop -m fast"
            ],
        },
        {
            "name": "AI 处理器",
            "keywords": ["asr_processor", "cut_processor"],
            "stop_commands": None,
        },
        {
            "name": "Momentry 服务",
            "keywords": ["momentry server", "momentry worker"],
            "stop_commands": None,
        },
    ]

    results = []

    # 运行测试
    for service in test_services:
        success = test_service_stopping(
            service["name"], service["keywords"], service.get("stop_commands")
        )
        results.append((service["name"], success))

    # 显示结果
    print("\n" + "=" * 60)
    print("测试结果")
    print("=" * 60)

    passed = 0
    failed = 0

    for service_name, success in results:
        if success:
            print(f"✅ {service_name}: 通过")
            passed += 1
        else:
            print(f"❌ {service_name}: 失败")
            failed += 1

    print(f"\n总计: {passed} 通过, {failed} 失败")

    # 建议
    if failed > 0:
        print("\n建议改进:")
        print("1. 增加停止命令的超时时间")
        print("2. 添加 sudo 权限支持")
        print("3. 实现进程树停止（父进程+子进程）")
        print("4. 添加强制停止选项（SIGKILL）")

    return passed == len(results)


if __name__ == "__main__":
    try:
        success = main()
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n测试被用户中断")
        sys.exit(130)
    except Exception as e:
        print(f"\n错误: {e}")
        sys.exit(1)