stack/tests/lead/test_load_strand5.py

"""Strand-5 files load into queryable tables. Absent series stay absent."""

from __future__ import annotations

import hashlib
import io
import zipfile
from pathlib import Path

import duckdb
import openpyxl

from bls.ces import read_series
from bls.oews import read_occupations
from lead.load.irs import choose_irs_members
from lead.load.strand5 import load_strand5

_TSV = "source\turl\tpath\tbytes\tsha256\tstatus\tnote\n"


def _sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _line(source: str, path: Path | str, status: str, *, note: str = "") -> str:
    if isinstance(path, Path):
        digest = _sha(path) if path.exists() else ""
        size = path.stat().st_size if path.exists() else ""
        shown = str(path)
    else:
        digest = ""
        size = ""
        shown = path
    return f"{source}\thttp://example.invalid\t{shown}\t{size}\t{digest}\t{status}\t{note}\n"


def _oews_zip(path: Path) -> None:
    wb = openpyxl.Workbook()
    sheet = wb.active
    sheet.append(
        ["AREA", "OCC_CODE", "OCC_TITLE", "O_GROUP", "TOT_EMP", "H_MEAN", "A_MEAN"]
    )
    sheet.append(["99", "29-1141", "Registered Nurses", "detailed", 100, 40, 80000])
    buf = io.BytesIO()
    wb.save(buf)
    with zipfile.ZipFile(path, "w") as zf:
        zf.writestr("oesm25nat/national_M2025_dl.xlsx", buf.getvalue())


def test_absent_rows_keep_distinct_urls(tmp_path):
    root = tmp_path
    (root / "files.tsv").write_text(
        _TSV
        + "qcew\thttps://data.bls.gov/cew/data/files/1975/csv/1975_annual_by_industry.zip\t\t\t\thttp_404\tabsent\n"
        + "qcew\thttps://data.bls.gov/cew/data/files/1976/csv/1976_annual_by_industry.zip\t\t\t\thttp_404\tabsent\n"
        + "fred\thttps://fred.stlouisfed.org/graph/fredgraph.csv?id=UNRATE\t\t\t\tunreachable\tno key\n"
    )
    con = duckdb.connect()
    try:
        load_strand5(con, root, actor="tester")
        load_strand5(con, root, actor="tester")
        rows = con.execute(
            """
            SELECT source_id, path, status
            FROM lead_public.load_log
            ORDER BY source_id, path
            """
        ).fetchall()
    finally:
        con.close()
    assert len(rows) == 3
    assert [row[1].rsplit("/", 1)[-1] for row in rows if row[0] == "qcew"] == [
        "1975_annual_by_industry.zip",
        "1976_annual_by_industry.zip",
    ]
    assert rows[0][0] == "fred" and rows[0][2] == "absent"


def test_choose_irs_members_keeps_national_and_cbsa_files():
    members = [
        "13incyallagi.csv",
        "13incyallnoagi.csv",
        "13incyall.xls",
        "13incyal.xls",
        "13incbsa.csv",
        "11incydocguide.doc",
    ]
    assert choose_irs_members(members) == [
        "13incyallagi.csv",
        "13incyallnoagi.csv",
        "13incbsa.csv",
    ]


def test_a_zipcode_csv_is_the_national_extract():
    assert choose_irs_members(
        [
            "zipcode07.csv",
            "ZIP Code 2007 TX.xls",
            "ZIP Code 2007 CA.xls",
        ]
    ) == ["zipcode07.csv"]


def test_a_state_only_zip_keeps_the_state_workbooks():
    assert choose_irs_members(
        ["1998ZIPCode/", "1998ZIPCode/98zp01al.xls", "1998ZIPCode/98zp02ak.xls"]
    ) == ["1998ZIPCode/98zp01al.xls", "1998ZIPCode/98zp02ak.xls"]


