sshclient/test_core.py

"""
核心模块 smoke test
跳过需要真实 SSH 服务的部分，验证：
- SSHConnection 实例化、错误处理
- SystemMonitor 解析逻辑（用伪造输出）
- AIAgent 工具注册正确
- ConnectionManager 增删改
"""
import sys
import json
from pathlib import Path

# 允许在 venv 之外运行
sys.path.insert(0, str(Path(__file__).parent))

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


def test_ssh_connection_failure():
    """连不通的地址应返回失败而不是抛异常"""
    c = SSHConnection("127.0.0.1", 1, "nope", "nope", timeout=2)
    ok, msg = c.connect()
    assert ok is False
    assert "失败" in msg or "认证" in msg
    print(f"  ✓ 连接失败正确处理: {msg[:60]}")


def test_monitor_parse():
    """SystemMonitor 解析伪造的脚本输出"""
    fake_out = """
===CPU===
CPU_USAGE=37
CPU_CORES=8
LOAD=0.50 0.40 0.30
UPTIME=86400
===MEM===
MEM_TOTAL=16777216
MEM_USED=8388608
MEM_AVAIL=8388608
SWAP_TOTAL=2097152
SWAP_USED=1048576
===DISK===
DISK|/dev/sda1|107374182400|53687091200|50%
DISK|/dev/sda2|536870912000|268435456000|50%
===NET===
NET|eth0|1234567|7654321
NET|eth1|100|200
===HOST===
HOSTNAME=test-server
KERNEL=5.15.0-test
OS=Ubuntu 22.04
"""
    # 用 monkey-patch 模拟 exec_command
    c = SSHConnection("dummy", 22, "u", "p")
    c.connected = True
    c.exec_command = lambda cmd, timeout=10: (0, fake_out, "")

    m = SystemMonitor.collect(c)
    assert m["cpu"] == 37.0, m["cpu"]
    assert m["cores"] == 8
    assert m["load1"] == 0.5
    assert m["mem_total"] == 16777216
    assert m["mem_percent"] == 50.0
    assert len(m["disks"]) == 2
    assert m["disks"][0]["mount"] == "/dev/sda1"
    assert m["disks"][0]["percent"] == 50
    assert len(m["net"]) == 2
    assert m["net"][0]["iface"] == "eth0"
    assert m["net"][0]["rx"] == 1234567
    assert m["hostname"] == "test-server"
    assert m["os"] == "Ubuntu 22.04"
    print(f"  ✓ 监控解析: CPU={m['cpu']}% MEM={m['mem_percent']:.1f}% disks={len(m['disks'])} nets={len(m['net'])}")


def test_ai_agent_tools():
    """AI Agent 工具注册"""
    assert len(TOOL_DEFINITIONS) >= 4
    names = {t["function"]["name"] for t in TOOL_DEFINITIONS}
    assert "exec_ssh_command" in names
    assert "get_system_metrics" in names
    assert "list_remote_files" in names
    assert "read_remote_file" in names
    assert "upload_local_file" in names
    a = AIAgent(api_key="test", model="gpt-4o-mini")
    assert a.api_key == "test"
    assert len(a.history) == 0
    print(f"  ✓ AI Agent 注册 {len(TOOL_DEFINITIONS)} 个工具")


def test_ai_agent_missing_key():
    """没配 API Key 时应抛出友好错误"""
    a = AIAgent(api_key="", model="gpt-4o")
    try:
        a.chat("hello", lambda n, a: "ok")
        assert False, "应抛出异常"
    except RuntimeError as e:
        assert "API Key" in str(e)
        print(f"  ✓ AI 缺 Key 抛错: {e}")


def test_manager(tmpdir=None):
    """管理器增删改查（用临时配置目录）"""
    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": "test", "host": "1.2.3.4", "port": 22,
                             "username": "u", "password": "p"})
        assert new_id
        assert len(m.hosts) == 1
        m.update_host(new_id, {"name": "renamed", "host": "5.6.7.8",
                               "port": 2222, "username": "u", "password": "p"})
        h = m.get_host(new_id)
        assert h["name"] == "renamed"
        assert h["host"] == "5.6.7.8"
        m.remove_host(new_id)
        assert m.get_host(new_id) is None
        print("  ✓ 管理器 CRUD 正常")


def test_format():
    """格式化函数"""
    assert SystemMonitor.format_bytes(1024) == "1.0KB"
    assert SystemMonitor.format_bytes(1024**3) == "1.0GB"
    assert "1天" in SystemMonitor.format_uptime(90000)
    assert "3分" in SystemMonitor.format_uptime(200)
    print("  ✓ 字节/时长格式化正确")


if __name__ == "__main__":
    print("[1/6] SSHConnection 失败处理")
    test_ssh_connection_failure()
    print("[2/6] SystemMonitor 解析")
    test_monitor_parse()
    print("[3/6] AI Agent 工具注册")
    test_ai_agent_tools()
    print("[4/6] AI Agent 缺 Key")
    test_ai_agent_missing_key()
    print("[5/6] ConnectionManager CRUD")
    test_manager()
    print("[6/6] 格式化函数")
    test_format()
    print("\n所有核心模块 smoke test 通过 ✓")