mcp-server/tests/unit/test_logs_tool.py

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

import pytest

from spark_executor.core import connection_store
from spark_executor.core.job_store import JobStore
from spark_executor.models import Connection, Job
from spark_executor.tools import logs
from spark_executor.tools import connections


@pytest.fixture
def fresh_stores(tmp_path: Path):
    """Wire up connection + job stores rooted in tmp_path. Per-test isolation
    so file-backed JobStore doesn't leak between cases."""
    store = connection_store.ConnectionStore(data_dir=str(tmp_path))
    connection_store.store = store
    connections.store = store
    logs.conn_store = store
    logs.store = JobStore(data_dir=str(tmp_path))


def _seed(job_id="abc", app_id="application_1"):
    store = connection_store.ConnectionStore()
    connection_store.store = store
    connections.store = store
    logs.conn_store = store
    logs.store = JobStore()
    logs.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088"))
    logs.store.put(
        Job(
            job_id=job_id,
            application_id=app_id,
            script_path="/tmp/j.py",
            queue="default",
            submit_time=datetime(2026, 6, 24),
            connection="prod",
            yarn_rm_url="http://rm:8088",
        )
    )


def test_get_job_logs_tails_to_default_5000():
    _seed()
    big = "x" * 8000 + "\nEND"
    with patch("spark_executor.tools.logs.get_application_logs", return_value=big):
        out = logs.get_job_logs("abc")
    assert out.endswith("END")
    assert len(out) == 5000


def test_get_job_logs_respects_custom_tail_chars():
    _seed()
    with patch(
        "spark_executor.tools.logs.get_application_logs",
        return_value="0123456789",
    ):
        out = logs.get_job_logs("abc", tail_chars=3)
    assert out == "789"


def test_get_job_logs_raises_for_unknown_job():
    _seed()
    with pytest.raises(KeyError):
        logs.get_job_logs("missing")


def test_get_job_logs_raises_400_for_external_application_id(fresh_stores):
    """Input that looks like a YARN application_id but is not in the local
    JobStore must raise ValueError (-> 400) with a hint to use the
    external tool, NOT a generic KeyError (-> 404)."""
    with pytest.raises(ValueError, match="get_external_job_logs"):
        logs.get_job_logs("application_17400000001_0001")


def test_get_job_logs_raises_when_connection_missing(fresh_stores):
    logs.store.put(
        Job(
            job_id="abc",
            application_id="application_1",
            script_path="/tmp/j.py",
            queue="default",
            submit_time=datetime(2026, 6, 24),
            connection="missing",
        )
    )
    with pytest.raises(KeyError, match="Connection not found"):
        logs.get_job_logs("abc")


# --- application_id accepted (regression: "agent passed wrong id" bug) ---

def test_get_job_logs_accepts_application_id():
    """The agent gets both job_id and application_id back from
    confirm_submit_job and routinely passes the wrong one. The tool must
    accept EITHER and return the same logs."""
    _seed(job_id="a1b2c3d4e5f6", app_id="application_17400000001_0001")
    with patch(
        "spark_executor.tools.logs.get_application_logs",
        return_value="logs here",
    ) as m:
        out = logs.get_job_logs("application_17400000001_0001")
    assert out == "logs here"
    # And the underlying YARN call used the YARN ID, not the local job_id.
    assert m.call_args.args[0] == "application_17400000001_0001"


def test_get_job_logs_unknown_error_message_mentions_both_ids():
    """The "neither matched" error should explicitly call out BOTH
    accepted id forms so the agent doesn't guess."""
    _seed()
    with pytest.raises(KeyError) as ei:
        logs.get_job_logs("totally-fake")
    msg = str(ei.value)
    assert "job_id" in msg
    assert "application_id" in msg