# 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