from unittest.mock import AsyncMock, patch

import pytest

from dobby.core.gpu_monitor import GPUMonitor


@pytest.fixture
def monitor():
    return GPUMonitor(poll_interval=1)


async def test_check_pid_running(monitor):
    mock_proc = AsyncMock()
    mock_proc.communicate = AsyncMock(return_value=(b"12345\n", b""))
    mock_proc.returncode = 0

    with patch("asyncio.create_subprocess_exec", return_value=mock_proc):
        running = await monitor.check_pid("user@gpu", 12345)
    assert running is True


async def test_check_pid_completed(monitor):
    mock_proc = AsyncMock()
    mock_proc.communicate = AsyncMock(return_value=(b"", b""))
    mock_proc.returncode = 1

    with patch("asyncio.create_subprocess_exec", return_value=mock_proc):
        running = await monitor.check_pid("user@gpu", 12345)
    assert running is False


async def test_fetch_log(monitor):
    mock_proc = AsyncMock()
    mock_proc.communicate = AsyncMock(return_value=(b"Epoch 3 loss=0.34", b""))
    mock_proc.returncode = 0

    with patch("asyncio.create_subprocess_exec", return_value=mock_proc):
        log = await monitor.fetch_log("user@gpu", "/home/user/train.log")
    assert "loss=0.34" in log


async def test_monitor_job_until_complete(monitor):
    call_count = 0
    async def mock_check_pid(host, pid):
        nonlocal call_count
        call_count += 1
        return call_count < 3

    monitor.check_pid = mock_check_pid
    monitor.fetch_log = AsyncMock(return_value="Training complete. Loss=0.5")

    completed = await monitor.wait_for_completion("user@gpu", 12345, "/tmp/log")
    assert completed is True
    assert call_count == 3
