mcp-server/tests/unit/test_pending_store.py

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

import pytest

from spark_executor.core import pending_store
from spark_executor.models import PendingSubmission


def _pending(pid: str = "p_1", created_at: datetime | None = None) -> PendingSubmission:
    return PendingSubmission(
        pending_id=pid,
        app_name="test",
        connection="prod",
        master="yarn",
        deploy_mode="cluster",
        script_path="/tmp/j.py",
        queue="default",
        executor_memory="2G",
        executor_cores=2,
        num_executors=2,
        spark_conf={},
        extra_args={},
        created_at=created_at or datetime(2026, 6, 24),
    )


@pytest.fixture(autouse=True)
def _fresh(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(pending_store, "store", pending_store.PendingStore())


def test_save_then_get_roundtrip():
    pending_store.store.save(_pending())
    out = pending_store.store.get("p_1")
    assert out is not None
    assert out.connection == "prod"
    assert out.master == "yarn"


def test_save_persists_to_date_file(tmp_path: Path):
    pending_store.store.save(_pending())
    path = tmp_path / "pending_jobs" / "2026-06-24.json"
    assert path.is_file()


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


def test_list_all_empty_when_file_absent():
    assert pending_store.store.list_all() == []


def test_list_all_returns_saved():
    pending_store.store.save(_pending("a"))
    pending_store.store.save(_pending("b"))
    assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"}


def test_list_all_spans_multiple_date_files():
    pending_store.store.save(_pending("a", created_at=datetime(2026, 6, 24)))
    pending_store.store.save(_pending("b", created_at=datetime(2026, 6, 25)))
    assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"}


def test_get_searches_date_files():
    pending_store.store.save(_pending("a", created_at=datetime(2026, 6, 24)))
    assert pending_store.store.get("a") is not None


def test_delete_removes_from_date_file_and_deletes_empty_file():
    pending_store.store.save(_pending("a", created_at=datetime(2026, 6, 24)))
    date_file = pending_store.store.dir_path / "2026-06-24.json"
    assert date_file.is_file()
    assert pending_store.store.delete("a") is True
    assert pending_store.store.get("a") is None
    assert not date_file.exists()


def _legacy_record_text() -> str:
    return (
        '{"p_legacy": {"pending_id": "p_legacy", "app_name": "legacy", '
        '"connection": "prod", "master": "yarn", "deploy_mode": "cluster", '
        '"script_path": "/tmp/j.py", "queue": "default", "executor_memory": "2G", '
        '"executor_cores": 2, "num_executors": 2, "spark_conf": {}, '
        '"created_at": "2026-06-23T00:00:00"}}'
    )


def test_legacy_file_is_readable(tmp_path: Path):
    legacy = tmp_path / "pending_jobs.json"
    legacy.write_text(_legacy_record_text(), encoding="utf-8")
    assert pending_store.store.get("p_legacy") is not None
    assert any(p.pending_id == "p_legacy" for p in pending_store.store.list_all())


def test_save_migrates_legacy_file(tmp_path: Path):
    legacy = tmp_path / "pending_jobs.json"
    legacy.write_text(_legacy_record_text(), encoding="utf-8")
    pending_store.store.save(_pending("p_new", created_at=datetime(2026, 6, 24)))
    assert not legacy.exists()
    assert pending_store.store.get("p_legacy") is not None
    migrated = pending_store.store.dir_path / "2026-06-23.json"
    assert migrated.is_file()


def test_save_upserts_existing():
    pending_store.store.save(_pending("a"))
    updated = _pending("a")
    updated.status = "SUBMITTED"
    updated.job_id = "j_1"
    pending_store.store.save(updated)
    assert pending_store.store.get("a").status == "SUBMITTED"
    assert pending_store.store.get("a").job_id == "j_1"


def test_delete_returns_true_when_present():
    pending_store.store.save(_pending())
    assert pending_store.store.delete("p_1") is True
    assert pending_store.store.get("p_1") is None


def test_delete_returns_false_when_absent():
    assert pending_store.store.delete("nope") is False


def test_loads_records_with_missing_or_null_app_name(tmp_path: Path):
    """Records persisted before app_name became required/optional must still load."""
    shard = tmp_path / "pending_jobs" / "2026-06-24.json"
    shard.parent.mkdir(parents=True, exist_ok=True)
    base = {
        "pending_id": "p_no_app",
        "connection": "prod",
        "master": "yarn",
        "deploy_mode": "cluster",
        "script_path": "/tmp/j.py",
        "queue": "default",
        "executor_memory": "2G",
        "executor_cores": 2,
        "num_executors": 2,
        "spark_conf": {},
        "created_at": "2026-06-24T00:00:00",
    }
    missing_app = dict(base, pending_id="p_no_app")
    null_app = dict(base, pending_id="p_null_app", app_name=None)
    shard.write_text(
        json.dumps({"p_no_app": missing_app, "p_null_app": null_app}),
        encoding="utf-8",
    )

    records = pending_store.store.list_all()
    assert len(records) == 2
    by_id = {r.pending_id: r for r in records}
    assert by_id["p_no_app"].app_name is None
    assert by_id["p_null_app"].app_name is None

    assert pending_store.store.get("p_no_app").app_name is None
    assert pending_store.store.get("p_null_app").app_name is None