"""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