mcp-server/tests/unit/test_yarn_client.py

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

import httpx
import pytest

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


RM = "http://rm:8088"


@pytest.fixture(autouse=True)
def _restore_settings():
    """Each test gets a clean copy of `settings` so env-var-style overrides
    in one test don't leak into the next."""
    snapshot = config.Settings(
        data_dir=config.settings.data_dir,
        jobs_dir=config.settings.jobs_dir,
        yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
        log_level=config.settings.log_level,
    )
    yield
    config.settings.data_dir = snapshot.data_dir
    config.settings.jobs_dir = snapshot.jobs_dir
    config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
    config.settings.log_level = snapshot.log_level


def _resp(status: int, *, json_data=None, text: str | None = None) -> httpx.Response:
    if json_data is not None:
        return httpx.Response(status, json=json_data)
    return httpx.Response(status, text=text or "")


# --- get_application_status ---

def test_status_parses_app_state():
    fake = _resp(200, json_data={"app": {"id": "application_1", "state": "RUNNING"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        state, raw = get_application_status("application_1", RM)
    assert state == "RUNNING"
    assert "RUNNING" in raw
    args = m.call_args.args
    assert args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1")


def test_status_raises_on_404():
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=_resp(404)):
        with pytest.raises(YarnError, match="not found"):
            get_application_status("application_x", RM)


def test_status_raises_on_5xx():
    with patch(
        "spark_executor.core.yarn_client.httpx.request",
        return_value=_resp(503, text="upstream down"),
    ):
        with pytest.raises(YarnError, match="503"):
            get_application_status("application_1", RM)


def test_status_raises_when_state_field_missing():
    fake = _resp(200, json_data={"app": {"id": "application_1"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake):
        with pytest.raises(YarnError, match="Could not parse YARN state"):
            get_application_status("application_1", RM)


def test_status_requires_rm_url():
    config.settings.yarn_resource_manager_url = None
    with pytest.raises(YarnConfigError):
        get_application_status("application_1", None)


def test_status_falls_back_to_settings_yarn_rm_url():
    config.settings.yarn_resource_manager_url = "http://env-rm:8088"
    fake = _resp(200, json_data={"app": {"state": "FINISHED"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        state, _ = get_application_status("application_1", None)
    assert state == "FINISHED"
    assert "env-rm:8088" in m.call_args.args[1]


def test_status_rejects_non_http_url():
    with pytest.raises(YarnConfigError, match="must start with"):
        get_application_status("application_1", "rm:8088")


# --- get_application_logs ---

def test_logs_returns_text():
    fake = _resp(200, text="log line 1\nlog line 2\n")
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        out = get_application_logs("application_1", RM)
    assert out == "log line 1\nlog line 2\n"
    assert m.call_args.args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1/aggregated-logs")


def test_logs_raises_on_404_with_explanation():
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=_resp(404)):
        with pytest.raises(YarnError, match="log-aggregation-enable"):
            get_application_logs("application_1", RM)


def test_logs_raises_on_5xx():
    with patch(
        "spark_executor.core.yarn_client.httpx.request",
        return_value=_resp(500, text="boom"),
    ):
        with pytest.raises(YarnError):
            get_application_logs("application_1", RM)


# --- kill_application ---

def test_kill_sends_put_with_killed_state():
    fake = _resp(200, json_data={"app": {"state": "KILLED"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        kill_application("application_1", RM)
    args = m.call_args.args
    assert args == ("PUT", f"{RM}/ws/v1/cluster/apps/application_1/state")
    assert m.call_args.kwargs["json"] == {"state": "KILLED"}


def test_kill_raises_on_5xx():
    with patch(
        "spark_executor.core.yarn_client.httpx.request",
        return_value=_resp(403, text="forbidden"),
    ):
        with pytest.raises(YarnError, match="403"):
            kill_application("application_1", RM)


# --- connection errors ---

def test_status_wraps_httpx_errors_as_yarn_error():
    with patch(
        "spark_executor.core.yarn_client.httpx.request",
        side_effect=httpx.ConnectError("connection refused"),
    ):
        with pytest.raises(YarnError, match="connection failed"):
            get_application_status("application_1", RM)