# 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"