def test_strand5_loads_each_family_and_skips_what_is_not_a_file(tmp_path):
    root = tmp_path / "strand5"
    (root / "bls").mkdir(parents=True)
    (root / "qcew").mkdir()
    (root / "irs_soi" / "county").mkdir(parents=True)
    (root / "irs_soi" / "zip").mkdir()
    (root / "ma_penetration" / "2024-06").mkdir(parents=True)
    (root / "fred").mkdir()

    ppi = root / "bls" / "ppi.csv"
    ppi.write_text(
        "series_id,year,period,value\nWPS512,2024,M01,100.5\nWPS512,2024,M02,\n"
    )
    jolts = root / "bls" / "jolts.csv"
    jolts.write_text(
        "series_id,year,period,value\nJTS620000000000000JOL,2024,M12,1410\n"
    )
    ces = root / "bls" / "ces.csv"
    ces.write_text(
        "series_id,year,period,value\nCES6562000001,2024,M01,17000\nCES6562000001,2024,M13,16900\n"
    )
    oews = tmp_path / "oesm25nat.zip"
    _oews_zip(oews)

    qcew = root / "qcew" / "2024_annual_by_industry.zip"
    with zipfile.ZipFile(qcew, "w") as zf:
        zf.writestr(
            "2024.annual.by_industry/2024.annual 10 Total, all industries.csv",
            "area_fips,year,qtr,annual_avg_emplvl\n01000,2024,A,10\n01001,2024,A,3\n",
        )

    county = root / "irs_soi" / "county" / "23incyallagi.csv"
    county.write_text("STATEFIPS,COUNTYFIPS,A00100\n01,001,100\n01,000,50\n")
    zcta = root / "irs_soi" / "zip" / "23zpallagi.csv"
    zcta.write_text("STATEFIPS,ZIPCODE,A00100\n01,30301,7\n")
    bundled = root / "irs_soi" / "county" / "county2023.zip"
    with zipfile.ZipFile(bundled, "w") as zf:
        zf.writestr("23incyallagi.csv", county.read_text())
        zf.writestr("al23.xls", b"not-read")

    pen = root / "ma_penetration" / "2024-06" / "county-penetration-june-2024.zip"
    with zipfile.ZipFile(pen, "w") as zf:
        zf.writestr("readme.txt", "notes\n")
        zf.writestr(
            "State_County_Penetration_MA_2024_06.csv",
            "FIPS,Eligibles,Enrolled,Penetration\n01001,100,40,0.4\n",
        )

    tsv = root / "files.tsv"
    tsv.write_text(
        _TSV
        + _line("bls_ppi", ppi, "ok")
        + _line("bls_jolts", jolts, "ok")
        + _line("bls_ces", ces, "ok")
        + _line("bls_oews", oews, "ok", note="May 2025 national")
        + _line("qcew", qcew, "ok")
        + _line("qcew", "1975_annual_by_industry.zip", "http_404", note="absent")
        + _line("irs_soi", county, "ok", note="county")
        + _line("irs_soi", zcta, "ok", note="zip")
        + _line("irs_soi", bundled, "ok", note="county")
        + _line("ma_penetration", pen, "ok")
        + _line("ma_penetration", pen, "http_22", note="curl 22")
        + _line("fred", "", "unreachable", note="no file")
        + _line(
            "cms_mds_frequency",
            "/no/such/mds.csv",
            "ok",
            note="covered by the enrichment manifest",
        )
    )
    # The unreachable and 404 rows have an empty path in the real inventory.
    text = tsv.read_text()
    text = text.replace(
        _line("qcew", "1975_annual_by_industry.zip", "http_404", note="absent"),
        "qcew\thttp://example.invalid\t\t\t\thttp_404\tabsent\n",
    )
    text = text.replace(
        _line("fred", "", "unreachable", note="no file"),
        "fred\thttp://example.invalid\t\t\t\tunreachable\tno file\n",
    )
    # The duplicate-hash http_22 row must not carry the ok row's digest as a
    # second load of the same bytes. Record it as a failed attempt with no hash.
    text = text.replace(
        _line("ma_penetration", pen, "http_22", note="curl 22"),
        f"ma_penetration\thttp://example.invalid\t{pen}\t\t\thttp_22\tcurl 22\n",
    )
    tsv.write_text(text)

    con = duckdb.connect()
    try:
        first = load_strand5(con, root, actor="tester")
        second = load_strand5(con, root, actor="tester")
        ppi_row = con.execute(
            "SELECT series_id, spine_year, spine_month, value FROM bls.ppi ORDER BY spine_month"
        ).fetchall()
        jolts_row = con.execute(
            "SELECT series_id, value, time_key FROM bls.jolts"
        ).fetchone()
        ces_rows = read_series(con, ["CES6562000001"])
        oews_rows = read_occupations(con, ["29-1141"])
        qcew_rows = con.execute(
            """
            SELECT geo_fips, spine_year, file_year, annual_avg_emplvl
            FROM bls.qcew
            ORDER BY geo_fips NULLS FIRST
            """
        ).fetchall()
        irs_rows = con.execute(
            """
            SELECT geo_fips, geo_zip, file_year, a00100
            FROM econ.irs_soi
            ORDER BY geo_fips NULLS LAST, geo_zip NULLS LAST
            """
        ).fetchall()
        pen_row = con.execute(
            """
            SELECT geo_fips, spine_year, spine_month, spine_quarter, time_key, penetration
            FROM lead_public.ma_county_penetration
            """
        ).fetchone()
        fred = con.execute(
            """
            SELECT COUNT(*) FROM information_schema.tables
            WHERE table_schema = 'econ' AND table_name = 'fred'
            """
        ).fetchone()[0]
        fred_log = con.execute(
            """
            SELECT status, reason, count(*)
            FROM lead_public.load_log
            WHERE source_id = 'fred'
            GROUP BY 1, 2
            """
        ).fetchall()
    finally:
        con.close()

    assert ppi_row == [("WPS512", 2024, 1, 100.5), ("WPS512", 2024, 2, None)]
    assert jolts_row == ("JTS620000000000000JOL", 1410.0, "2024Q4")
    assert [(r.year, r.month, r.value) for r in ces_rows] == [
        (2024, 1, 17000.0),
        (2024, 13, 16900.0),
    ]
    assert len(oews_rows) == 1 and oews_rows[0].year == 2025
    assert qcew_rows == [(None, 2024, 2024, "10"), ("01001", 2024, 2024, "3")]
    assert irs_rows == [
        ("01001", None, 2023, "100"),
        (None, "30301", 2023, "7"),
        (None, None, 2023, "50"),
    ]
    assert pen_row == ("01001", 2024, 6, 2, "2024Q2", "0.4")
    assert fred == 0
    assert fred_log == [("absent", "unreachable", 1)]
    by_status = {}
    for row in first:
        by_status.setdefault(row.status, []).append(row.source_id)
    assert "fred" in by_status["absent"] or "fred" in by_status.get("unreachable", [])
    assert "qcew" in by_status["absent"] or any(
        row.source_id == "qcew" and row.status == "absent" for row in first
    )
    assert any(
        row.source_id == "cms_mds_frequency" and row.status == "skipped"
        for row in first
    )
    assert any(
        row.source_id == "ma_penetration" and row.status == "skipped" for row in first
    )
    assert sum(row.rows for row in second) == 0
    assert all(row.status in {"loaded", "skipped", "absent"} for row in second)