SupplyAgent/tests/supply_infra/pipeline/test_step_claim.py

from __future__ import annotations

from datetime import timedelta

from sqlalchemy import create_engine
from sqlalchemy.orm import Session

from supply_infra.db.models.pipeline_lock import PipelineLock
from supply_infra.db.models.pipeline_run import PipelineRun
from supply_infra.db.models.pipeline_step_run import PipelineStepRun
from supply_infra.db.repositories.pipeline_step_run_repo import (
    PipelineStepRunRepository,
)
from supply_infra.pipeline.dates import china_now


def _session() -> Session:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    PipelineRun.__table__.create(engine)
    PipelineStepRun.__table__.create(engine)
    PipelineLock.__table__.create(engine)
    return Session(engine, expire_on_commit=False)


def _run(run_id: str) -> PipelineRun:
    now = china_now()
    return PipelineRun(
        run_id=run_id,
        dedupe_key=f"supply_pipeline:20260727:{run_id}",
        pipeline_key="supply_pipeline",
        biz_dt="20260727",
        trigger_type="test",
        run_mode="full",
        dry_run=True,
        status="queued",
        deadline_at=now + timedelta(days=1),
        config_snapshot_json={},
        date_snapshot_json={},
    )


def _step(
    *,
    step_run_id: str,
    run_id: str,
    status: str,
    order: int,
) -> PipelineStepRun:
    return PipelineStepRun(
        step_run_id=step_run_id,
        run_id=run_id,
        step_key=f"step-{step_run_id}",
        step_order=order,
        attempt=1,
        status=status,
        critical=True,
        dependency_snapshot_json=[],
        input_snapshot_json={},
        timeout_seconds=10,
        max_attempts=1,
        retryable=False,
    )


def test_claim_respects_global_active_step_cap() -> None:
    with _session() as session:
        session.add_all(
            [
                _run("run-1"),
                _run("run-2"),
                _step(
                    step_run_id="running",
                    run_id="run-1",
                    status="running",
                    order=1,
                ),
                _step(
                    step_run_id="ready",
                    run_id="run-2",
                    status="ready",
                    order=1,
                ),
            ]
        )
        session.flush()

        claimed = PipelineStepRunRepository(session).claim_next_ready(
            owner="worker-2",
            lease_seconds=120,
            max_active_steps=1,
            now=china_now(),
        )

        assert claimed is None
        ready = session.get(PipelineStepRun, "ready")
        assert ready is not None
        assert ready.status == "ready"


def test_claim_does_not_overlap_a_different_pipeline_run() -> None:
    with _session() as session:
        session.add_all(
            [
                _run("run-1"),
                _run("run-2"),
                _step(
                    step_run_id="running",
                    run_id="run-1",
                    status="running",
                    order=1,
                ),
                _step(
                    step_run_id="ready",
                    run_id="run-2",
                    status="ready",
                    order=1,
                ),
            ]
        )
        session.flush()

        claimed = PipelineStepRunRepository(session).claim_next_ready(
            owner="worker-2",
            lease_seconds=120,
            max_active_steps=2,
            now=china_now(),
        )

        assert claimed is None