Aislo/resources/tester/test_618_precompute_split.py

"""PLAN 618-1 — 전국 미리 계산 두 PC 나눔(`tools/precompute_terrain.py`) 지킴 시험.

나눔 겹침 0 · 합치면 전체 · 이웃 일감 한 PC 만 · PC 별 기록 이어 받기 · CPU 상한(가짜 값) ·
0차곡 남의 몫 이웃 없으면 미룸 · 지정 코어 인자 풀기 ·
셈 함수는 가짜로 바꿔 서버 · 전국 파일 없이 돎.
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import pytest

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

from B04_PreProcess.B04_PreProcess_Engine_MapSheet import neighbors_3x3  # noqa: E402
from tools import precompute_terrain as tool  # noqa: E402

SHEETS = [f"36901{i:03d}" for i in range(1, 101)] + [f"35810{i:03d}" for i in range(1, 101)]


def test_parse_part_and_cpus() -> None:
    assert tool.parse_part("") is None
    assert tool.parse_part("0-69/100") == (0, 69, 100)
    assert tool.parse_part("70-99/100") == (70, 99, 100)
    for bad in ("50-10/100", "0-100/100"):
        with pytest.raises(ValueError):
            tool.parse_part(bad)
    assert tool.parse_cpus("12-19") == list(range(12, 20))
    assert tool.parse_cpus("1, 3-4") == [1, 3, 4]
    assert tool.parse_cpus("") == []


def test_split_no_overlap_and_union_is_all() -> None:
    for spans in (("0-69/100", "70-99/100"), ("0-49/100", "50-99/100")):
        a, b = (tool.parse_part(s) for s in spans)
        mine = {no for no in SHEETS if tool.in_part(no, a)}
        other = {no for no in SHEETS if tool.in_part(no, b)}
        assert not mine & other and mine | other == set(SHEETS)
        ja, jb = set(tool.dem_jobs(SHEETS, a)), set(tool.dem_jobs(SHEETS, b))
        assert not ja & jb and ja | jb == set(tool.dem_jobs(SHEETS, None))
    big = {no for no in SHEETS if tool.in_part(no, (0, 69, 100))}
    assert 0.6 < len(big) / len(SHEETS) < 0.8
    assert tool.in_part("36901001@5186", (0, 99, 100))


def test_dem_jobs_cover_neighbors_once() -> None:
    jobs = tool.dem_jobs(["36901055"], None)
    assert jobs[0] == "36901055" and set(jobs) == set(neighbors_3x3("36901055"))
    # 띠 경계(128° · 5186 | 5187) 넘는 이웃 = 가운데 좌표계로 옮겨 담는 일감
    west = [no for no in tool.dem_jobs(["36809001"], None) if "@" in no]
    assert west and all(j.endswith(f"@{tool.epsg_of('36809001')}") for j in west)


@pytest.fixture()
def fake_store(tmp_path, monkeypatch):
    monkeypatch.setattr(tool.store, "TERRAIN_ROOT", tmp_path)
    monkeypatch.setattr(tool.store, "_gpkg_stamp", lambda: "stamp-1")
    monkeypatch.setattr(tool, "PC_NAME", "LAPTOP")
    monkeypatch.setattr(tool, "CPU_SAMPLE_S", 0.0)
    log: list = []

    def fake_run(step: str, job: str, part=None):
        log.append(("run", step, job))
        return step, job, {"ok": True, "s": 0.01, "fresh": True, "empty": False}

    monkeypatch.setattr(tool, "run_step", fake_run)
    return tmp_path, log


def test_progress_per_pc_and_legacy_resume(fake_store, monkeypatch) -> None:
    root, log = fake_store
    legacy = {"gpkg": "stamp-1", "cell_m": 2.0, "steps": {"dem": {"36901001": {"ok": True}}}}
    (root / tool.PROGRESS_NAME).write_text(json.dumps(legacy), encoding="utf-8")
    tool.run(["36901001", "36901002"], ["e07"], 1, tool.load_progress(), say=lambda _: None)
    assert (root / "_precompute_progress_LAPTOP.json").is_file()
    assert json.loads((root / tool.PROGRESS_NAME).read_text(encoding="utf-8")) == legacy  # 안 씀
    log.clear()
    tool.run(["36901001", "36901002"], ["dem"], 1, tool.load_progress(), say=lambda _: None)
    assert ("run", "dem", "36901001") not in log  # 옛 기록 성공 = 이어 받음

    monkeypatch.setattr(tool, "PC_NAME", "DESKTOP")
    desk = tool.load_progress()
    assert desk["steps"]["e07"] == {}  # 다른 PC 기록은 run 에선 안 섞음
    desk["steps"]["e07"]["36901003"] = {"ok": True, "s": 1.0, "fresh": True}
    tool.save_progress(desk)
    merged = tool.load_progress(every_pc=True)["steps"]["e07"]
    assert {"36901001", "36901002", "36901003"} <= set(merged)


def test_cpu_cap_holds_new_work(fake_store) -> None:
    _, log = fake_store
    values = iter([95, 90, 60, 80, 40])

    def cpu() -> float:
        value = next(values)
        log.append(("cpu", value))
        return value

    tool.run(
        ["36901001", "36901002"],
        ["e07"],
        1,
        tool.load_progress(),
        cpu_cap=70,
        cpu=cpu,
        say=lambda _: None,
    )
    assert [x[1] if x[0] == "cpu" else x[2] for x in log] == [
        95,
        90,
        60,
        "36901001",
        80,
        40,
        "36901002",
    ]


def test_zero_defers_when_foreign_neighbor_missing(tmp_path, monkeypatch) -> None:
    monkeypatch.setattr(tool.store, "TERRAIN_ROOT", tmp_path)
    center = "36901055"
    part = (0, 49, 100)
    foreign = [no for no in neighbors_3x3(center) if not tool.in_part(no, part)]
    assert foreign  # 고른 가운데는 남의 몫 이웃이 있음
    stored: set[str] = set()
    monkeypatch.setattr(
        tool.store, "_stored_ok", lambda no, e, c: {"k": 1} if no in stored else None
    )
    ensured: list[str] = []
    monkeypatch.setattr(tool.store, "ensure_sheet_terrain", lambda no, e, c: ensured.append(no))
    _, _, info = tool.run_step("zero", center, part)
    assert info["ok"] is False and info["wait"] and not ensured  # 셈 안 함

    stored.update(foreign)  # 다른 PC 가 셈 · 동기 옴
    from D03_PlanAnalysis import D03_PlanAnalysis_ZeroOrder as zero

    monkeypatch.setattr(zero, "sheet_accumulation", lambda no, e, c: tmp_path / "acc.npz")
    _, _, info = tool.run_step("zero", center, part)
    assert info["ok"] is True and len(ensured) == 9


def test_deferred_rounds_retry_until_done(fake_store, monkeypatch) -> None:
    root, log = fake_store
    tries: dict[str, int] = {}

    def fake_run(step: str, job: str, part=None):
        tries[job] = tries.get(job, 0) + 1
        if job == "36901002" and tries[job] < 3:
            return step, job, {"ok": False, "wait": True, "error": "이웃 지형 없음"}
        return step, job, {"ok": True, "s": 0.01, "fresh": True, "empty": False}

    monkeypatch.setattr(tool, "run_step", fake_run)
    monkeypatch.setattr(tool.time, "sleep", lambda s: None)
    progress = tool.load_progress()
    tool.run(["36901001", "36901002"], ["zero"], 1, progress, say=lambda _: None, defer_wait=1)
    assert tries == {"36901001": 1, "36901002": 3}
    out = tool.summary(["36901001", "36901002"], progress)["zero"]
    assert out["ok"] == 2 and out["failed"] == {} and out["waiting"] == 0


def test_clean_dead_only_own_part(tmp_path) -> None:
    part = (0, 49, 100)
    mine = next(no for no in SHEETS if tool.in_part(no, part))
    other = next(no for no in SHEETS if not tool.in_part(no, part))
    for no in (mine, other):
        (tmp_path / no).mkdir()
        (tmp_path / no / ".dem_5187_2m.lock").write_text("999999999 0", encoding="utf-8")
    assert tool.clean_dead(tmp_path, part) == 1
    assert (tmp_path / other / ".dem_5187_2m.lock").exists()


def test_worker_init_pins_first_e_workers(monkeypatch, tmp_path) -> None:
    import multiprocessing

    import psutil

    pinned: list[list[int]] = []

    class FakeProc:
        def nice(self, _):
            pass

        def cpu_affinity(self, cpus):
            pinned.append(cpus)

    monkeypatch.setattr(psutil, "Process", FakeProc)
    monkeypatch.setattr(tool.store, "TERRAIN_ROOT", tool.store.TERRAIN_ROOT)  # 끝나면 되돌림
    counter = multiprocessing.Value("i", 0)
    for _ in range(5):
        tool.init_worker(str(tmp_path), counter, [12, 13], 3)
    assert pinned == [[12, 13]] * 3 and tool.store.TERRAIN_ROOT == tmp_path