xwlb/tests/test_tunnel.py

"""隧道自愈回归测试（不依赖 pytest，不碰真实隧道）

    python tests/test_tunnel.py

覆盖：
- 端口探测 / 本机地址判定；
- 各分支状态：禁用、非本机、已可连、systemd 托管中等待、执行脚本后恢复、脚本不存在；
- **端口已通时绝不执行 autossh.sh**（避免在正常运行时多起进程抢端口）；
- systemd 单元 active 时只等待、不抢端口；
- dry-run 不执行脚本。

测试用临时 config.yml（`XWLB_CONFIG_FILE`）+ 临时端口的假脚本，
不会触碰真实的 13306 与项目里的 autossh.sh。
"""
import os
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))


def free_port():
    with socket.socket() as s:
        s.bind(('127.0.0.1', 0))
        return s.getsockname()[1]


def write_cfg(path, port, script, unit='xwlb-tunnel.service', enabled=1):
    path.write_text(f"""mysql:
  tunnel:
    enabled: {enabled}
    script: {script}
    systemd_unit: {unit}
    wait_seconds: 4
    connect_timeout: 1
  host: 127.0.0.1
  port: {port}
  user: myquant
  database: myquant
paths:
  video_dir: xwlb_video
  audio_dir: audio_processing
models:
  asr_model: paraformer-realtime-v2
  correct_model: qwen3.8-flash
  split_model: deepseek-flash
""", encoding='utf-8')


def make_fake_script(dirpath, port, marker):
    """假 autossh.sh：起一个 TCP 监听并把"被执行过"记录到 marker 文件"""
    p = dirpath / 'fake_tunnel.sh'
    p.write_text(f"""#!/bin/bash
touch "{marker}"
setsid .venv/bin/python -m http.server {port} --bind 127.0.0.1 >/dev/null 2>&1 &
""", encoding='utf-8')
    p.chmod(0o755)
    return p


def kill_listener(port):
    subprocess.run(['pkill', '-f', f'http[.]server {port}'], capture_output=True)
    time.sleep(0.3)


def main():
    failed = 0
    tmp = Path(tempfile.mkdtemp(prefix='xwlb_tunnel_test_'))

    def check(desc, got, want):
        nonlocal failed
        if got == want:
            print(f"  ✓ {desc}: {got!r}")
        else:
            failed += 1
            print(f"  ✗ {desc}: 期望 {want!r}，实际 {got!r}")

    # 临时配置 + 重新加载（用环境变量指定，不污染真实 config.yml）
    port = free_port()
    script = make_fake_script(tmp, port, tmp / 'executed')
    cfg = tmp / 'config.yml'
    write_cfg(cfg, port, str(script))
    os.environ['XWLB_CONFIG_FILE'] = str(cfg)

    import config
    config.reload_config()
    import tunnel

    try:
        print("端口探测:")
        check('关闭的端口判定为不通', tunnel.port_open('127.0.0.1', free_port(), timeout=0.5), False)
        srv = subprocess.Popen([sys.executable, '-m', 'http.server', str(port), '--bind', '127.0.0.1'],
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(1.0)
        check('已监听端口判定为可连', tunnel.port_open('127.0.0.1', port, timeout=1), True)
        srv.terminate()
        srv.wait(timeout=5)

        print("\n本机地址判定:")
        for h in ('localhost', '127.0.0.1', '::1'):
            check(f'{h} 视为本机', tunnel.is_local(h), True)
        check('远端地址不视为本机', tunnel.is_local('db.example.com'), False)

        print("\n分支：端口已通（不得执行脚本）:")
        if (tmp / 'executed').exists():
            os.remove(tmp / 'executed')
        srv = subprocess.Popen([sys.executable, '-m', 'http.server', str(port), '--bind', '127.0.0.1'],
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(1.0)
        r = tunnel.ensure_tunnel()
        check('状态 = 已可连', r['status'], tunnel.ST_ALREADY_OPEN)
        check('未执行隧道脚本', (tmp / 'executed').exists(), False)
        srv.terminate()
        srv.wait(timeout=5)

        print("\n分支：端口不通 + systemd 托管中（只等待，不抢端口）:")
        if (tmp / 'executed').exists():
            os.remove(tmp / 'executed')
        tunnel.systemd_unit_active = lambda unit, timeout=5: True
        r = tunnel.ensure_tunnel()
        tunnel.systemd_unit_active = lambda unit, timeout=5: False  # 还原
        check('状态 = 等待 systemd', r['status'], tunnel.ST_UNAVAILABLE)
        check('未执行隧道脚本（不抢 systemd 的端口）', (tmp / 'executed').exists(), False)

        print("\n分支：端口不通 + 执行脚本后恢复:")
        if (tmp / 'executed').exists():
            os.remove(tmp / 'executed')
        r = tunnel.ensure_tunnel(wait_seconds=6)
        check('状态 = 执行脚本后端口就绪', r['status'], tunnel.ST_SCRIPT_STARTED)
        check('脚本确实被执行', (tmp / 'executed').exists(), True)
        check('判定为已恢复', r['healed'], True)
        check('端口现在可连', tunnel.port_open('127.0.0.1', port), True)
        kill_listener(port)

        print("\n分支：脚本不存在 / 禁用 / 非本机:")
        write_cfg(cfg, free_port(), str(tmp / '不存在.sh'))
        config.reload_config()
        r = tunnel.ensure_tunnel()
        check('脚本不存在 → 不可用', r['status'], tunnel.ST_UNAVAILABLE)
        check('且未判定为已恢复', r['healed'], False)

        write_cfg(cfg, free_port(), str(script), enabled=0)
        config.reload_config()
        check('enabled=0 → 跳过探测', tunnel.ensure_tunnel()['status'], tunnel.ST_DISABLED)

        write_cfg(cfg, free_port(), str(script))
        config.reload_config()
        real_db_config = config.db_config
        config.db_config = lambda *a, **k: {'host': 'db.example.com', 'port': 3306,
                                            'username': 'u', 'password': 'p', 'database': 'd'}
        try:
            check('远端地址 → 不适用隧道', tunnel.ensure_tunnel()['status'], tunnel.ST_NOT_LOCAL)
        finally:
            config.db_config = real_db_config  # 必须还原，否则污染后续用例

        print("\n分支：dry-run 不执行脚本:")
        port2 = free_port()
        write_cfg(cfg, port2, str(script))
        config.reload_config()
        if (tmp / 'executed').exists():
            os.remove(tmp / 'executed')
        r = tunnel.ensure_tunnel(dry_run=True)
        check('dry-run 报告将执行脚本', 'dry-run' in r['detail'], True)
        check('dry-run 未真的执行', (tmp / 'executed').exists(), False)

    finally:
        os.environ.pop('XWLB_CONFIG_FILE', None)
        kill_listener(port)
        kill_listener(port)
        import shutil
        shutil.rmtree(tmp, ignore_errors=True)

    print(f"\n结果: {'全部通过' if failed == 0 else f'{failed} 项失败'}")
    return 1 if failed else 0


if __name__ == '__main__':
    sys.exit(main())