PaperPod/tests/test_ocr_parsing.py

"""Unit tests for the regex/heuristic parsing logic in paperpod.ocr.extract.

These test the parsing functions directly on strings so they stay
deterministic and don't depend on Tesseract's OCR accuracy. End-to-end OCR
accuracy against rendered receipts is exercised in test_ocr_integration.py.
"""

from paperpod.ocr.extract import (
    clean_vendor,
    extract_date,
    extract_doc_title,
    extract_form_code,
    extract_line_items,
    extract_org_name,
    extract_page_marker,
    extract_party,
    extract_store_name,
    extract_tax_year,
    extract_time,
    extract_total,
)


def test_extract_date_iso():
    assert extract_date("2023-03-14 09:41 #4821") == "2023-03-14"


def test_extract_date_slash_mdy():
    assert extract_date("Purchased on 03/14/2023 at noon") == "2023-03-14"


def test_extract_date_month_name():
    assert extract_date("Invoice date: March 14, 2023") == "2023-03-14"
    assert extract_date("Mar 14 2023") == "2023-03-14"
    # OCR often reads the comma as a period ("July 13. 2026").
    assert extract_date("Date: July 13. 2025") == "2025-07-13"


def test_extract_date_none_when_absent():
    assert extract_date("no date here, just text") is None


def test_extract_date_rejects_implausible_values():
    # Not a real month/day -> not treated as a date.
    assert extract_date("13/45/2023") is None


def test_extract_total_basic():
    text = "SUBTOTAL 72.14\nHST 13% 9.38\nTOTAL 81.52\n"
    assert extract_total(text) == "81.52"


def test_extract_total_ignores_subtotal():
    text = "SUBTOTAL 30.63\n"
    assert extract_total(text) is None


def test_extract_total_word_boundary_with_prefix():
    # "FUEL TOTAL" should still match "total" as its own word.
    text = "FUEL TOTAL 65.10\n"
    assert extract_total(text) == "65.10"


def test_extract_total_amount_on_next_line():
    text = "TOTAL\n81.52\n"
    assert extract_total(text) == "81.52"


def test_extract_total_none_when_absent():
    assert extract_total("no totals here") is None


def test_extract_line_items_skips_totals_and_keeps_skus():
    text = (
        "COSTCO WHOLESALE\n"
        "1234567 LEGO CLASSIC SET 24.99 H\n"
        "KIRKLAND PAPER TOWEL 18.99\n"
        "SUBTOTAL 43.98\n"
        "HST 5.72\n"
        "TOTAL 49.70\n"
        "MASTERCARD 49.70\n"
    )
    items = extract_line_items(text)
    assert items == [
        ("LEGO CLASSIC SET", "24.99"),
        ("KIRKLAND PAPER TOWEL", "18.99"),
    ]


def test_extract_time_hms():
    assert extract_time("03/20/23  15:25:35") == "15:25"


def test_extract_time_hm():
    assert extract_time("Transaction at 09:05 today") == "09:05"


def test_extract_time_zero_pads_single_digit_hour():
    assert extract_time("open at 9:05am") == "09:05"


def test_extract_time_none_when_absent():
    assert extract_time("no time here, just $103.60 total") is None


def test_extract_time_ignores_colonless_numbers():
    assert extract_time("TC# 3821 2498 6394 0609 7307") is None


def test_extract_form_code_t4():
    text = "Employer's name\nYORK UNIVERSITY\nT4\nStatement of Remuneration Paid"
    assert extract_form_code(text) == "T4"


def test_extract_form_code_prefers_more_specific_match():
    # T4A should win over the more generic T4 substring it contains.
    text = "Payer's name\nT4A\nStatement of Pension, Retirement, Annuity"
    assert extract_form_code(text) == "T4A"


def test_extract_form_code_t5008():
    text = "BOX 14:\nT5008\nSUMMARY OF SECURITY DISPOSITIONS 2023"
    assert extract_form_code(text) == "T5008"


def test_extract_form_code_none_for_receipt():
    assert extract_form_code("WALMART\nSUBTOTAL 21.97\nTOTAL 23.71") is None


def test_extract_form_code_requires_slip_boilerplate():
    # Garbled OCR of a bank statement once misread "LEVIT I" as "T4";
    # without CRA boilerplate nearby the code must not count.
    assert extract_form_code("HORT NAMY nV T4\nRAN AMOUNY HATL") is None


def test_extract_org_name_finds_employer():
    text = (
        "Employer's name\nYORK UNIVERSITY Year 2023 Statement of Remuneration Paid\n"
        "4700 KEELE STREET\nTORONTO ON M3J 1P3"
    )
    assert extract_org_name(text) == "YORK UNIVERSITY"


