mcp-server/tests/unit/test_yarn_client.py

# coding=utf-8
import base64

from unittest.mock import patch

import httpx
import pytest

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


RM = "http://rm:8088"
HTTPS_RM = "https://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,
        ssl_verify_default=config.settings.ssl_verify_default,
        ssl_ca_bundle_default=config.settings.ssl_ca_bundle_default,
    )
    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
    config.settings.ssl_verify_default = snapshot.ssl_verify_default
    config.settings.ssl_ca_bundle_default = snapshot.ssl_ca_bundle_default


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", YarnClientConfig(yarn_rm_url=RM))
    assert state == "RUNNING"
    assert "RUNNING" in raw
    args = m.call_args.args
    assert args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1")
    assert m.call_args.kwargs["verify"] is True


def test_request_sends_accept_json_header():
    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:
        get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM))
    assert m.call_args.kwargs["headers"] == {"Accept": "application/json"}


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", YarnClientConfig(yarn_rm_url=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", YarnClientConfig(yarn_rm_url=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", YarnClientConfig(yarn_rm_url=RM))


def test_status_requires_rm_url():
    config.settings.yarn_resource_manager_url = None
    with pytest.raises(YarnConfigError):
        get_application_status("application_1", YarnClientConfig(yarn_rm_url=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", YarnClientConfig(yarn_rm_url=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", YarnClientConfig(yarn_rm_url="rm:8088"))


# --- get_application_logs ---

def test_logs_returns_text_on_h3_aggregated_endpoint():
    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", YarnClientConfig(yarn_rm_url=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",
    )
    assert m.call_args.kwargs["verify"] is True


def test_logs_falls_back_to_am_container_on_404():
    responses = [
        _resp(404),  # aggregated-logs endpoint missing
        _resp(
            200,
            json_data={
                "app": {
                    "amContainerLogs": "http://nm1:8042/node/containerlogs/container_1/hdfs",
                }
            },
        ),
        _resp(200, text="am container log content"),
    ]
    with patch(
        "spark_executor.core.yarn_client.httpx.request", side_effect=responses
    ) as m:
        out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
    assert out == "am container log content"
    calls = m.call_args_list
    assert calls[0].args == (
        "GET",
        f"{RM}/ws/v1/cluster/apps/application_1/aggregated-logs",
    )
    assert calls[1].args == ("GET", f"{RM}/ws/v1/cluster/apps/application_1")
    assert calls[2].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs/stdout",
    )


def test_logs_raises_on_404_when_am_container_field_missing():
    responses = [
        _resp(404),  # aggregated-logs endpoint missing
        _resp(200, json_data={"app": {}}),
    ]
    with patch(
        "spark_executor.core.yarn_client.httpx.request", side_effect=responses
    ):
        with pytest.raises(YarnError, match="log-aggregation-enable"):
            get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))


def test_logs_returns_am_container_stdout():
    responses = [
        _resp(404),  # aggregated-logs endpoint missing
        _resp(
            200,
            json_data={
                "app": {
                    "amContainerLogs": "http://nm1:8042/node/containerlogs/container_1/hdfs",
                }
            },
        ),
        _resp(200, text="driver stdout"),
    ]
    with patch(
        "spark_executor.core.yarn_client.httpx.request", side_effect=responses
    ) as m:
        out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
    assert out == "driver stdout"
    assert m.call_args_list[2].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs/stdout",
    )


def test_logs_5xx_on_aggregated_endpoint_raises_immediately():
    with patch(
        "spark_executor.core.yarn_client.httpx.request",
        return_value=_resp(500, text="boom"),
    ) as m:
        with pytest.raises(YarnError, match="500"):
            get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))
    assert m.call_count == 1


