SupplyAgent/tests/supply_infra/scheduler/test_grade_demand_pool.py

"""Current demand-grade scheduler wrapper tests."""
from __future__ import annotations

from unittest.mock import patch

from supply_infra.scheduler.jobs.grade_demand_pool import (
    execute_plan_tasks_until_complete,
    grade_demand_pool,
)


@patch("supply_infra.scheduler.jobs.grade_demand_pool._execute_plan_tasks")
def test_plan_execution_stops_when_no_groups_are_claimable(mock_execute) -> None:
    snapshot = {
        "execution_complete": True,
        "assigned_category_ids": [],
        "planned_groups": 0,
        "group_status": {},
    }
    mock_execute.return_value = {
        "attempted_groups": 0,
        "groups_run": 0,
        "workers": 0,
        "final_snapshot": snapshot,
        "execution_complete": True,
    }

    result = execute_plan_tasks_until_complete(
        "20260721",
        workers=5,
        max_demands_per_batch=20,
    )

    assert result["rounds"] == 1
    assert result["execution_complete"] is True
    mock_execute.assert_called_once()


@patch("supply_infra.scheduler.jobs.grade_demand_pool._grade_demand_pool_impl")
@patch("supply_infra.scheduler.jobs.grade_demand_pool._resolve_biz_dt", return_value="20260721")
def test_grade_job_never_raises_to_pipeline(_mock_resolve, mock_impl) -> None:
    mock_impl.side_effect = RuntimeError("boom")

    result = grade_demand_pool("20260721")

    assert result["success"] is False
    assert result["biz_dt"] == "20260721"
    assert result["error"] == "boom"