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"] * 4565758def test_impossible_years_are_quarantined_not_deleted() -> None:59    ind = indicators_by_id()["gdp-per-capita"]60    df = _rows("gdp-per-capita", ind.unit, {"CAN": [100.0, 101.0, 102.0, 103.0, 104.0, 105.0]}, 2020)  # 2020–202561    bad = pl.DataFrame({"country_id": ["CAN", "CAN"], "indicator_id": ["gdp-per-capita"] * 2, "period": [date(1700, 1, 1), date(2099, 1, 1)],62                        "year": [1700, 2099], "frequency": ["A", "A"], "value": [50.0, 500.0], "unit": [ind.unit] * 2,63                        "source_id": ["worldbank"] * 2, "source_dataset": ["WDI"] * 2, "source_series_code": ["X"] * 2,64                        "is_estimate": [False] * 2, "is_forecast": [False] * 2, "retrieved_at": [datetime(2026, 1, 1, tzinfo=UTC).replace(tzinfo=None)] * 2,65                        "source_updated_at": [None, None], "status": ["imported"] * 2, "metadata": [None, None]}).cast(df.schema)66    gv = validate_frame(pl.concat([df, bad]), ind, now=datetime(2026, 9, 1, tzinfo=UTC))67    assert gv.frame.height == 8  # nothing deleted68    by_year = dict(zip(gv.frame["year"].to_list(), gv.frame["status"].to_list(), strict=True))69    assert by_year[1700] == "quarantined" and by_year[2099] == "quarantined" and by_year[2024] == "imported"70    assert gv.n_quarantined == 2 and not gv.quarantine_dataset71    assert {i.code for i in gv.issues} >= {"impossible_year"}72    # a projection 6 years out is fine73    assert dict(zip(gv.frame["year"].to_list(), gv.frame["status"].to_list(), strict=True))[2025] == "imported"747576def test_duplicate_keys_quarantine_dataset() -> None:77    ind = indicators_by_id()["gdp-per-capita"]78    df = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]})79    gv = validate_frame(pl.concat([df, df.tail(1)]), ind)80    assert gv.quarantine_dataset and any(i.code == "duplicate" for i in gv.issues)81    assert gv.frame.height == 4  # rows kept for audit828384def test_schema_change_quarantines_dataset() -> None:85    ind = indicators_by_id()["gdp-per-capita"]86    df = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]})87    gv = validate_frame(df.drop("period"), ind)88    assert gv.quarantine_dataset and gv.issues[0].code == "schema_change" and "period" in gv.issues[0].message89    gv2 = validate_frame(df.with_columns(pl.col("value").cast(pl.Utf8)), ind)90    assert gv2.quarantine_dataset and "dtype" in gv2.issues[0].message919293def test_null_spike_and_vintage_shift_are_warnings_only() -> None:94    from countryatlas.pipeline.validate import PreviousStats9596    ind = indicators_by_id()["gdp-per-capita"]97    series = {c: [100.0 + k for k in range(25)] for c in ("CAN", "FRA", "DEU")}98    prev = _rows("gdp-per-capita", ind.unit, series)99    # null spike: 20 % of values null now vs 0 % before100    cur = prev.with_columns(pl.when(pl.col("year") % 5 == 0).then(None).otherwise(pl.col("value")).alias("value"))101    gv = validate_frame(cur, ind, previous=PreviousStats.from_frame(prev))102    assert not gv.quarantine_dataset and any(i.code == "null_spike" for i in gv.issues)103    # vintage shift: every overlapping value 40 % higher → warning, dataset kept104    shifted = prev.with_columns((pl.col("value") * 1.4).alias("value"))105    gv2 = validate_frame(shifted, ind, previous=PreviousStats.from_frame(prev))106    codes = {i.code for i in gv2.issues}107    assert "vintage_shift" in codes and not gv2.quarantine_dataset and gv2.frame.height == shifted.height108    # a 3 % revision is normal109    gv3 = validate_frame(prev.with_columns((pl.col("value") * 1.03).alias("value")), ind, previous=PreviousStats.from_frame(prev))110    assert not any(i.code in ("vintage_shift", "null_spike") for i in gv3.issues)111    # partial download via PreviousStats too112    gv4 = validate_frame(prev.head(5), ind, previous=PreviousStats.from_frame(prev))113    assert gv4.quarantine_dataset and any(i.code == "partial_download" for i in gv4.issues)114