bambuddy/backend/tests/unit/services/test_spoolbuddy_ssh.py

"""Unit tests for SpoolBuddy SSH update service."""

import os
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from backend.app.services.spoolbuddy_ssh import (
    _get_ssh_key_dir,
    _run_ssh_command,
    detect_current_branch,
    get_or_create_keypair,
    get_public_key,
    perform_ssh_update,
)

# -- _get_ssh_key_dir ---------------------------------------------------------


def test_get_ssh_key_dir_creates_directory(tmp_path):
    with patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings:
        mock_settings.base_dir = tmp_path
        key_dir = _get_ssh_key_dir()
        assert key_dir == tmp_path / "spoolbuddy" / "ssh"
        assert key_dir.exists()


def test_get_ssh_key_dir_returns_existing(tmp_path):
    ssh_dir = tmp_path / "spoolbuddy" / "ssh"
    ssh_dir.mkdir(parents=True)
    with patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings:
        mock_settings.base_dir = tmp_path
        assert _get_ssh_key_dir() == ssh_dir


# -- get_or_create_keypair -----------------------------------------------------


@pytest.mark.asyncio
async def test_get_or_create_keypair_returns_existing(tmp_path):
    ssh_dir = tmp_path / "spoolbuddy" / "ssh"
    ssh_dir.mkdir(parents=True)
    priv = ssh_dir / "id_ed25519"
    pub = ssh_dir / "id_ed25519.pub"
    priv.write_text("PRIVATE")
    pub.write_text("PUBLIC")

    with patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings:
        mock_settings.base_dir = tmp_path
        result = await get_or_create_keypair()
        assert result == (priv, pub)


@pytest.mark.asyncio
async def test_get_or_create_keypair_generates_new(tmp_path):
    with patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings:
        mock_settings.base_dir = tmp_path

        ssh_dir = tmp_path / "spoolbuddy" / "ssh"

        async def fake_keygen(*args, **kwargs):
            # Simulate ssh-keygen creating the files
            ssh_dir.mkdir(parents=True, exist_ok=True)
            (ssh_dir / "id_ed25519").write_text("PRIVATE")
            (ssh_dir / "id_ed25519.pub").write_text("PUBLIC")
            mock_proc = AsyncMock()
            mock_proc.communicate = AsyncMock(return_value=(b"", b""))
            mock_proc.returncode = 0
            return mock_proc

        with patch("asyncio.create_subprocess_exec", side_effect=fake_keygen) as mock_exec:
            priv, pub = await get_or_create_keypair()

            mock_exec.assert_called_once()
            args = mock_exec.call_args[0]
            assert "ssh-keygen" in args
            assert "-t" in args
            assert "ed25519" in args


@pytest.mark.asyncio
async def test_get_or_create_keypair_raises_on_failure(tmp_path):
    with patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings:
        mock_settings.base_dir = tmp_path

        mock_proc = AsyncMock()
        mock_proc.communicate = AsyncMock(return_value=(b"", b"keygen error"))
        mock_proc.returncode = 1

        with (
            patch("asyncio.create_subprocess_exec", return_value=mock_proc),
            pytest.raises(RuntimeError, match="ssh-keygen failed"),
        ):
            await get_or_create_keypair()


# -- get_public_key ------------------------------------------------------------


@pytest.mark.asyncio
async def test_get_public_key(tmp_path):
    ssh_dir = tmp_path / "spoolbuddy" / "ssh"
    ssh_dir.mkdir(parents=True)
    (ssh_dir / "id_ed25519").write_text("PRIVATE")
    (ssh_dir / "id_ed25519.pub").write_text("ssh-ed25519 AAAA bambuddy-spoolbuddy\n")

    with patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings:
        mock_settings.base_dir = tmp_path
        key = await get_public_key()
        assert key == "ssh-ed25519 AAAA bambuddy-spoolbuddy"


# -- detect_current_branch ----------------------------------------------------


def test_detect_branch_from_git(tmp_path):
    (tmp_path / ".git").mkdir()
    with (
        patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings,
        patch("subprocess.run") as mock_run,
    ):
        mock_settings.base_dir = tmp_path
        mock_run.return_value = MagicMock(returncode=0, stdout="dev\n")
        assert detect_current_branch() == "dev"


