stack/tests/aco/test_lake_jobs.py

"""Tests for aco.lake.jobs — programmatic job management."""

from __future__ import annotations

from unittest.mock import MagicMock

from aco.lake.jobs import JobManager, RunStatus


class TestRunStatus:
    def test_terminal_states(self):
        assert RunStatus(run_id=1, state="TERMINATED").is_terminal
        assert RunStatus(run_id=1, state="SKIPPED").is_terminal
        assert not RunStatus(run_id=1, state="RUNNING").is_terminal

    def test_success(self):
        assert RunStatus(run_id=1, state="TERMINATED", result_state="SUCCESS").succeeded
        assert not RunStatus(
            run_id=1, state="TERMINATED", result_state="FAILED"
        ).succeeded


class TestJobManager:
    def _make_manager(self):
        client = MagicMock()
        return JobManager(client, catalog="aco_dev"), client

    def test_create_pipeline_job(self):
        mgr, client = self._make_manager()
        client._ws.jobs.create.return_value = MagicMock(job_id=42)

        job_id = mgr.create_pipeline_job("readmissions")
        assert job_id == 42
        client._ws.jobs.create.assert_called_once()

        # Verify task structure
        call_kwargs = client._ws.jobs.create.call_args.kwargs
        assert call_kwargs["name"] == "stack_readmissions"
        assert len(call_kwargs["tasks"]) >= 1
        task = call_kwargs["tasks"][0]
        assert task.task_key == "readmissions"
        assert task.python_wheel_task.package_name == "stack"

    def test_create_with_schedule(self):
        mgr, client = self._make_manager()
        client._ws.jobs.create.return_value = MagicMock(job_id=99)

        mgr.create_pipeline_job("core", schedule="0 0 6 * * ?")
        call_kwargs = client._ws.jobs.create.call_args.kwargs
        assert call_kwargs["schedule"] is not None

    def test_run_now(self):
        mgr, client = self._make_manager()
        client._ws.jobs.run_now.return_value = MagicMock(run_id=123)

        run_id = mgr.run_now(42)
        assert run_id == 123
        client._ws.jobs.run_now.assert_called_once_with(job_id=42)

    def test_get_run_status(self):
        mgr, client = self._make_manager()
        state = MagicMock()
        state.life_cycle_state = "RUNNING"
        state.result_state = None
        state.state_message = "In progress"
        run = MagicMock()
        run.state = state
        client._ws.jobs.get_run.return_value = run

        status = mgr.get_run_status(123)
        assert status.state == "RUNNING"
        assert not status.is_terminal

    def test_list_jobs(self):
        mgr, client = self._make_manager()
        job = MagicMock()
        job.job_id = 1
        job.settings.name = "stack_core"
        client._ws.jobs.list.return_value = [job]

        jobs = mgr.list_jobs()
        assert len(jobs) == 1
        assert jobs[0]["name"] == "stack_core"

    def test_delete_job(self):
        mgr, client = self._make_manager()
        mgr.delete_job(42)
        client._ws.jobs.delete.assert_called_once_with(job_id=42)