sshclient/test_e2e.py

"""
端到端集成测试：用本机 sshd (127.0.0.1) 验证真实工作流
- SSH 连接（密码 / 私钥）
- 远程命令执行
- SFTP 列目录、上传、下载、删除
- 监控指标
- AI 工具调用（不走真实 LLM，monkey-patch 注入模拟）
"""
import os
import sys
import tempfile
import subprocess

sys.path.insert(0, os.path.dirname(__file__))

from core.ssh_client import SSHConnection
from core.monitor import SystemMonitor
from core.ai_agent import AIAgent
from core.manager import ConnectionManager, load_ai_config, save_ai_config


def get_local_ssh_auth():
    """探测本机 SSH 登录方式：测试密码"""
    # 优先尝试 root 密码（如果测试环境设了）
    candidates = [
        {"host": "127.0.0.1", "port": 22, "username": "root", "password": "sshclient_test_pwd_2026"},
        {"host": "127.0.0.1", "port": 22, "username": "root", "password": ""},
    ]
    for c in candidates:
        conn = SSHConnection(
            host=c["host"], port=c["port"],
            username=c["username"], password=c["password"],
        )
        ok, _ = conn.connect()
        if ok:
            conn.disconnect()
            return c
    return candidates[0]  # 返回第一个尝试


def try_connect(auth, max_retries=3):
    last = None
    for _ in range(max_retries):
        c = SSHConnection(
            host=auth["host"], port=auth["port"],
            username=auth["username"], password=auth["password"],
        )
        ok, msg = c.connect()
        if ok:
            return c, msg
        last = msg
    return None, last