def test_logs_falls_back_to_directory_listing_when_stdout_returns_html():
    """Knox / wrapped YARN: GET {amContainerLogs}/stdout returns an HTML
    directory listing instead of raw log bytes. The fallback should fetch the
    bare URL, parse the listing, and pull each log file individually."""
    listing = (
        "<html><body>"
        '<a href="../">Parent Directory</a>'
        '<a href="?C=N;O=D">Name</a>'
        '<a href="stdout">stdout</a>'
        '<a href="stderr">stderr</a>'
        '<a href="syslog">syslog</a>'
        "</body></html>"
    )
    html_stdout = (
        "<!doctype html><html><body>wrapper around /stdout endpoint</body></html>"
    )
    responses = [
        _resp(404),  # aggregated-logs missing
        _resp(
            200,
            json_data={
                "app": {
                    "amContainerLogs": "http://nm1:8042/node/containerlogs/container_1/hdfs",
                }
            },
        ),
        _resp(200, text=html_stdout),                       # /stdout -> HTML wrapper
        _resp(200, text=listing),                           # bare URL -> directory listing
        _resp(200, text="driver stdout line\n"),            # fetch stdout
        _resp(200, text="driver stderr line\n"),            # fetch stderr
        _resp(200, text="syslog line\n"),                   # fetch syslog
    ]
    with patch(
        "spark_executor.core.yarn_client.httpx.request", side_effect=responses
    ) as m:
        out = get_application_logs("application_1", YarnClientConfig(yarn_rm_url=RM))

    # Header + body for each file in document order.
    assert "=== stdout ===" in out
    assert "=== stderr ===" in out
    assert "=== syslog ===" in out
    assert "driver stdout line" in out
    assert "driver stderr line" in out
    assert "syslog line" in out
    # stdout must come before stderr (document order preserved).
    assert out.index("=== stdout ===") < out.index("=== stderr ===") < out.index("=== syslog ===")

    # 7 HTTP calls: aggregated-logs, app, /stdout (HTML), bare URL, stdout file, stderr file, syslog file
    assert m.call_count == 7
    assert m.call_args_list[2].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs/stdout",
    )
    assert m.call_args_list[3].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs",
    )
    assert m.call_args_list[4].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs/stdout",
    )
    assert m.call_args_list[5].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs/stderr",
    )
    assert m.call_args_list[6].args == (
        "GET",
        "http://nm1:8042/node/containerlogs/container_1/hdfs/syslog",
    )


def test_logs_html_filter_strips_sort_links_and_parent_dir():
    """_parse_log_listing_html must drop ../, sort-toggle query strings, and
    absolute URLs, while keeping plain log-file names."""
    html = (
        '<a href="../">Parent</a>'
        '<a href="?C=N;O=D">Name</a>'
        '<a href="?C=M;O=A">Last modified</a>'
        '<a href="?C=S;O=A">Size</a>'
        '<a href="?C=D;O=A">Description</a>'
        '<a href="http://example.com/somewhere">absolute</a>'
        '<a href="/absolute/path">rooted</a>'
        '<a href="stdout">stdout</a>'
        '<a href="stderr">stderr</a>'
        '<a href="syslog">syslog</a>'
        '<a href="stdout">stdout-dup</a>'  # duplicate
    )
    assert _parse_log_listing_html(html) == ["stdout", "stderr", "syslog"]


def test_logs_html_filter_handles_prelaunch_files():
    """Some YARN versions list prelaunch.out / prelaunch.err / launch_container.sh
    in the directory listing alongside stdout/stderr. Keep them — we want to
    surface every log file the container produced."""
    html = (
        '<a href="prelaunch.out">prelaunch.out</a>'
        '<a href="prelaunch.err">prelaunch.err</a>'
        '<a href="launch_container.sh">launch_container.sh</a>'
        '<a href="directory-info">directory-info</a>'
    )
    assert _parse_log_listing_html(html) == [
        "prelaunch.out",
        "prelaunch.err",
        "launch_container.sh",
        "directory-info",
    ]


def test_logs_raises_when_directory_listing_empty():
    """If /stdout returns HTML AND the directory listing has no parseable log
    file names, raise _logs_unavailable_error."""
    html_stdout = "<!doctype html><html>wrapper</html>"
    empty_listing = '<html><body><a href="../">Parent</a></body></html>'
    responses = [
        _resp(404),                                  # aggregated-logs missing
        _resp(
            200,
            json_data={
                "app": {
                    "amContainerLogs": "http://nm1:8042/node/containerlogs/container_1/hdfs",
                }
            },
        ),
        _resp(200, text=html_stdout),                # /stdout -> HTML
        _resp(200, text=empty_listing),              # bare URL -> only parent link
    ]
    with patch(
        "spark_executor.core.yarn_client.httpx.request", side_effect=responses
    ):
        with pytest.raises(YarnError, match="log-aggregation-enable"):
            get_application_logs("application_1", YarnClientConfig(yarn_rm_url=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", YarnClientConfig(yarn_rm_url=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"}
    assert m.call_args.kwargs["verify"] is True


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", YarnClientConfig(yarn_rm_url=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", YarnClientConfig(yarn_rm_url=RM))


