spb/countryatlas

1from __future__ import annotations23from datetime import UTC, date, datetime45import polars as pl67from countryatlas.models import NormalizedObservation8from countryatlas.pipeline.staging import rows_to_frame9from countryatlas.pipeline.validate import validate_frame10from countryatlas.registry import indicators_by_id111213def _rows(ind: str, unit: str, series: dict[str, list[float]], start: int = 2000) -> pl.DataFrame:14    now = datetime.now(UTC)15    out = []16    for c, vals in series.items():17        for k, v in enumerate(vals):18            out.append(NormalizedObservation(country_id=c, indicator_id=ind, period=date(start + k, 1, 1), year=start + k,19                                             frequency="A", value=v, unit=unit, source_id="worldbank", source_dataset="WDI",20                                             source_series_code="X", retrieved_at=now))21    return rows_to_frame(out)222324def test_bounds_quarantine_and_jump_warning() -> None:25    ind = indicators_by_id()["life-expectancy"]  # bounds roughly [20, 100]26    df = _rows("life-expectancy", ind.unit, {"CAN": [78, 78.5, 79, 79.4, 79.9, 80.2, 80.6, 81.0, 81.3, 250.0],27                                               "FRA": [78, 78.4, 78.9, 79.3, 79.7, 80.1, 80.5, 60.0, 81.2, 81.5]}, 2010)28    gv = validate_frame(df, ind, now=datetime(2021, 6, 1, tzinfo=UTC))29    st = dict(zip(gv.frame["country_id"].to_list(), gv.frame["status"].to_list(), strict=True))  # last per country wins30    assert gv.frame.filter((pl.col("country_id") == "CAN") & (pl.col("year") == 2019))["status"][0] == "quarantined"31    assert gv.frame.filter((pl.col("country_id") == "FRA") & (pl.col("year") == 2017))["status"][0] == "warning"32    assert gv.n_quarantined == 1 and gv.n_warning >= 133    assert not gv.quarantine_dataset34    assert {i.code for i in gv.issues} >= {"out_of_bounds", "extreme_jump"}35    assert st  # sanity363738def test_unit_mismatch_and_partial_download_quarantine_dataset() -> None:39    ind = indicators_by_id()["gdp-per-capita"]40    df = _rows("gdp-per-capita", "wrong unit", {"CAN": [1.0, 2.0, 3.0]})41    gv = validate_frame(df, ind)42    assert gv.quarantine_dataset and any(i.code == "unit_mismatch" for i in gv.issues)43    df2 = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]})44    gv2 = validate_frame(df2, ind, previous_rows=100)45    assert gv2.quarantine_dataset and any(i.code == "partial_download" for i in gv2.issues)46    gv3 = validate_frame(df2, ind, previous_rows=5)47    assert not gv3.quarantine_dataset484950def test_stale_flag_on_latest_row_only() -> None:51    ind = indicators_by_id()["gdp-per-capita"]52    df = _rows("gdp-per-capita", ind.unit, {"CAN": [100.0, 101.0, 102.0, 103.0, 104.0]}, 2010)  # latest 201453    gv = validate_frame(df, ind, now=datetime(2026, 9, 1, tzinfo=UTC))54    statuses = gv.frame.sort("period")["status"].to_list()55    assert statuses[-1] == "stale" and statuses[:-1] == ["imported"] * 456