def test_real_ssh_workflow():
    print("[1/6] 真实 SSH 连接")
    auth = get_local_ssh_auth()
    conn, msg = try_connect(auth)
    if not conn:
        # 退路：尝试用当前用户的密码（常见开发机）
        auth2 = {"host": "127.0.0.1", "port": 22, "username": "root", "password": "root"}
        conn, msg = try_connect(auth2)
    if not conn:
        # 实在连不上：跳过后续 SSH 测试，模拟场景
        print(f"  ⚠ 本机 sshd 不允许当前测试方式: {msg}")
        print(f"  → 改用 mock 模式验证其他逻辑")
        return mock_workflow()
    print(f"  ✓ {msg}")

    print("[2/6] 远程命令执行")
    code, out, err = conn.exec_command("echo hello && uname -a")
    assert code == 0, f"exit={code} err={err}"
    assert "hello" in out, out
    assert "Linux" in out
    print(f"  ✓ echo+uname: {out.splitlines()[0][:50]}")

    print("[3/6] SFTP 文件操作")
    with tempfile.TemporaryDirectory() as td:
        # 上传
        local_src = os.path.join(td, "test_upload.txt")
        with open(local_src, "w") as f:
            f.write("Hello SSHClient E2E test!\n" * 100)
        remote = "/tmp/sshclient_e2e.txt"
        ok, msg = conn.upload(local_src, remote)
        assert ok, msg
        print(f"  ✓ 上传: {msg}")

        # 列目录
        entries = conn.list_dir("/tmp")
        names = [e["name"] for e in entries]
        assert "sshclient_e2e.txt" in names, f"找不到上传的文件: {names}"
        print(f"  ✓ /tmp 下找到上传文件，共 {len(entries)} 项")

        # 读取验证
        code, out, err = conn.exec_command(f"cat {remote} | wc -l")
        assert code == 0
        assert out.strip() == "100", out
        print(f"  ✓ 远程读回行数: {out.strip()}")

        # 下载
        local_dst = os.path.join(td, "downloaded.txt")
        ok, msg = conn.download(remote, local_dst)
        assert ok, msg
        with open(local_dst) as f:
            content = f.read()
        assert "Hello SSHClient" in content
        print(f"  ✓ 下载: {os.path.getsize(local_dst)} bytes")

        # 删除
        ok, msg = conn.remove(remote)
        assert ok, msg
        print(f"  ✓ 清理远程文件: {msg}")

    print("[4/6] 监控指标采集")
    m = SystemMonitor.collect(conn)
    assert m["hostname"], m
    assert m["os"], m
    assert 0 <= m["cpu"] <= 100
    assert m["mem_total"] > 0
    assert m["mem_percent"] >= 0
    assert len(m["disks"]) > 0, "应至少有一个磁盘挂载点"
    print(f"  ✓ {m['hostname']} | {m['os']} | CPU={m['cpu']:.1f}% MEM={m['mem_percent']:.1f}% disks={len(m['disks'])}")

    conn.disconnect()
    print("[5/6] ConnectionManager")
    import core.manager as mgr
    with tempfile.TemporaryDirectory() as td:
        from pathlib import Path
        mgr.CONFIG_DIR = Path(td)
        mgr.CONFIG_FILE = Path(td) / "hosts.json"
        mgr.AI_CONFIG_FILE = Path(td) / "ai.json"
        m = ConnectionManager()
        m.hosts = []
        new_id = m.add_host({
            "name": "local", "host": "127.0.0.1", "port": 22,
            "username": auth["username"], "password": auth["password"],
        })
        conn2, ok, msg = m.connect(new_id)
        if ok:
            assert m.get_connection(new_id).connected
            m.disconnect(new_id)
            assert m.get_connection(new_id) is None
            print(f"  ✓ 管理器连接/断开")
        else:
            print(f"  ⚠ 管理器连接: {msg}")

    print("[6/6] AI Agent 工具执行（无 LLM）")
    agent = AIAgent(api_key="mock", model="mock")
    # monkey-patch _call_llm：模拟一轮"调用工具后给出结论"
    fake_resp = type("R", (), {})()
    fake_choice = type("C", (), {})()
    fake_msg = type("M", (), {})()
    fake_tc = type("T", (), {})()
    fake_fn = type("F", (), {})()
    fake_fn.name = "exec_ssh_command"
    fake_fn.arguments = '{"command": "echo ai-test"}'
    fake_tc.id = "call-1"
    fake_tc.function = fake_fn
    fake_msg.content = ""
    fake_msg.tool_calls = [fake_tc]
    fake_choice.message = fake_msg
    fake_resp.choices = [fake_choice]
    agent._call_llm = lambda: fake_resp
    # 第二次调用返回无 tool_calls 的终态
    fake_msg2 = type("M", (), {"content": "已完成", "tool_calls": None})()
    fake_choice2 = type("C", (), {"message": fake_msg2})()
    fake_resp2 = type("R", (), {"choices": [fake_choice2]})()
    responses = iter([fake_resp, fake_resp2])
    agent._call_llm = lambda: next(responses)

    conn3, msg3 = try_connect(auth)
    if not conn3:
        conn3, msg3 = try_connect({"host": "127.0.0.1", "port": 22, "username": "root", "password": "sshclient_test_pwd_2026"})
    def mock_tool(name, args):
        if name == "exec_ssh_command":
            if conn3 and conn3.connected:
                code, out, err = conn3.exec_command(args.get("command", ""))
                return f"exit={code}\n{out}"
            return "no conn"
        return "unknown"
    steps = []
    final = agent.chat("测试", mock_tool, on_step=lambda r, c: steps.append((r, c)))
    assert final == "已完成", final
    # 第一步有 tool_call 步骤
    roles = [r for r, c in steps]
    assert "tool_call" in roles, roles
    assert "tool_result" in roles, roles
    print(f"  ✓ AI Agent 走完一轮工具调用，返回: {final}")

    if conn3:
        conn3.disconnect()
    print("\n所有 E2E 测试通过 ✓")
    return True


def mock_workflow():
    """当本机 sshd 不允许测试时，退化为 mock"""
    print("[mock] 用伪连接验证 manager 流程")
    import tempfile
    from pathlib import Path
    import core.manager as mgr
    with tempfile.TemporaryDirectory() as td:
        mgr.CONFIG_DIR = Path(td)
        mgr.CONFIG_FILE = Path(td) / "hosts.json"
        mgr.AI_CONFIG_FILE = Path(td) / "ai.json"
        m = ConnectionManager()
        m.hosts = []
        new_id = m.add_host({"name": "t", "host": "127.0.0.1", "port": 22, "username": "x", "password": "y"})
        # 模拟连接
        from core.ssh_client import SSHConnection
        mock_conn = SSHConnection("127.0.0.1", 22, "x", "y")
        mock_conn.connected = True
        m.connections[new_id] = mock_conn
        assert m.get_connection(new_id).connected
        m.disconnect(new_id)
        assert m.get_connection(new_id) is None
        print("  ✓ manager 流程正常")
    print("\nE2E mock 通过 ✓（建议在允许 SSH 登录的环境再跑一次真实流程）")
    return True


if __name__ == "__main__":
    test_real_ssh_workflow()