# --- SSL/TLS configuration ---

def test_request_passes_ssl_verify_true():
    config.settings.ssl_verify_default = True
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status(
            "application_1",
            YarnClientConfig(yarn_rm_url=RM, ssl_verify=True),
        )
    assert m.call_args.kwargs["verify"] is True


def test_request_passes_ssl_ca_bundle_path():
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status(
            "application_1",
            YarnClientConfig(
                yarn_rm_url=HTTPS_RM,
                ssl_ca_bundle="/path/to/ca.crt",
            ),
        )
    assert m.call_args.kwargs["verify"] == "/path/to/ca.crt"


def test_request_uses_global_ssl_default_when_connection_unset():
    config.settings.ssl_verify_default = True
    config.settings.ssl_ca_bundle_default = "/global/ca.crt"
    config.settings.yarn_resource_manager_url = "https://rm:8088"
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    conn = Connection(name="prod", master="yarn", ssl_verify=None, ssl_ca_bundle=None)
    cfg = YarnClientConfig.from_connection(conn)
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", cfg)
    assert m.call_args.kwargs["verify"] == "/global/ca.crt"


def test_request_per_connection_ssl_overrides_global():
    config.settings.ssl_verify_default = True
    config.settings.ssl_ca_bundle_default = "/global/ca.crt"
    config.settings.yarn_resource_manager_url = "https://rm:8088"
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    conn = Connection(
        name="prod",
        master="yarn",
        ssl_verify=False,
        ssl_ca_bundle="/conn/ca.crt",
    )
    cfg = YarnClientConfig.from_connection(conn)
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", cfg)
    assert m.call_args.kwargs["verify"] == "/conn/ca.crt"


def test_ssl_verify_false_without_ca_bundle_uses_global_default():
    """If a connection only sets ssl_verify=False and no CA bundle is
    configured globally, httpx should skip certificate verification."""
    config.settings.ssl_verify_default = True
    config.settings.ssl_ca_bundle_default = None
    config.settings.yarn_resource_manager_url = "https://rm:8088"
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    conn = Connection(name="prod", master="yarn", ssl_verify=False)
    cfg = YarnClientConfig.from_connection(conn)
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", cfg)
    assert m.call_args.kwargs["verify"] is False


# --- authentication configuration ---

def test_status_no_auth_for_none():
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM, auth_type="none"))
    assert m.call_args.kwargs["auth"] is None

def test_status_basic_auth_header():
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    cfg = YarnClientConfig(
        yarn_rm_url=RM,
        auth_type="basic",
        auth_user="alice",
        auth_password="pw",
    )
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", cfg)
    auth = m.call_args.kwargs["auth"]
    assert isinstance(auth, httpx.BasicAuth)
    encoded = auth._auth_header.split(" ", 1)[1]
    assert base64.b64decode(encoded).decode() == "alice:pw"


def test_status_no_auth_for_simple():
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", YarnClientConfig(yarn_rm_url=RM, auth_type="simple"))
    assert m.call_args.kwargs["auth"] is None

def test_status_kerberos_auth():
    fake = _resp(200, json_data={"app": {"state": "RUNNING"}})
    cfg = YarnClientConfig(yarn_rm_url=RM, auth_type="kerberos")
    with patch("spark_executor.core.yarn_client.httpx.request", return_value=fake) as m:
        get_application_status("application_1", cfg)
    auth = m.call_args.kwargs["auth"]
    assert auth is not None
    assert "Kerberos" in type(auth).__name__

def test_basic_auth_missing_user_raises():
    cfg = YarnClientConfig(yarn_rm_url=RM, auth_type="basic")
    with pytest.raises(YarnConfigError, match="auth_type='basic' requires auth_user"):
        get_application_status("application_1", cfg)