# coding=utf-8
from unittest.mock import MagicMock, patch

import pytest

from spark_executor.core.yarn_client import (
    YarnError,
    get_application_logs,
    get_application_status,
    kill_application,
)


def _fake_proc(returncode: int, stdout: str = "", stderr: str = ""):
    p = MagicMock()
    p.returncode = returncode
    p.stdout = stdout
    p.stderr = stderr
    return p


def test_status_parses_state_line():
    fake = _fake_proc(
        0,
        stdout="Application Report :\n  State : RUNNING\n  ...\n",
    )
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
        state, raw = get_application_status("application_1")
    assert state == "RUNNING"
    assert "RUNNING" in raw


def test_status_raises_on_nonzero_return():
    fake = _fake_proc(1, stderr="not found")
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
        with pytest.raises(YarnError):
            get_application_status("application_x")


def test_logs_returns_stdout():
    fake = _fake_proc(0, stdout="log line 1\nlog line 2\n")
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
        out = get_application_logs("application_1")
    assert out == "log line 1\nlog line 2\n"
    args = m.call_args.args[0]
    assert args[:3] == ["yarn", "logs", "-applicationId"]
    assert args[3] == "application_1"


def test_kill_invokes_yarn_application_kill():
    fake = _fake_proc(0)
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
        kill_application("application_1")
    args = m.call_args.args[0]
    assert args[:3] == ["yarn", "application", "-kill"]
    assert args[3] == "application_1"


def test_kill_raises_on_nonzero_return():
    fake = _fake_proc(1, stderr="denied")
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
        with pytest.raises(YarnError):
            kill_application("application_1")