def test_detect_branch_env_fallback(tmp_path):
    with (
        patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings,
        patch.dict(os.environ, {"GIT_BRANCH": "staging"}),
    ):
        mock_settings.base_dir = tmp_path
        assert detect_current_branch() == "staging"


def test_detect_branch_default_main(tmp_path):
    with (
        patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings,
        patch.dict(os.environ, {}, clear=True),
    ):
        mock_settings.base_dir = tmp_path
        # Remove GIT_BRANCH if present
        os.environ.pop("GIT_BRANCH", None)
        assert detect_current_branch() == "main"


# -- _run_ssh_command ----------------------------------------------------------


@pytest.mark.asyncio
async def test_run_ssh_command_success(tmp_path):
    key_file = tmp_path / "key"
    key_file.write_text("KEY")

    mock_proc = AsyncMock()
    mock_proc.communicate = AsyncMock(return_value=(b"hello\n", b""))
    mock_proc.returncode = 0

    with patch("asyncio.create_subprocess_exec", return_value=mock_proc) as mock_exec:
        rc, stdout, stderr = await _run_ssh_command("10.0.0.1", "echo hello", key_file)

    assert rc == 0
    assert stdout == "hello\n"
    assert stderr == ""
    args = mock_exec.call_args[0]
    assert "spoolbuddy@10.0.0.1" in args
    assert "echo hello" in args
    assert "BatchMode=yes" in args


@pytest.mark.asyncio
async def test_run_ssh_command_failure(tmp_path):
    key_file = tmp_path / "key"
    key_file.write_text("KEY")

    mock_proc = AsyncMock()
    mock_proc.communicate = AsyncMock(return_value=(b"", b"Connection refused"))
    mock_proc.returncode = 255

    with patch("asyncio.create_subprocess_exec", return_value=mock_proc):
        rc, stdout, stderr = await _run_ssh_command("10.0.0.1", "echo hello", key_file)

    assert rc == 255
    assert "Connection refused" in stderr


@pytest.mark.asyncio
async def test_run_ssh_command_timeout(tmp_path):
    key_file = tmp_path / "key"
    key_file.write_text("KEY")

    mock_proc = AsyncMock()
    mock_proc.communicate = AsyncMock(return_value=(b"", b""))
    mock_proc.kill = MagicMock()

    async def fake_wait_for(coro, timeout):
        # Consume the coroutine to avoid warning
        coro.close()
        raise TimeoutError

    with (
        patch("asyncio.create_subprocess_exec", return_value=mock_proc),
        patch("backend.app.services.spoolbuddy_ssh.asyncio.wait_for", side_effect=fake_wait_for),
    ):
        rc, stdout, stderr = await _run_ssh_command("10.0.0.1", "sleep 999", key_file, timeout=1)

    assert rc == -1
    assert "timed out" in stderr
    mock_proc.kill.assert_called_once()


# -- perform_ssh_update --------------------------------------------------------


def _make_update_mocks(tmp_path):
    """Create common mocks for perform_ssh_update tests."""
    mock_db_device = MagicMock()
    mock_db_device.update_status = None
    mock_db_device.update_message = None
    mock_db_device.pending_command = None

    mock_result = MagicMock()
    mock_result.scalar_one_or_none.return_value = mock_db_device

    mock_session = AsyncMock()
    mock_session.execute = AsyncMock(return_value=mock_result)
    mock_session.commit = AsyncMock()

    mock_ctx = AsyncMock()
    mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
    mock_ctx.__aexit__ = AsyncMock(return_value=False)

    mock_ws = MagicMock()
    mock_ws.broadcast = AsyncMock()

    return mock_db_device, mock_ctx, mock_ws


