diff --git a/perf/perfkit/gen/grow.py b/perf/perfkit/gen/grow.py index 005cd0b..0675a3f 100644 --- a/perf/perfkit/gen/grow.py +++ b/perf/perfkit/gen/grow.py @@ -60,11 +60,27 @@ async def grow_vsp_sheet(api: Api, form_id: int, sheet: str, need: int, org_inde return need +def form4_ids(rows: list[dict]) -> tuple[list[int], list[int]]: + """(vsp_ids, project_ids) из строк листа FORM_4; отсутствующие ключи (data/header/vsp_id/project_id) пропускаются.""" + vsp_ids: set[int] = set() + project_ids: set[int] = set() + for r in rows: + header = (r.get("data") or {}).get("header") or {} + if r.get("row_type") == "INPUT": + vid = header.get("vsp_id") + if vid is not None: + vsp_ids.add(vid) + elif r.get("row_type") == "ITEM": + pid = header.get("project_id") + if pid is not None: + project_ids.add(pid) + return sorted(vsp_ids), sorted(project_ids) + + async def grow_form4_sheet(api: Api, form_id: int, sheet: str, need: int) -> int: path = f"/form/{form_id}/sheet/{sheet}" rows = (await api.get(path))["result"] - vsp_ids = sorted({r["data"]["header"]["vsp_id"] for r in rows if r.get("row_type") == "INPUT" and r["data"]["header"].get("vsp_id")}) - project_ids = sorted({r["data"]["header"]["project_id"] for r in rows if r.get("row_type") == "ITEM"}) + vsp_ids, project_ids = form4_ids(rows) if not project_ids: raise RuntimeError(f"{sheet}: нет проектов") leaves = await leaf_items(api, "FORM_4", sheet) diff --git a/perf/tests/test_gen_grow.py b/perf/tests/test_gen_grow.py index 1b4aa4e..70e7d6e 100644 --- a/perf/tests/test_gen_grow.py +++ b/perf/tests/test_gen_grow.py @@ -1,6 +1,6 @@ import pytest -from perfkit.gen.grow import rows_to_add, TARGETS +from perfkit.gen.grow import form4_ids, rows_to_add, TARGETS @pytest.mark.parametrize("cur,target,exp", [(31, 500, 469), (110, 300, 190), (500, 500, 0), (700, 500, 0), (0, 5, 5)]) @@ -12,3 +12,15 @@ def test_targets_do_not_touch_small_pair_sheets(): # AHR_RENT и FORM_4/AHR остаются малыми (пара для @big) assert ("FORM_2", "AHR_RENT") not in [(f, s) for f, s, _ in TARGETS] assert ("FORM_4", "AHR") not in [(f, s) for f, s, _ in TARGETS] + + +def test_form4_ids_tolerates_absent_keys(): + rows = [ + {"row_type": "INPUT", "data": {"header": {"vsp_id": 3}}}, + {"row_type": "INPUT", "data": {"header": {}}}, + {"row_type": "INPUT", "data": {}}, + {"row_type": "ITEM", "data": {"header": {"project_id": 11}}}, + {"row_type": "ITEM", "data": {"header": {}}}, + {"row_type": "ITEM", "data": {}}, + ] + assert form4_ids(rows) == ([3], [11])