def test_extract_org_name_skips_addresses():
    text = "MR ILIA DOBKIN\n153 NIAGARA DR\nOSHAWA ON L1G 8A6"
    assert extract_org_name(text) == "MR ILIA DOBKIN"


def test_extract_org_name_skips_boilerplate_and_barcodes():
    text = (
        "JTA9500902-0301613-26631-0004-0004-00-\n"
        "Canada Revenue Agence du revenu\n"
        "Agency du Canada\n"
        "AMC ENTERTAINMENT HOLDINGS INC AMC PRFRD EQT"
    )
    assert extract_org_name(text) == "AMC ENTERTAINMENT HOLDINGS INC AMC"


def test_extract_org_name_none_when_nothing_plausible():
    assert extract_org_name("Canada Revenue Agency\nStatement of Remuneration Paid") is None


def test_extract_tax_year_from_t4():
    text = "YORK UNIVERSITY Year 2023 Statement of Remuneration Paid"
    assert extract_tax_year(text) == "2023"


def test_extract_tax_year_from_t5008_title():
    text = "T5008\nSUMMARY OF SECURITY DISPOSITIONS 2023"
    assert extract_tax_year(text) == "2023"


def test_extract_store_name_costco():
    assert extract_store_name("COSTCO WHOLESALE\nVaughan #547") == "Costco"


def test_extract_store_name_walmart():
    assert extract_store_name("HOW DID WE DO TODAY?\nWalmart") == "Walmart"


def test_extract_date_rejects_far_future_ocr_misreads():
    # "2053" is a misread of "2023" — a document can't be dated decades ahead.
    assert extract_date("PERIOD ENDING 2053-06-30") is None


def test_extract_date_statement_period_end_wins():
    # The period end is the statement's document date, not the first
    # transaction row.
    text = "PERIOD: FROM : 01/ 01/ 2025 TO : 01/ 31/ 2025\n01/02/2025 PYT TO"
    assert extract_date(text) == "2025-01-31"


def test_extract_date_tolerates_ocr_spaces_and_letter_o():
    assert extract_date("TO : O1/ 31/2025") == "2025-01-31"
    assert extract_date("12/ 31/2024") == "2024-12-31"


def test_extract_page_marker_english():
    assert extract_page_marker("DEPOSIT ACCOUNT HISTORY   PAGE  1 OF   2") == (1, 2)
    assert extract_page_marker("PAGE 13 OF 17") == (13, 17)


def test_extract_page_marker_french():
    assert extract_page_marker("HISTORIQUE DES OPERATIONS PAGE 2 DE 3") == (2, 3)


def test_extract_page_marker_rejects_nonsense():
    assert extract_page_marker("PAGE 5 OF 2") is None
    assert extract_page_marker("no marker here") is None


def test_extract_doc_title_transaction_history():
    assert extract_doc_title("TD Canada Trust\nTransaction History") == "transaction_history"
    assert extract_doc_title("Historique des opérations") == "transaction_history"


def test_extract_doc_title_direct_deposit():
    assert (
        extract_doc_title("Government of Canada - Direct Deposit Enrolment")
        == "direct_deposit"
    )


def test_extract_doc_title_none_for_receipt():
    assert extract_doc_title("WALMART\nSUBTOTAL 21.97") is None


def test_extract_party_shortname():
    text = "BR # : 1076 ACCOUNT: 6254276 MBA - MIN SHORTNAME : LEVIT I\nPERIOD: FROM"
    assert extract_party(text) == "LEVIT I"


def test_extract_party_none_when_absent():
    assert extract_party("WALMART\nTOTAL 23.71") is None


def test_clean_vendor_strips_edge_junk():
    assert clean_vendor("| ») Transaction History i") == "Transaction History"
    assert clean_vendor("7D TD Canada Trust") == "TD Canada Trust"
    assert clean_vendor("+7 TD Canada Trus") == "TD Canada Trus"
    assert clean_vendor(": TD Canada Trust") == "TD Canada Trust"


def test_clean_vendor_rejects_gibberish():
    assert clean_vendor("BH tc rn te So Gee EA Nie PR ae UE aa a RR Pia") is None


def test_clean_vendor_keeps_real_names():
    assert clean_vendor("TD Canada Trust") == "TD Canada Trust"
    assert clean_vendor("Costco") == "Costco"
    assert clean_vendor("Historique des opérations") == "Historique des opérations"


def test_clean_vendor_none_passthrough():
    assert clean_vendor(None) is None
    assert clean_vendor("|| 123") is None