@pytest.mark.asyncio
async def test_perform_ssh_update_success(tmp_path):
    """Full update flow: all SSH commands succeed."""
    ssh_dir = tmp_path / "spoolbuddy" / "ssh"
    ssh_dir.mkdir(parents=True)
    (ssh_dir / "id_ed25519").write_text("PRIVATE")
    (ssh_dir / "id_ed25519.pub").write_text("PUBLIC")

    ssh_calls = []

    async def mock_ssh(ip, cmd, key, timeout=60):
        ssh_calls.append(cmd)
        return 0, "ok", ""

    _, mock_ctx, mock_ws = _make_update_mocks(tmp_path)

    with (
        patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings,
        patch("backend.app.services.spoolbuddy_ssh._run_ssh_command", side_effect=mock_ssh),
        patch("backend.app.services.spoolbuddy_ssh.detect_current_branch", return_value="dev"),
        patch("backend.app.core.database.async_session", return_value=mock_ctx),
        patch("backend.app.api.routes.spoolbuddy.ws_manager", mock_ws),
    ):
        mock_settings.base_dir = tmp_path
        await perform_ssh_update("sb-test", "10.0.0.1")

    # Should have run: echo ok, git fetch, git checkout+reset, pip install,
    # systemctl restart, find (SW cleanup), systemctl restart getty
    assert len(ssh_calls) == 7
    assert "echo ok" in ssh_calls[0]
    assert "fetch" in ssh_calls[1]
    assert "checkout" in ssh_calls[2]
    assert "pip" in ssh_calls[3]
    assert "spoolbuddy.service" in ssh_calls[4]
    assert "Service Worker" in ssh_calls[5]
    assert "getty" in ssh_calls[6]

    assert mock_ws.broadcast.call_count >= 4


@pytest.mark.asyncio
async def test_perform_ssh_update_ssh_failure(tmp_path):
    """SSH connectivity check fails — should set error status."""
    ssh_dir = tmp_path / "spoolbuddy" / "ssh"
    ssh_dir.mkdir(parents=True)
    (ssh_dir / "id_ed25519").write_text("PRIVATE")
    (ssh_dir / "id_ed25519.pub").write_text("PUBLIC")

    async def mock_ssh(ip, cmd, key, timeout=60):
        if "echo ok" in cmd:
            return 255, "", "Connection refused"
        return 0, "", ""

    mock_device, mock_ctx, mock_ws = _make_update_mocks(tmp_path)

    with (
        patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings,
        patch("backend.app.services.spoolbuddy_ssh._run_ssh_command", side_effect=mock_ssh),
        patch("backend.app.services.spoolbuddy_ssh.detect_current_branch", return_value="main"),
        patch("backend.app.core.database.async_session", return_value=mock_ctx),
        patch("backend.app.api.routes.spoolbuddy.ws_manager", mock_ws),
    ):
        mock_settings.base_dir = tmp_path
        await perform_ssh_update("sb-test", "10.0.0.1")

    # Should broadcast error status
    error_broadcasts = [c for c in mock_ws.broadcast.call_args_list if c[0][0].get("update_status") == "error"]
    assert len(error_broadcasts) >= 1
    assert "SSH connection failed" in error_broadcasts[0][0][0]["update_message"]


@pytest.mark.asyncio
async def test_perform_ssh_update_git_fetch_failure(tmp_path):
    """Git fetch fails — should set error and stop."""
    ssh_dir = tmp_path / "spoolbuddy" / "ssh"
    ssh_dir.mkdir(parents=True)
    (ssh_dir / "id_ed25519").write_text("PRIVATE")
    (ssh_dir / "id_ed25519.pub").write_text("PUBLIC")

    ssh_calls = []

    async def mock_ssh(ip, cmd, key, timeout=60):
        ssh_calls.append(cmd)
        if "fetch" in cmd:
            return 1, "", "fatal: could not read from remote"
        return 0, "ok", ""

    _, mock_ctx, mock_ws = _make_update_mocks(tmp_path)

    with (
        patch("backend.app.services.spoolbuddy_ssh.settings") as mock_settings,
        patch("backend.app.services.spoolbuddy_ssh._run_ssh_command", side_effect=mock_ssh),
        patch("backend.app.services.spoolbuddy_ssh.detect_current_branch", return_value="main"),
        patch("backend.app.core.database.async_session", return_value=mock_ctx),
        patch("backend.app.api.routes.spoolbuddy.ws_manager", mock_ws),
    ):
        mock_settings.base_dir = tmp_path
        await perform_ssh_update("sb-test", "10.0.0.1")

    # Should stop after git fetch — no checkout, pip, restart
    assert len(ssh_calls) == 2  # echo ok + git fetch
    assert not any("checkout" in c for c in ssh_calls)