from __future__ import annotations from datetime import UTC, date, datetime import polars as pl from countryatlas.models import NormalizedObservation from countryatlas.pipeline.staging import rows_to_frame from countryatlas.pipeline.validate import validate_frame from countryatlas.registry import indicators_by_id def _rows(ind: str, unit: str, series: dict[str, list[float]], start: int = 2000) -> pl.DataFrame: now = datetime.now(UTC) out = [] for c, vals in series.items(): for k, v in enumerate(vals): out.append(NormalizedObservation(country_id=c, indicator_id=ind, period=date(start + k, 1, 1), year=start + k, frequency="A", value=v, unit=unit, source_id="worldbank", source_dataset="WDI", source_series_code="X", retrieved_at=now)) return rows_to_frame(out) def test_bounds_quarantine_and_jump_warning() -> None: ind = indicators_by_id()["life-expectancy"] # bounds roughly [20, 100] 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], "FRA": [78, 78.4, 78.9, 79.3, 79.7, 80.1, 80.5, 60.0, 81.2, 81.5]}, 2010) gv = validate_frame(df, ind, now=datetime(2021, 6, 1, tzinfo=UTC)) st = dict(zip(gv.frame["country_id"].to_list(), gv.frame["status"].to_list(), strict=True)) # last per country wins assert gv.frame.filter((pl.col("country_id") == "CAN") & (pl.col("year") == 2019))["status"][0] == "quarantined" assert gv.frame.filter((pl.col("country_id") == "FRA") & (pl.col("year") == 2017))["status"][0] == "warning" assert gv.n_quarantined == 1 and gv.n_warning >= 1 assert not gv.quarantine_dataset assert {i.code for i in gv.issues} >= {"out_of_bounds", "extreme_jump"} assert st # sanity def test_unit_mismatch_and_partial_download_quarantine_dataset() -> None: ind = indicators_by_id()["gdp-per-capita"] df = _rows("gdp-per-capita", "wrong unit", {"CAN": [1.0, 2.0, 3.0]}) gv = validate_frame(df, ind) assert gv.quarantine_dataset and any(i.code == "unit_mismatch" for i in gv.issues) df2 = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]}) gv2 = validate_frame(df2, ind, previous_rows=100) assert gv2.quarantine_dataset and any(i.code == "partial_download" for i in gv2.issues) gv3 = validate_frame(df2, ind, previous_rows=5) assert not gv3.quarantine_dataset def test_stale_flag_on_latest_row_only() -> None: ind = indicators_by_id()["gdp-per-capita"] df = _rows("gdp-per-capita", ind.unit, {"CAN": [100.0, 101.0, 102.0, 103.0, 104.0]}, 2010) # latest 2014 gv = validate_frame(df, ind, now=datetime(2026, 9, 1, tzinfo=UTC)) statuses = gv.frame.sort("period")["status"].to_list() assert statuses[-1] == "stale" and statuses[:-1] == ["imported"] * 4 def test_impossible_years_are_quarantined_not_deleted() -> None: ind = indicators_by_id()["gdp-per-capita"] df = _rows("gdp-per-capita", ind.unit, {"CAN": [100.0, 101.0, 102.0, 103.0, 104.0, 105.0]}, 2020) # 2020–2025 bad = pl.DataFrame({"country_id": ["CAN", "CAN"], "indicator_id": ["gdp-per-capita"] * 2, "period": [date(1700, 1, 1), date(2099, 1, 1)], "year": [1700, 2099], "frequency": ["A", "A"], "value": [50.0, 500.0], "unit": [ind.unit] * 2, "source_id": ["worldbank"] * 2, "source_dataset": ["WDI"] * 2, "source_series_code": ["X"] * 2, "is_estimate": [False] * 2, "is_forecast": [False] * 2, "retrieved_at": [datetime(2026, 1, 1, tzinfo=UTC).replace(tzinfo=None)] * 2, "source_updated_at": [None, None], "status": ["imported"] * 2, "metadata": [None, None]}).cast(df.schema) gv = validate_frame(pl.concat([df, bad]), ind, now=datetime(2026, 9, 1, tzinfo=UTC)) assert gv.frame.height == 8 # nothing deleted by_year = dict(zip(gv.frame["year"].to_list(), gv.frame["status"].to_list(), strict=True)) assert by_year[1700] == "quarantined" and by_year[2099] == "quarantined" and by_year[2024] == "imported" assert gv.n_quarantined == 2 and not gv.quarantine_dataset assert {i.code for i in gv.issues} >= {"impossible_year"} # a projection 6 years out is fine assert dict(zip(gv.frame["year"].to_list(), gv.frame["status"].to_list(), strict=True))[2025] == "imported" def test_duplicate_keys_quarantine_dataset() -> None: ind = indicators_by_id()["gdp-per-capita"] df = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]}) gv = validate_frame(pl.concat([df, df.tail(1)]), ind) assert gv.quarantine_dataset and any(i.code == "duplicate" for i in gv.issues) assert gv.frame.height == 4 # rows kept for audit def test_schema_change_quarantines_dataset() -> None: ind = indicators_by_id()["gdp-per-capita"] df = _rows("gdp-per-capita", ind.unit, {"CAN": [1.0, 2.0, 3.0]}) gv = validate_frame(df.drop("period"), ind) assert gv.quarantine_dataset and gv.issues[0].code == "schema_change" and "period" in gv.issues[0].message gv2 = validate_frame(df.with_columns(pl.col("value").cast(pl.Utf8)), ind) assert gv2.quarantine_dataset and "dtype" in gv2.issues[0].message def test_null_spike_and_vintage_shift_are_warnings_only() -> None: from countryatlas.pipeline.validate import PreviousStats ind = indicators_by_id()["gdp-per-capita"] series = {c: [100.0 + k for k in range(25)] for c in ("CAN", "FRA", "DEU")} prev = _rows("gdp-per-capita", ind.unit, series) # null spike: 20 % of values null now vs 0 % before cur = prev.with_columns(pl.when(pl.col("year") % 5 == 0).then(None).otherwise(pl.col("value")).alias("value")) gv = validate_frame(cur, ind, previous=PreviousStats.from_frame(prev)) assert not gv.quarantine_dataset and any(i.code == "null_spike" for i in gv.issues) # vintage shift: every overlapping value 40 % higher → warning, dataset kept shifted = prev.with_columns((pl.col("value") * 1.4).alias("value")) gv2 = validate_frame(shifted, ind, previous=PreviousStats.from_frame(prev)) codes = {i.code for i in gv2.issues} assert "vintage_shift" in codes and not gv2.quarantine_dataset and gv2.frame.height == shifted.height # a 3 % revision is normal gv3 = validate_frame(prev.with_columns((pl.col("value") * 1.03).alias("value")), ind, previous=PreviousStats.from_frame(prev)) assert not any(i.code in ("vintage_shift", "null_spike") for i in gv3.issues) # partial download via PreviousStats too gv4 = validate_frame(prev.head(5), ind, previous=PreviousStats.from_frame(prev)) assert gv4.quarantine_dataset and any(i.code == "partial_download" for i in gv4.issues)