SupplyAgent/tests/supply_infra/pipeline/test_resume_cancelled_run.py

from __future__ import annotations

from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from supply_infra.db.models.pipeline_run import PipelineRun
from supply_infra.db.models.pipeline_step_run import PipelineStepRun
from supply_infra.pipeline import run_service


def _session_factory() -> sessionmaker[Session]:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    PipelineRun.__table__.create(engine)
    PipelineStepRun.__table__.create(engine)
    return sessionmaker(bind=engine, expire_on_commit=False)


def _run(status: str = "cancelled") -> PipelineRun:
    return PipelineRun(
        run_id="run-cancelled",
        dedupe_key="supply_pipeline:20260727:full",
        pipeline_key="supply_pipeline",
        biz_dt="20260727",
        trigger_type="api",
        run_mode="full",
        dry_run=True,
        status=status,
        current_step=None,
        deadline_at=None,
        config_snapshot_json={},
        date_snapshot_json={},
        summary_json={"old": "summary"},
        error_code="cancellation_requested",
        error_message="cancelled",
    )


def _step(
    *,
    step_run_id: str,
    key: str,
    order: int,
    status: str,
) -> PipelineStepRun:
    return PipelineStepRun(
        step_run_id=step_run_id,
        run_id="run-cancelled",
        step_key=key,
        step_order=order,
        attempt=1,
        status=status,
        critical=True,
        dependency_snapshot_json=[],
        input_snapshot_json={},
        timeout_seconds=60,
        max_attempts=1,
        retryable=False,
        error_code="cancelled" if status == "cancelled" else None,
        error_message="Cancellation requested" if status == "cancelled" else None,
    )


def _patch_session(monkeypatch, factory: sessionmaker[Session]) -> None:
    @contextmanager
    def get_test_session():
        session = factory()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    monkeypatch.setattr(run_service, "get_session", get_test_session)


def test_cancelled_run_resumes_from_first_cancelled_step(monkeypatch) -> None:
    factory = _session_factory()
    _patch_session(monkeypatch, factory)
    with factory() as session:
        session.add_all(
            [
                _run(),
                _step(
                    step_run_id="step-1",
                    key="global_tree_sync",
                    order=1,
                    status="succeeded",
                ),
                _step(
                    step_run_id="step-2",
                    key="source_video_sync",
                    order=2,
                    status="cancelled",
                ),
                _step(
                    step_run_id="step-3",
                    key="demand_grade",
                    order=3,
                    status="cancelled",
                ),
            ]
        )
        session.commit()

    assert run_service.resume_pipeline_run("run-cancelled") is True

    with factory() as session:
        run = session.get(PipelineRun, "run-cancelled")
        assert run is not None
        assert run.status == "queued"
        assert run.current_step == "source_video_sync"
        assert run.finished_at is None
        assert run.summary_json is None
        assert run.error_code is None
        assert run.error_message is None

        steps = (
            session.query(PipelineStepRun)
            .filter(PipelineStepRun.run_id == run.run_id)
            .order_by(PipelineStepRun.step_order, PipelineStepRun.attempt)
            .all()
        )
        assert [(step.step_key, step.attempt, step.status) for step in steps] == [
            ("global_tree_sync", 1, "succeeded"),
            ("source_video_sync", 1, "cancelled"),
            ("source_video_sync", 2, "ready"),
            ("demand_grade", 1, "pending"),
        ]


def test_cancelling_run_cannot_resume_before_cancellation_finishes(monkeypatch) -> None:
    factory = _session_factory()
    _patch_session(monkeypatch, factory)
    with factory() as session:
        session.add_all(
            [
                _run(status="cancelling"),
                _step(
                    step_run_id="step-1",
                    key="source_video_sync",
                    order=1,
                    status="running",
                ),
            ]
        )
        session.commit()

    assert run_service.resume_pipeline_run("run-cancelled") is False