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) -> None:20    now = datetime.now(UTC)21    rows = []22    for k, c in enumerate(COUNTRIES):23        for y in years:24            v = base * (1 + 0.05 * k) * (growth ** (y - years.start))25            rows.append(NormalizedObservation(country_id=c, indicator_id=spec.indicator_id, period=date(y, 1, 1), year=y,26                                              frequency="A", value=v, unit=unit, source_id=spec.connector,27                                              source_dataset=spec.dataset, source_series_code=spec.code, retrieved_at=now,28                                              is_forecast=bool(forecast_from and y >= forecast_from)))29    p = spec_paths(spec)30    write_parquet_atomic(rows_to_frame(rows), p["parquet"])31    write_run(spec, ImportRun(run_id="T1", connector=spec.connector, dataset=spec.dataset, started_at=now, finished_at=now,32                              status="ok", rows_norm=len(rows), rows_valid=len(rows)))333435def _tiny_staging() -> None:36    wb = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD")37    _stage(wb, "current US$", 10_000, 1.03, range(2000, 2025))38    # a lower-priority alternative source for the same indicator → observations_alt39    imf = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="imf", dataset="WEO", code="NGDPDPC", priority=2)40    _stage(imf, "current US$", 10_100, 1.03, range(2000, 2027), forecast_from=2025)41    pop = IndicatorSourceSpec(indicator_id="population", connector="worldbank", dataset="WDI", code="SP.POP.TOTL")42    _stage(pop, "people", 5_000_000, 1.01, range(2000, 2025))43    le = IndicatorSourceSpec(indicator_id="life-expectancy", connector="worldbank", dataset="WDI", code="SP.DYN.LE00.IN")44    _stage(le, "years", 70, 1.002, range(2000, 2024))45    ma = IndicatorSourceSpec(indicator_id="median-age", connector="owid", dataset="grapher", code="median-age")46    _stage(ma, "years", 25, 1.005, range(2000, 2024))47    fr = IndicatorSourceSpec(indicator_id="fertility-rate", connector="worldbank", dataset="WDI", code="SP.DYN.TFRT.IN")48    _stage(fr, "births per woman", 3.0, 0.99, range(2000, 2024))495051@pytest.mark.usefixtures("data_dir")52def test_build_tiny_staging_produces_all_tables() -> None:53    _tiny_staging()54    r = build_mod.build(run_id="20260101T000000Z", strict=False)55    assert settings.db_path.exists()56    assert r.snapshot_path is not None and r.snapshot_path.exists()57    con = duckdb.connect(str(settings.db_path), read_only=True)58    try:59        tables = {t[0] for t in con.execute("SELECT table_name FROM information_schema.tables").fetchall()}60        for t in ("countries", "groups", "group_members", "sources", "indicators", "indicator_sources", "observations",61                  "observations_alt", "observation_revisions", "latest", "rankings", "changes", "events", "similarity",62                  "insights", "country_dna", "coverage", "import_runs", "validation_issues", "search_index", "meta"):63            assert t in tables, t64        n_obs = con.execute("SELECT count(*) FROM observations").fetchone()[0]65        n_alt = con.execute("SELECT count(*) FROM observations_alt").fetchone()[0]66        assert n_obs > 0 and n_alt > 0  # IMF rows lose the priority race where WB has a value67        # forecasts kept in observations but excluded from latest/rankings68        assert con.execute("SELECT count(*) FROM observations WHERE is_forecast").fetchone()[0] > 069        assert con.execute("SELECT count(*) FROM latest WHERE is_forecast").fetchone()[0] == 070        assert con.execute("SELECT max(year) FROM latest WHERE indicator_id='gdp-per-capita'").fetchone()[0] == 202471        lat = con.execute("SELECT rank_world, n_world, change_10y_pct FROM latest WHERE country_id='CAN' AND indicator_id='gdp-per-capita'").fetchone()72        assert lat[0] is not None and lat[1] == len(COUNTRIES) and lat[2] is not None73        assert con.execute("SELECT count(*) FROM rankings WHERE indicator_id='gdp-per-capita' AND year=2024").fetchone()[0] == len(COUNTRIES)74        # rank 1 = highest value (higher_is_better true for gdp-per-capita)75        top = con.execute("SELECT country_id FROM rankings WHERE indicator_id='gdp-per-capita' AND year=2024 AND rank=1").fetchone()[0]76        assert top == COUNTRIES[-1]77        assert con.execute("SELECT count(*) FROM coverage").fetchone()[0] == len(con.execute("SELECT * FROM countries").fetchall())78        assert con.execute("SELECT count(*) FROM search_index WHERE type='country'").fetchone()[0] >= 20079        assert con.execute("SELECT count(*) FROM insights").fetchone()[0] > 080        assert con.execute("SELECT count(*) FROM country_dna").fetchone()[0] > 081        assert con.execute("SELECT count(*) FROM import_runs").fetchone()[0] == 682        meta = dict(con.execute("SELECT key, value FROM meta").fetchall())83        assert meta["schema_version"] == "1" and meta["build_run_id"] == "20260101T000000Z"84        assert int(meta["observation_count"]) == n_obs85        ind = con.execute("SELECT n_countries, last_year, primary_source_id FROM indicators WHERE id='gdp-per-capita'").fetchone()86        assert ind == (len(COUNTRIES), 2024, "worldbank")87    finally:88        con.close()899091@pytest.mark.usefixtures("data_dir")92def test_failed_build_keeps_live_db_and_revisions_carry_forward(monkeypatch: pytest.MonkeyPatch) -> None:93    _tiny_staging()94    first = build_mod.build(run_id="20260101T000000Z", strict=False)95    before = settings.db_path.stat()96    # 1) a failure in the derived step must not touch the live DB97    def boom(con):98        raise RuntimeError("synthetic failure")99100    monkeypatch.setattr(build_mod.derived, "build_latest", boom)101    with pytest.raises(RuntimeError, match="synthetic"):102        build_mod.build(run_id="20260102T000000Z", strict=False)103    after = settings.db_path.stat()104    assert (after.st_ino, after.st_mtime_ns, after.st_size) == (before.st_ino, before.st_mtime_ns, before.st_size)105    assert not list(settings.build_dir.glob("atlas-20260102*"))106    monkeypatch.undo()107    # 2) an integrity failure in strict mode also keeps the live DB108    monkeypatch.setattr(build_mod, "HEADLINE_MIN_COUNTRIES", 10_000)109    with pytest.raises(build_mod.IntegrityError):110        build_mod.build(run_id="20260103T000000Z", strict=True)111    assert settings.db_path.stat().st_ino == before.st_ino112    monkeypatch.undo()113    # 3) a changed value is recorded in observation_revisions on the next successful build114    spec = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD")115    p = spec_paths(spec)["parquet"]116    import polars as pl117118    df = pl.read_parquet(p)119    df = df.with_columns(pl.when((pl.col("country_id") == "CAN") & (pl.col("year") == 2020)).then(pl.col("value") * 1.5)120                         .otherwise(pl.col("value")).alias("value"))121    write_parquet_atomic(df, p)122    second = build_mod.build(run_id="20260104T000000Z", strict=False)123    assert second.counts["revisions_new"] == 1124    con = duckdb.connect(str(settings.db_path), read_only=True)125    rev = con.execute("SELECT country_id, old_value, new_value, run_id FROM observation_revisions").fetchall()126    con.close()127    assert rev[0][0] == "CAN" and rev[0][2] == pytest.approx(rev[0][1] * 1.5) and rev[0][3] == "20260104T000000Z"128    snaps = sorted(settings.snapshots_dir.glob("atlas-*.duckdb"))129    assert [s.name for s in snaps] == [f"atlas-{first.run_id}.duckdb", f"atlas-{second.run_id}.duckdb"]130131132@pytest.mark.usefixtures("data_dir")133def test_export_helpers(tmp_path: Path) -> None:134    _tiny_staging()135    build_mod.build(run_id="20260101T000000Z", strict=False)136    from countryatlas.pipeline.export import export_country, export_indicator137138    p = export_indicator("gdp-per-capita", "json", tmp_path)139    doc = json.loads(p.read_bytes())140    assert doc["meta"]["run_id"] == "20260101T000000Z" and len(doc["rows"]) == len(COUNTRIES) * 27141    c = export_country("CAN", "csv", tmp_path)142    assert c.exists() and c.read_text().count("\n") > 50143