# coding=utf-8 from datetime import datetime from spark_executor.models import Connection, Job, JobStatus, PendingSubmission, SubmitResult def test_job_roundtrip(): job = Job( job_id="abc123", application_id="application_17400000001", script_path="/tmp/jobs/job_001.py", queue="default", submit_time=datetime(2026, 6, 24, 10, 0, 0), connection="prod-yarn", yarn_rm_url="http://rm:8088", ) dumped = job.model_dump() assert dumped["job_id"] == "abc123" assert dumped["application_id"] == "application_17400000001" assert dumped["connection"] == "prod-yarn" assert dumped["yarn_rm_url"] == "http://rm:8088" def test_job_yarn_rm_url_optional(): job = Job( job_id="j1", application_id="application_1", script_path="/tmp/x.py", queue="default", submit_time=datetime(2026, 6, 24), connection="dev", ) assert job.yarn_rm_url is None def test_job_status_default_raw(): s = JobStatus(application_id="application_1", state="RUNNING") assert s.raw == "" def test_submit_result_tracking_url_optional(): r = SubmitResult(job_id="j1", application_id="application_1", tracking_url=None) assert r.tracking_url is None def test_connection_defaults(): c = Connection(name="prod", master="yarn") assert c.deploy_mode == "cluster" assert c.yarn_rm_url is None assert c.spark_conf == {} def test_connection_with_spark_conf(): c = Connection( name="dev", master="spark://master:7077", deploy_mode="client", spark_conf={"spark.sql.shuffle.partitions": "200"}, ) assert c.spark_conf["spark.sql.shuffle.partitions"] == "200" def test_pending_submission_defaults_to_pending_status(): p = PendingSubmission( pending_id="p_1", connection="prod", master="yarn", deploy_mode="cluster", script_path="/tmp/j.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, spark_conf={}, created_at=datetime(2026, 6, 24), ) assert p.status == "PENDING" assert p.error is None assert p.job_id is None assert p.application_id is None assert p.yarn_rm_url is None # default for connections without one set def test_pending_submission_carries_yarn_rm_url_snapshot(): """snapshotting yarn_rm_url lets prepare/confirm survive connection edits.""" p = PendingSubmission( pending_id="p_x", connection="prod", master="yarn", deploy_mode="cluster", yarn_rm_url="http://rm-prod:8088", script_path="/tmp/j.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, spark_conf={}, created_at=datetime(2026, 6, 24), ) assert p.yarn_rm_url == "http://rm-prod:8088" def test_pending_submission_can_record_outcome(): p = PendingSubmission( pending_id="p_2", connection="prod", master="yarn", deploy_mode="cluster", script_path="/tmp/j.py", queue="default", executor_memory="4G", executor_cores=2, num_executors=2, spark_conf={}, created_at=datetime(2026, 6, 24), status="SUBMITTED", job_id="j_abc", application_id="application_17400000001", ) assert p.status == "SUBMITTED" assert p.job_id == "j_abc" assert p.application_id == "application_17400000001"