# coding=utf-8
from datetime import datetime
from unittest.mock import patch

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


def test_get_job_status_returns_state():
    status.store = JobStore()
    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
    # yarn_rm_url is forwarded to the REST client
    assert m.call_args.args == ("application_1", "http://rm:8088")


def test_get_job_status_raises_for_unknown_job():
    status.store = JobStore()
    import pytest
    with pytest.raises(KeyError):
        status.get_job_status("missing")