# coding=utf-8 from datetime import datetime from pathlib import Path from unittest.mock import patch import pytest from spark_executor.core.job_store import JobStore from spark_executor.core import connection_store from spark_executor.models import Connection, Job from spark_executor.tools import connections, status from spark_executor.core.yarn_client import YarnClientConfig @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 status.conn_store = store status.store = JobStore(data_dir=str(tmp_path)) def _fresh_stores(): """Reset job store and connection store singletons for a single test.""" store = connection_store.ConnectionStore() connection_store.store = store connections.store = store status.conn_store = store status.store = JobStore() def test_get_job_status_returns_state(): _fresh_stores() status.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")) status.store.put( Job( job_id="abc", application_id="application_1", script_path="/tmp/j.py", queue="default", submit_time=datetime(2026, 6, 24), connection="prod", yarn_rm_url="http://rm:8088", ) ) with patch( "spark_executor.tools.status.get_application_status", return_value=("RUNNING", "State : RUNNING\n"), ) as m: out = status.get_job_status("abc") assert out.application_id == "application_1" assert out.state == "RUNNING" assert "RUNNING" in out.raw # resolved YarnClientConfig is forwarded to the REST client args = m.call_args.args assert args[0] == "application_1" assert args[1].yarn_rm_url == "http://rm:8088" def test_get_job_status_raises_for_unknown_job(): _fresh_stores() with pytest.raises(KeyError): status.get_job_status("missing") def test_get_job_status_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_status"): status.get_job_status("application_17400000001_0001") def test_get_job_status_raises_when_connection_missing(): _fresh_stores() status.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"): status.get_job_status("abc") def test_get_job_status_passes_auth_config(): _fresh_stores() status.conn_store.save( Connection( name="secure", master="yarn", yarn_rm_url="http://rm:8088", auth_type="basic", auth_user="hdfs", auth_password="secret", ) ) status.store.put( Job( job_id="abc", application_id="application_1", script_path="/tmp/j.py", queue="default", submit_time=datetime(2026, 6, 24), connection="secure", yarn_rm_url="http://rm:8088", ) ) with patch( "spark_executor.tools.status.get_application_status", return_value=("RUNNING", "State : RUNNING\n"), ) as m: status.get_job_status("abc") args = m.call_args.args assert args[0] == "application_1" assert isinstance(args[1], YarnClientConfig) assert args[1].auth_type == "basic" assert args[1].auth_user == "hdfs" # --- application_id accepted (regression: "agent passed wrong id" bug) --- def test_get_job_status_accepts_application_id(fresh_stores): status.conn_store.save(Connection(name="prod", master="yarn", yarn_rm_url="http://rm:8088")) status.store.put( Job( job_id="a1b2c3d4e5f6", application_id="application_17400000001_0001", script_path="/tmp/j.py", queue="default", submit_time=datetime(2026, 6, 24), connection="prod", yarn_rm_url="http://rm:8088", ) ) with patch( "spark_executor.tools.status.get_application_status", return_value=("RUNNING", "{}"), ): out = status.get_job_status("application_17400000001_0001") assert out.state == "RUNNING" assert out.application_id == "application_17400000001_0001" def test_get_job_status_unknown_error_message_mentions_both_ids(fresh_stores): with pytest.raises(KeyError) as ei: status.get_job_status("totally-fake") msg = str(ei.value) assert "job_id" in msg assert "application_id" in msg