shanxia2.0/backend/tests/test_pl_plan_arrange_copy_flags.py

"""Regression: arrange copy flags must not N+1 query load/delivery per plan."""
from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock

from app.services.pl_plan_arrange_copy_flags import load_plan_arrange_copy_flags


def test_load_plan_arrange_copy_flags_batches_stop_lookups():
    plans = [
        SimpleNamespace(
            id=f"p{i}",
            transport_status="arranged",
            follow_status="pending_arrange",
            transport_id=f"t{i}",
            plan_type="direct",
            supplier_name="外部供方",
            customer_name="外部需方",
            destination_name="",
        )
        for i in range(5)
    ]

    transport_rows = [
        SimpleNamespace(
            id=f"t{i}",
            transport_status="arranged",
            driver_access_token="tok",
            site_access_token="",
        )
        for i in range(5)
    ]

    call_count = {"n": 0}

    def query_side_effect(model, *cols):
        call_count["n"] += 1
        q = MagicMock()
        q.filter.return_value = q
        q.distinct.return_value = q
        # PlTransport full objects
        name = getattr(model, "__name__", "") or str(model)
        if "PlTransport" in name and "Link" not in name and "Db" not in name and "Site" not in name:
            q.all.return_value = transport_rows
        elif "SiteLink" in name:
            q.all.return_value = []
        elif "DbLink" in name:
            q.all.return_value = []
        elif "Inbound" in name or "Outbound" in name:
            q.all.return_value = []
        elif "PlPlanLoad" in name or "Load" in name:
            q.all.return_value = []
        elif "Delivery" in name:
            q.all.return_value = []
        else:
            q.all.return_value = []
        return q

    db = MagicMock()
    db.query.side_effect = query_side_effect

    result = load_plan_arrange_copy_flags(db, "shanxia2.0", plans)
    assert len(result) == 5
    # 5 plans must not become ~15+ stop queries; expect a small fixed number of batch queries
    assert call_count["n"] <= 8
    for plan in plans:
        assert result[plan.id].can_copy_to_driver is True