pylib_ai_arm_statements/arm_statements/tests/test_segments.py

from decimal import Decimal as D

from arm_statements.header import parse_header
from arm_statements.loader import Rules
from arm_statements.model import RowRole
from arm_statements.segments import find_layout
from arm_statements.tests.builders import excel, hdr, ocr, place

DATES3 = ["На 30 сентября 2025 г.", "На 31 декабря 2024 г.", "На 31 декабря 2023 г."]


def test_npz_sheet_with_own_header_maps_by_date(rules: Rules) -> None:
    first = [
        place({1: "Наименование показателя", 12: "Код", 16: DATES3[0], 21: DATES3[1], 24: DATES3[2]}),
        place({1: "Нематериальные активы", 12: "1110", 16: D("10"), 21: D("9"), 24: D("8")}),
        place({1: "Итого по разделу I", 12: "1100", 16: D("10"), 21: D("9"), 24: D("8")}),
    ]
    sheet4 = [
        place({0: "Наименование показателя", 6: "Код", 8: DATES3[0], 9: DATES3[1], 10: DATES3[2]}),
        place({0: "Заемные средства", 6: "1510", 8: D("3"), 9: D("2"), 10: D("1")}),
    ]
    cand = excel([("Бухгалтерский баланс", first), ("Бухгалтерский баланс. Лист 4", sheet4)])
    layout = find_layout(cand, parse_header(hdr(report_date="30.09.2025"), rules), rules)
    assert [c.column_index for c in layout.columns] == [16, 21, 24]
    assert layout.segments[1].value_cols == ((8, 16), (9, 21), (10, 24))
    assert not [w for w in layout.warnings if w.code.startswith("SEGMENT_")]


def test_repeat_header_mid_table(rules: Rules) -> None:
    rows = [
        ["Наименование показателя", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Запасы", "1210", "5", "4"],
        ["Наименование показателя", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Итого", "1200", "5", "4"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [s.first_row for s in layout.segments] == [0, 2]
    assert layout.roles[2] is RowRole.HEADER


def test_repeat_header_without_dates_maps_by_order(rules: Rules) -> None:
    rows = [
        ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Запасы", "1210", "5", "4"],
        ["Наименование", "Код", "отчетн", "прошл"],
        ["Итого", "1200", "5", "4"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert layout.segments[1].value_cols == ((2, 2), (3, 3))
    assert layout.warnings == ()


def test_wider_segment_is_mismatch(rules: Rules) -> None:
    rows = [
        ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024", "На 31.12.2023"],
        ["Запасы", "1210", "5", "4", "3"],
        ["Наименование", "Код", "графа 3", "графа 4", "графа 5", "графа 6"],
        ["Итого", "1200", "5", "4", "3", "2"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    (warning,) = [w for w in layout.warnings if w.code == "SEGMENT_COLUMNS_MISMATCH"]
    assert warning.severity.value == "error" and "5" in warning.message
    assert layout.segments[1].value_cols == ((2, 2), (3, 3), (4, 4)) and layout.segments[1].extra_cols == (5,)


def test_second_sheet_single_code_in_other_column(rules: Rules) -> None:
    first = [
        place({0: "Наименование", 3: "Код", 4: "За 9 месяцев 2025 г.", 5: "За 9 месяцев 2024 г."}),
        place({0: "Выручка", 3: "2110", 4: D("100"), 5: D("90")}),
        place({0: "Себестоимость", 3: "2120", 4: D("-60"), 5: D("-50")}),
        place({0: "Валовая прибыль", 3: "2100", 4: D("40"), 5: D("40")}),
    ]
    second = [place({0: "Чистая прибыль", 1: "2400", 2: D("7"), 3: D("6")})]
    cand = excel([("ОФР", first), ("ОФР_2", second)])
    layout = find_layout(cand, parse_header(hdr(), rules), rules)
    assert layout.segments[1].code_col == 1
    assert layout.segments[1].value_cols == ((2, 4), (3, 5))
    assert not [w for w in layout.warnings if w.code.startswith(("SEGMENT_", "HEADER_"))]


def test_header_not_found_anywhere(rules: Rules) -> None:
    rows = [["Запасы", "1210", "5", "4"], ["Дебиторы", "1230", "3", "2"], ["Итого", "1200", "8", "6"]]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [w.code for w in layout.warnings] == ["HEADER_ROW_NOT_FOUND"]
    assert [(c.column_index, c.label) for c in layout.columns] == [(2, ""), (3, "")]


def test_segment_without_codes(rules: Rules) -> None:
    first = [["Наименование", "Код", "На 30.09.2025"], ["Запасы", "1210", "5"], ["Итого", "1200", "5"]]
    tail = [["Примечание к форме", "", "7"]]
    layout = find_layout(excel([("Баланс", first), ("Хвост", tail)]), parse_header(hdr(), rules), rules)
    assert [w.code for w in layout.warnings] == ["SEGMENT_WITHOUT_CODES"]
    assert layout.roles[3] is RowRole.DATA and layout.code_found


def test_first_segment_without_header_columns_from_first_headed(rules: Rules) -> None:
    """P1-18: колонки документа задаёт первый сегмент с заголовком, у первого значения — по порядку."""
    rows = [
        ["Запасы", "1210", "5", "4"],
        ["Дебиторы", "1230", "3", "2"],
        ["Итого", "1200", "8", "6"],
        ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Прочее", "1250", "1", "1"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [s.first_row for s in layout.segments] == [0, 3]
    assert [c.column_index for c in layout.columns] == [2, 3]
    assert layout.segments[0].value_cols == ((2, 2), (3, 3))
    assert layout.warnings == ()


def test_no_codes_anywhere_has_no_segment_warnings(rules: Rules) -> None:
    layout = find_layout(ocr([["Просто", "текст"], ["ещё", "строка"]]), parse_header(hdr(), rules), rules)
    assert not layout.code_found and layout.columns == () and layout.warnings == ()
    assert layout.roles == (RowRole.DATA, RowRole.DATA)