import polars as pl

from wp39.instruments import build_crosswalk, sic_division


def test_sic_divisions():
    assert sic_division(3571) == "D_manufacturing" and sic_division(6021) == "H_finance_insurance_realestate"
    assert sic_division(None) == "unknown"


def test_crosswalk_keeps_all_entities_and_flags_missing_cik():
    ents = pl.DataFrame({"entity_id": ["frd:AAA", "frd:ZZZ"], "security_id": ["firstrate:stock:AAA", "firstrate:stock:ZZZ"],
                         "ticker": ["AAA", "ZZZ"]})
    sec = pl.DataFrame({"ticker": ["AAA"], "cik": [1], "sec_name": ["AAA Corp"]})
    subs = pl.DataFrame({"cik": [1], "sic": [3571], "sic_division": ["D_manufacturing"]})
    xw = build_crosswalk(ents, sec, subs)
    assert xw.height == 2 and xw.filter(pl.col("ticker") == "ZZZ")["has_cik"][0] is False