spb/countryatlas

1from __future__ import annotations23import json4from datetime import UTC, date, datetime5from pathlib import Path67import duckdb8import pytest910from countryatlas.config import settings11from countryatlas.models import ImportRun, IndicatorSourceSpec, NormalizedObservation12from countryatlas.pipeline import build as build_mod13from countryatlas.pipeline.staging import rows_to_frame, spec_paths, write_parquet_atomic, write_run1415COUNTRIES = ["CAN", "USA", "FRA", "DEU", "JPN", "BRA", "IND", "NGA", "AUS", "MEX", "KOR", "ITA", "ESP", "GBR", "CHN",16             "ZAF", "EGY", "TUR", "ARG", "IDN", "SWE", "NOR", "CHL", "POL"]171819def _stage(spec: IndicatorSourceSpec, unit: str, base: float, growth: float, years: range, forecast_from: int | None = None,20           skip: tuple[str, ...] = (), stop_at: dict[str, int] | None = None) -> None:21    now = datetime.now(UTC)22    rows = []23    for k, c in enumerate(COUNTRIES):24        if c in skip:25            continue26        for y in years:27            if stop_at and c in stop_at and y > stop_at[c]:28                continue29            v = base * (1 + 0.05 * k) * (growth ** (y - years.start))30            rows.append(NormalizedObservation(country_id=c, indicator_id=spec.indicator_id, period=date(y, 1, 1), year=y,31                                              frequency="A", value=v, unit=unit, source_id=spec.connector,32                                              source_dataset=spec.dataset, source_series_code=spec.code, retrieved_at=now,33                                              is_forecast=bool(forecast_from and y >= forecast_from)))34    p = spec_paths(spec)35    write_parquet_atomic(rows_to_frame(rows), p["parquet"])36    write_run(spec, ImportRun(run_id="T1", connector=spec.connector, dataset=spec.dataset, started_at=now, finished_at=now,37                              status="ok", rows_norm=len(rows), rows_valid=len(rows)))383940def _tiny_staging() -> None:41    # WB: no data for POL (IMF takes the whole series); CHL stops in 2018 (IMF is > 3 years fresher → "fresher")42    wb = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD")43    _stage(wb, "current US$", 10_000, 1.03, range(2000, 2025), skip=("POL",), stop_at={"CHL": 2018})44    # a lower-priority alternative source for the same indicator (complete series → observations_alt where WB wins)45    imf = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="imf", dataset="WEO", code="NGDPDPC", priority=2)46    _stage(imf, "current US$", 10_100, 1.03, range(2000, 2027), forecast_from=2025)47    pop = IndicatorSourceSpec(indicator_id="population", connector="worldbank", dataset="WDI", code="SP.POP.TOTL")48    _stage(pop, "people", 5_000_000, 1.01, range(2000, 2025))49    le = IndicatorSourceSpec(indicator_id="life-expectancy", connector="worldbank", dataset="WDI", code="SP.DYN.LE00.IN")50    _stage(le, "years", 70, 1.002, range(2000, 2024))51    ma = IndicatorSourceSpec(indicator_id="median-age", connector="owid", dataset="grapher", code="median-age")52    _stage(ma, "years", 25, 1.005, range(2000, 2024))53    fr = IndicatorSourceSpec(indicator_id="fertility-rate", connector="worldbank", dataset="WDI", code="SP.DYN.TFRT.IN")54    _stage(fr, "births per woman", 3.0, 0.99, range(2000, 2024))55    # annual inflation for everyone + a MONTHLY FRED series for the USA (must stay out of latest/rankings)56    infl = IndicatorSourceSpec(indicator_id="inflation", connector="worldbank", dataset="WDI", code="FP.CPI.TOTL.ZG")57    _stage(infl, "annual %", 2.0, 1.01, range(2000, 2026))58    # CAN alone already has 2026 → its rank must fall back to 2025 (the last year with ≥ 20 ranked countries)59    import polars as pl6061    p = spec_paths(infl)["parquet"]62    df = pl.read_parquet(p)63    extra = df.filter((pl.col("country_id") == "CAN") & (pl.col("year") == 2025)).with_columns(64        pl.lit(date(2026, 1, 1)).alias("period"), pl.lit(2026, dtype=pl.Int32).alias("year"), (pl.col("value") + 1.0).alias("value"))65    write_parquet_atomic(pl.concat([df, extra]), p)66    fred = IndicatorSourceSpec(indicator_id="inflation", connector="fred", dataset="FRED", code="CPIAUCSL", priority=2,67                               countries=["USA"], frequency="M")68    now = datetime.now(UTC)69    monthly = [NormalizedObservation(country_id="USA", indicator_id="inflation", period=date(2026, m, 1), year=2026, frequency="M",70                                     value=3.0 + m / 10, unit="annual %", source_id="fred", source_dataset="FRED",71                                     source_series_code="CPIAUCSL", retrieved_at=now) for m in range(1, 8)]72    write_parquet_atomic(rows_to_frame(monthly), spec_paths(fred)["parquet"])73    write_run(fred, ImportRun(run_id="T1", connector="fred", dataset="FRED", started_at=now, finished_at=now, status="ok",74                              rows_norm=len(monthly), rows_valid=len(monthly)))757677@pytest.mark.usefixtures("data_dir")78def test_build_tiny_staging_produces_all_tables() -> None:79    _tiny_staging()80    r = build_mod.build(run_id="20260101T000000Z", strict=False)81    assert settings.db_path.exists()82    assert r.snapshot_path is not None and r.snapshot_path.exists()83    con = duckdb.connect(str(settings.db_path), read_only=True)84    try:85        tables = {t[0] for t in con.execute("SELECT table_name FROM information_schema.tables").fetchall()}86        for t in ("countries", "groups", "group_members", "sources", "indicators", "indicator_sources", "observations",87                  "observations_alt", "observation_revisions", "latest", "rankings", "changes", "events", "similarity",88                  "insights", "country_dna", "coverage", "import_runs", "validation_issues", "search_index", "meta"):89            assert t in tables, t90        n_obs = con.execute("SELECT count(*) FROM observations").fetchone()[0]91        n_alt = con.execute("SELECT count(*) FROM observations_alt").fetchone()[0]92        assert n_obs > 0 and n_alt > 0  # IMF series lose the priority race where WB has data93        # one source per (country, indicator, frequency) series — never spliced94        assert con.execute("SELECT count(*) FROM (SELECT country_id, indicator_id, frequency FROM observations "95                           "GROUP BY ALL HAVING count(DISTINCT source_id) > 1)").fetchone()[0] == 096        assert con.execute("SELECT count(*) FROM observations_alt WHERE country_id='CAN' AND source_id='imf'").fetchone()[0] == 2797        src = dict(con.execute("SELECT country_id, source_id FROM observations WHERE indicator_id='gdp-per-capita' "98                               "GROUP BY ALL").fetchall())99        assert src["CAN"] == "worldbank" and src["POL"] == "imf" and src["CHL"] == "imf"100        reasons = dict(con.execute("SELECT country_id, json_extract_string(metadata, '$.merge_reason') FROM observations "101                                   "WHERE indicator_id='gdp-per-capita' GROUP BY ALL").fetchall())102        assert reasons["CAN"] == "priority" and reasons["POL"] == "priority" and reasons["CHL"] == "fresher"103        # forecasts (from the chosen source) kept in observations but excluded from latest/rankings104        assert con.execute("SELECT count(*) FROM observations WHERE is_forecast").fetchone()[0] == 2 * 2105        assert con.execute("SELECT count(*) FROM latest WHERE is_forecast").fetchone()[0] == 0106        assert con.execute("SELECT max(year) FROM latest WHERE indicator_id='gdp-per-capita'").fetchone()[0] == 2024107        lat = con.execute("SELECT rank_world, n_world, change_10y_pct FROM latest WHERE country_id='CAN' AND indicator_id='gdp-per-capita'").fetchone()108        assert lat[0] is not None and lat[1] == len(COUNTRIES) and lat[2] is not None109        assert con.execute("SELECT count(*) FROM rankings WHERE indicator_id='gdp-per-capita' AND year=2024").fetchone()[0] == len(COUNTRIES)110        # rank 1 = highest value (higher_is_better true for gdp-per-capita)111        top = con.execute("SELECT country_id FROM rankings WHERE indicator_id='gdp-per-capita' AND year=2024 AND rank=1").fetchone()[0]112        assert top == COUNTRIES[-1]113        assert con.execute("SELECT count(*) FROM coverage").fetchone()[0] == len(con.execute("SELECT * FROM countries").fetchall())114        assert con.execute("SELECT count(*) FROM search_index WHERE type='country'").fetchone()[0] >= 200115        assert con.execute("SELECT count(*) FROM insights").fetchone()[0] > 0116        assert con.execute("SELECT count(*) FROM country_dna").fetchone()[0] > 0117        assert con.execute("SELECT count(*) FROM import_runs").fetchone()[0] == 8118        assert r.counts["series_fresher_source"] == 1119        # canonical frequency only: the USA monthly FRED series is in observations but not in latest/rankings120        assert con.execute("SELECT count(*) FROM observations WHERE frequency='M'").fetchone()[0] == 7121        usa = con.execute("SELECT year, frequency, source_id, n_world FROM latest WHERE country_id='USA' AND indicator_id='inflation'").fetchone()122        assert usa == (2025, "A", "worldbank", len(COUNTRIES))123        assert con.execute("SELECT count(*) FROM latest l JOIN indicators i ON i.id=l.indicator_id WHERE l.frequency<>i.frequency").fetchone()[0] == 0124        assert con.execute("SELECT count(*) FROM rankings WHERE indicator_id='inflation' AND year=2026").fetchone()[0] == 0125        # rank fallback: CAN's latest inflation is 2026 (only country) → ranked within 2025 among all 24126        can = con.execute("SELECT year, rank_year, n_world FROM latest WHERE country_id='CAN' AND indicator_id='inflation'").fetchone()127        assert can[0] == 2026 and can[1] == 2025 and can[2] == len(COUNTRIES)128        meta = dict(con.execute("SELECT key, value FROM meta").fetchall())129        assert meta["schema_version"] == "1" and meta["build_run_id"] == "20260101T000000Z"130        assert int(meta["observation_count"]) == n_obs131        ind = con.execute("SELECT n_countries, last_year, primary_source_id FROM indicators WHERE id='gdp-per-capita'").fetchone()132        assert ind == (len(COUNTRIES), 2024, "worldbank")133    finally:134        con.close()135136137@pytest.mark.usefixtures("data_dir")138def test_failed_build_keeps_live_db_and_revisions_carry_forward(monkeypatch: pytest.MonkeyPatch) -> None:139    _tiny_staging()140    first = build_mod.build(run_id="20260101T000000Z", strict=False)141    before = settings.db_path.stat()142    # 1) a failure in the derived step must not touch the live DB143    def boom(con):144        raise RuntimeError("synthetic failure")145146    with pytest.MonkeyPatch.context() as mp:  # scoped: must not undo the data_dir fixture patch147        mp.setattr(build_mod.derived, "build_latest", boom)148        with pytest.raises(RuntimeError, match="synthetic"):149            build_mod.build(run_id="20260102T000000Z", strict=False)150    after = settings.db_path.stat()151    assert (after.st_ino, after.st_mtime_ns, after.st_size) == (before.st_ino, before.st_mtime_ns, before.st_size)152    assert not list(settings.build_dir.glob("atlas-20260102*"))153    # 2) an integrity failure in strict mode also keeps the live DB154    with pytest.MonkeyPatch.context() as mp:155        mp.setattr(build_mod, "HEADLINE_MIN_COUNTRIES", 10_000)156        with pytest.raises(build_mod.IntegrityError, match="only 24 countries"):157            build_mod.build(run_id="20260103T000000Z", strict=True)158    assert settings.db_path.stat().st_ino == before.st_ino159    # 3) a changed value is recorded in observation_revisions on the next successful build160    spec = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD")161    p = spec_paths(spec)["parquet"]162    import polars as pl163164    df = pl.read_parquet(p)165    df = df.with_columns(pl.when((pl.col("country_id") == "CAN") & (pl.col("year") == 2020)).then(pl.col("value") * 1.5)166                         .otherwise(pl.col("value")).alias("value"))167    write_parquet_atomic(df, p)168    second = build_mod.build(run_id="20260104T000000Z", strict=False)169    assert second.counts["revisions_new"] == 1170    con = duckdb.connect(str(settings.db_path), read_only=True)171    rev = con.execute("SELECT country_id, old_value, new_value, run_id FROM observation_revisions").fetchall()172    con.close()173    assert rev[0][0] == "CAN" and rev[0][2] == pytest.approx(rev[0][1] * 1.5) and rev[0][3] == "20260104T000000Z"174    snaps = sorted(settings.snapshots_dir.glob("atlas-*.duckdb"))175    assert [s.name for s in snaps] == [f"atlas-{first.run_id}.duckdb", f"atlas-{second.run_id}.duckdb"]176177178@pytest.mark.usefixtures("data_dir")179def test_export_helpers(tmp_path: Path) -> None:180    _tiny_staging()181    build_mod.build(run_id="20260101T000000Z", strict=False)182    from countryatlas.pipeline.export import export_country, export_indicator183184    p = export_indicator("gdp-per-capita", "json", tmp_path)185    doc = json.loads(p.read_bytes())186    # 22 WB series × 25 years + POL and CHL taken whole from IMF (27 rows each, incl. 2 forecasts)187    assert doc["meta"]["run_id"] == "20260101T000000Z" and len(doc["rows"]) == 22 * 25 + 2 * 27188    c = export_country("CAN", "csv", tmp_path)189    assert c.exists() and c.read_text().count("\n") > 50190