test_ssh.py

"""Tests for hytop.core.ssh.collect_from_host (subprocess mocked)."""

from __future__ import annotations

import subprocess
from unittest.mock import MagicMock, patch

from hytop.core.ssh import SSHOptions, collect_from_host, collect_python_from_host


def _make_proc(returncode=0, stdout="", stderr=""):
    m = MagicMock()
    m.returncode = returncode
    m.stdout = stdout
    m.stderr = stderr
    return m


class TestCollectFromHostLocal:
    @patch("hytop.core.ssh.subprocess.run")
    def test_success_returns_no_error(self, mock_run):
        mock_run.return_value = _make_proc(stdout='{"card0":{}}')
        result = collect_from_host(
            "localhost", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"]
        )
        assert result.error is None
        assert result.host == "localhost"

    @patch("hytop.core.ssh.subprocess.run")
    def test_local_invokes_hy_smi_directly(self, mock_run):
        mock_run.return_value = _make_proc()
        collect_from_host("localhost", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"])
        cmd = mock_run.call_args[0][0]
        assert cmd[0] == "hy-smi"
        assert "ssh" not in cmd

    @patch("hytop.core.ssh.subprocess.run")
    def test_127_0_0_1_treated_as_local(self, mock_run):
        mock_run.return_value = _make_proc()
        collect_from_host("127.0.0.1", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"])
        cmd = mock_run.call_args[0][0]
        assert cmd[0] == "hy-smi"

    @patch("hytop.core.ssh.subprocess.run")
    def test_nonzero_exit_returns_error(self, mock_run):
        mock_run.return_value = _make_proc(returncode=1, stderr="permission denied")
        result = collect_from_host(
            "localhost", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"]
        )
        assert result.error is not None
        assert "exit 1" in result.error

    @patch(
        "hytop.core.ssh.subprocess.run",
        side_effect=subprocess.TimeoutExpired("cmd", 10),
    )
    def test_timeout_returns_error(self, mock_run):
        result = collect_from_host(
            "localhost", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"]
        )
        assert result.error is not None
        assert "timeout" in result.error

    @patch("hytop.core.ssh.subprocess.run", side_effect=OSError("no such file"))
    def test_oserror_returns_error(self, mock_run):
        result = collect_from_host(
            "localhost", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"]
        )
        assert result.error is not None
        assert "no such file" in result.error


class TestCollectFromHostRemote:
    @patch("hytop.core.ssh.subprocess.run")
    def test_remote_uses_ssh(self, mock_run):
        mock_run.return_value = _make_proc(stdout="{}")
        collect_from_host("node01", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"])
        cmd = mock_run.call_args[0][0]
        assert cmd[0] == "ssh"

    @patch("hytop.core.ssh.subprocess.run")
    def test_remote_hostname_in_cmd(self, mock_run):
        mock_run.return_value = _make_proc(stdout="{}")
        collect_from_host("node01", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"])
        cmd = mock_run.call_args[0][0]
        assert "node01" in cmd

    @patch("hytop.core.ssh.subprocess.run")
    def test_remote_batch_mode_set(self, mock_run):
        mock_run.return_value = _make_proc(stdout="{}")
        collect_from_host("node01", ssh_timeout=5, cmd_timeout=10, hy_smi_args=["--json"])
        cmd = mock_run.call_args[0][0]
        assert "BatchMode=yes" in cmd

    @patch("hytop.core.ssh.subprocess.run")
    def test_hy_smi_args_forwarded(self, mock_run):
        mock_run.return_value = _make_proc(stdout="{}")
        collect_from_host(
            "node01",
            ssh_timeout=5,
            cmd_timeout=10,
            hy_smi_args=["--json", "--showtemp"],
        )
        cmd = mock_run.call_args[0][0]
        assert "--json" in cmd
        assert "--showtemp" in cmd


class TestCollectPythonFromHost:
    @patch("hytop.core.ssh.subprocess.run")
    def test_remote_uses_ssh(self, mock_run):
        mock_run.return_value = _make_proc(stdout='{"ok":1}')
        collect_python_from_host("node01", ssh_timeout=5, cmd_timeout=10, python_code="print('ok')")
        cmd = mock_run.call_args[0][0]
        assert cmd[0] == "ssh"
        assert any(str(part).startswith("python3 -c ") for part in cmd)

    @patch("hytop.core.ssh.subprocess.run")
    def test_mux_options_applied(self, mock_run):
        mock_run.return_value = _make_proc(stdout='{"ok":1}')
        collect_python_from_host(
            "node01",
            ssh_timeout=5,
            cmd_timeout=10,
            python_code="print('ok')",
            ssh_options=SSHOptions(use_mux=True),
        )
        cmd = mock_run.call_args[0][0]
        assert "ControlMaster=auto" in cmd
        assert any(str(item).startswith("ControlPersist=") for item in cmd)

    @patch("hytop.core.ssh.subprocess.run")
    def test_multiline_script_is_single_remote_arg(self, mock_run):
        mock_run.return_value = _make_proc(stdout='{"ok":1}')
        collect_python_from_host(
            "node01",
            ssh_timeout=5,
            cmd_timeout=10,
            python_code="import json\nprint(json.dumps({'ok': True}))",
        )
        cmd = mock_run.call_args[0][0]
        remote_parts = [
            part for part in cmd if isinstance(part, str) and part.startswith("python3 -c ")
        ]
        assert len(remote_parts) == 1