mcp-server/tests/unit/test_job_store.py

# coding=utf-8
from datetime import datetime
from pathlib import Path

import pytest

from spark_executor.core.job_store import JobStore
from spark_executor.models import Job


def _job(jid: str, app_id: str | None = None) -> Job:
    return Job(
        job_id=jid,
        application_id=app_id or f"application_{jid}",
        script_path="/tmp/j.py",
        queue="default",
        submit_time=datetime(2026, 6, 24),
        connection="prod",
    )


@pytest.fixture
def store(tmp_path: Path) -> JobStore:
    """Per-test file-backed JobStore rooted in tmp_path. No cross-test leakage."""
    return JobStore(data_dir=str(tmp_path))


# --- Basic CRUD (was the entire file before the persistence fix) ---

def test_put_then_get_roundtrip(store):
    store.put(_job("a"))
    assert store.get("a") is not None
    assert store.get("a").application_id == "application_a"


def test_get_missing_returns_none(store):
    assert store.get("nope") is None


def test_list_returns_all_jobs(store):
    store.put(_job("a"))
    store.put(_job("b"))
    assert {j.job_id for j in store.list()} == {"a", "b"}


# --- Persistence: writes go to disk and survive a fresh instance ---

def test_data_persists_across_instances(tmp_path: Path):
    """The whole reason this used to be in-memory: gunicorn workers don't
    share memory. A second JobStore pointed at the same data_dir MUST see
    the records the first one wrote, otherwise we're back to the
    "Unknown job_id" bug from the multi-worker setup."""
    writer = JobStore(data_dir=str(tmp_path))
    writer.put(_job("a1b2c3d4e5f6", app_id="application_17400000001_0001"))

    reader = JobStore(data_dir=str(tmp_path))
    assert reader.get("a1b2c3d4e5f6") is not None
    assert reader.get("a1b2c3d4e5f6").application_id == "application_17400000001_0001"


def test_data_file_is_human_readable_json(tmp_path: Path):
    """If we're going to disk at all, the file should be inspectable
    without the application running — saves an ops engineer a forensics
    trip at 2am."""
    import json
    store = JobStore(data_dir=str(tmp_path))
    store.put(_job("abc", app_id="application_1"))
    raw = json.loads((tmp_path / "jobs.json").read_text(encoding="utf-8"))
    assert "abc" in raw
    assert raw["abc"]["application_id"] == "application_1"
    assert raw["abc"]["connection"] == "prod"


def test_corrupt_file_does_not_crash(store, tmp_path):
    """A partial write (e.g. killed mid-dump, full disk) shouldn't take
    the whole tool surface down — log it and treat as empty."""
    (tmp_path / "jobs.json").write_text("{not valid json", encoding="utf-8")
    # Should not raise; should return None / empty.
    assert store.get("anything") is None
    assert store.list() == []
    # And we should still be able to write through it.
    store.put(_job("x"))
    assert store.get("x") is not None


# --- get_either: accept job_id OR application_id ---

def test_get_either_finds_by_job_id(store):
    store.put(_job("a1b2c3d4e5f6", app_id="application_1"))
    job = store.get_either("a1b2c3d4e5f6")
    assert job is not None
    assert job.application_id == "application_1"


def test_get_either_finds_by_application_id(store):
    """The fix for 'agent passed the wrong id and got Unknown job_id':
    if the caller has a YARN application_id on hand, look it up by that."""
    store.put(_job("a1b2c3d4e5f6", app_id="application_17400000001_0001"))
    job = store.get_either("application_17400000001_0001")
    assert job is not None
    assert job.job_id == "a1b2c3d4e5f6"


def test_get_either_prefers_job_id_on_collision(store):
    """If a job_id and an application_id collide (unlikely but possible —
    e.g. someone seeds both), the direct job_id lookup wins."""
    store.put(_job("collide", app_id="not-collide"))
    store.put(_job("other", app_id="collide"))
    job = store.get_either("collide")
    assert job is not None
    assert job.job_id == "collide"


def test_get_either_returns_none_for_unknown(store):
    assert store.get_either("nothing-here") is None


def test_get_by_application_id_is_distinct(store):
    """get_by_application_id should NOT match by job_id (it's the explicit
    application_id-only lookup). get_either is the forgiving one."""
    store.put(_job("a1b2c3d4e5f6", app_id="application_1"))
    assert store.get_by_application_id("a1b2c3d4e5f6") is None
    assert store.get_by_application_id("application_1") is not None


# --- put is idempotent (re-put same job_id replaces, not duplicates) ---

def test_put_replaces_existing_job(store):
    """confirm_submit_job may retry; the second put of the same job_id
    must replace, not append, so list() doesn't grow on every retry."""
    store.put(_job("a", app_id="application_1"))
    store.put(_job("a", app_id="application_2"))  # same job_id, new app_id
    assert len(store.list()) == 1
    assert store.get("a").application_id == "application_2"