from __future__ import annotations import json from datetime import UTC, date, datetime from pathlib import Path import duckdb import pytest from countryatlas.config import settings from countryatlas.models import ImportRun, IndicatorSourceSpec, NormalizedObservation from countryatlas.pipeline import build as build_mod from countryatlas.pipeline.staging import rows_to_frame, spec_paths, write_parquet_atomic, write_run COUNTRIES = ["CAN", "USA", "FRA", "DEU", "JPN", "BRA", "IND", "NGA", "AUS", "MEX", "KOR", "ITA", "ESP", "GBR", "CHN", "ZAF", "EGY", "TUR", "ARG", "IDN", "SWE", "NOR", "CHL", "POL"] def _stage(spec: IndicatorSourceSpec, unit: str, base: float, growth: float, years: range, forecast_from: int | None = None, skip: tuple[str, ...] = (), stop_at: dict[str, int] | None = None) -> None: now = datetime.now(UTC) rows = [] for k, c in enumerate(COUNTRIES): if c in skip: continue for y in years: if stop_at and c in stop_at and y > stop_at[c]: continue v = base * (1 + 0.05 * k) * (growth ** (y - years.start)) rows.append(NormalizedObservation(country_id=c, indicator_id=spec.indicator_id, period=date(y, 1, 1), year=y, frequency="A", value=v, unit=unit, source_id=spec.connector, source_dataset=spec.dataset, source_series_code=spec.code, retrieved_at=now, is_forecast=bool(forecast_from and y >= forecast_from))) p = spec_paths(spec) write_parquet_atomic(rows_to_frame(rows), p["parquet"]) write_run(spec, ImportRun(run_id="T1", connector=spec.connector, dataset=spec.dataset, started_at=now, finished_at=now, status="ok", rows_norm=len(rows), rows_valid=len(rows))) def _tiny_staging() -> None: # WB: no data for POL (IMF takes the whole series); CHL stops in 2018 (IMF is > 3 years fresher → "fresher") wb = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD") _stage(wb, "current US$", 10_000, 1.03, range(2000, 2025), skip=("POL",), stop_at={"CHL": 2018}) # a lower-priority alternative source for the same indicator (complete series → observations_alt where WB wins) imf = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="imf", dataset="WEO", code="NGDPDPC", priority=2) _stage(imf, "current US$", 10_100, 1.03, range(2000, 2027), forecast_from=2025) pop = IndicatorSourceSpec(indicator_id="population", connector="worldbank", dataset="WDI", code="SP.POP.TOTL") _stage(pop, "people", 5_000_000, 1.01, range(2000, 2025)) le = IndicatorSourceSpec(indicator_id="life-expectancy", connector="worldbank", dataset="WDI", code="SP.DYN.LE00.IN") _stage(le, "years", 70, 1.002, range(2000, 2024)) ma = IndicatorSourceSpec(indicator_id="median-age", connector="owid", dataset="grapher", code="median-age") _stage(ma, "years", 25, 1.005, range(2000, 2024)) fr = IndicatorSourceSpec(indicator_id="fertility-rate", connector="worldbank", dataset="WDI", code="SP.DYN.TFRT.IN") _stage(fr, "births per woman", 3.0, 0.99, range(2000, 2024)) # annual inflation for everyone + a MONTHLY FRED series for the USA (must stay out of latest/rankings) infl = IndicatorSourceSpec(indicator_id="inflation", connector="worldbank", dataset="WDI", code="FP.CPI.TOTL.ZG") _stage(infl, "annual %", 2.0, 1.01, range(2000, 2026)) # CAN alone already has 2026 → its rank must fall back to 2025 (the last year with ≥ 20 ranked countries) import polars as pl p = spec_paths(infl)["parquet"] df = pl.read_parquet(p) extra = df.filter((pl.col("country_id") == "CAN") & (pl.col("year") == 2025)).with_columns( pl.lit(date(2026, 1, 1)).alias("period"), pl.lit(2026, dtype=pl.Int32).alias("year"), (pl.col("value") + 1.0).alias("value")) write_parquet_atomic(pl.concat([df, extra]), p) fred = IndicatorSourceSpec(indicator_id="inflation", connector="fred", dataset="FRED", code="CPIAUCSL", priority=2, countries=["USA"], frequency="M") now = datetime.now(UTC) monthly = [NormalizedObservation(country_id="USA", indicator_id="inflation", period=date(2026, m, 1), year=2026, frequency="M", value=3.0 + m / 10, unit="annual %", source_id="fred", source_dataset="FRED", source_series_code="CPIAUCSL", retrieved_at=now) for m in range(1, 8)] write_parquet_atomic(rows_to_frame(monthly), spec_paths(fred)["parquet"]) write_run(fred, ImportRun(run_id="T1", connector="fred", dataset="FRED", started_at=now, finished_at=now, status="ok", rows_norm=len(monthly), rows_valid=len(monthly))) @pytest.mark.usefixtures("data_dir") def test_build_tiny_staging_produces_all_tables() -> None: _tiny_staging() r = build_mod.build(run_id="20260101T000000Z", strict=False) assert settings.db_path.exists() assert r.snapshot_path is not None and r.snapshot_path.exists() con = duckdb.connect(str(settings.db_path), read_only=True) try: tables = {t[0] for t in con.execute("SELECT table_name FROM information_schema.tables").fetchall()} for t in ("countries", "groups", "group_members", "sources", "indicators", "indicator_sources", "observations", "observations_alt", "observation_revisions", "latest", "rankings", "changes", "events", "similarity", "insights", "country_dna", "coverage", "import_runs", "validation_issues", "search_index", "meta"): assert t in tables, t n_obs = con.execute("SELECT count(*) FROM observations").fetchone()[0] n_alt = con.execute("SELECT count(*) FROM observations_alt").fetchone()[0] assert n_obs > 0 and n_alt > 0 # IMF series lose the priority race where WB has data # one source per (country, indicator, frequency) series — never spliced assert con.execute("SELECT count(*) FROM (SELECT country_id, indicator_id, frequency FROM observations " "GROUP BY ALL HAVING count(DISTINCT source_id) > 1)").fetchone()[0] == 0 assert con.execute("SELECT count(*) FROM observations_alt WHERE country_id='CAN' AND source_id='imf'").fetchone()[0] == 27 src = dict(con.execute("SELECT country_id, source_id FROM observations WHERE indicator_id='gdp-per-capita' " "GROUP BY ALL").fetchall()) assert src["CAN"] == "worldbank" and src["POL"] == "imf" and src["CHL"] == "imf" reasons = dict(con.execute("SELECT country_id, json_extract_string(metadata, '$.merge_reason') FROM observations " "WHERE indicator_id='gdp-per-capita' GROUP BY ALL").fetchall()) assert reasons["CAN"] == "priority" and reasons["POL"] == "priority" and reasons["CHL"] == "fresher" # forecasts (from the chosen source) kept in observations but excluded from latest/rankings assert con.execute("SELECT count(*) FROM observations WHERE is_forecast").fetchone()[0] == 2 * 2 assert con.execute("SELECT count(*) FROM latest WHERE is_forecast").fetchone()[0] == 0 assert con.execute("SELECT max(year) FROM latest WHERE indicator_id='gdp-per-capita'").fetchone()[0] == 2024 lat = con.execute("SELECT rank_world, n_world, change_10y_pct FROM latest WHERE country_id='CAN' AND indicator_id='gdp-per-capita'").fetchone() assert lat[0] is not None and lat[1] == len(COUNTRIES) and lat[2] is not None assert con.execute("SELECT count(*) FROM rankings WHERE indicator_id='gdp-per-capita' AND year=2024").fetchone()[0] == len(COUNTRIES) # rank 1 = highest value (higher_is_better true for gdp-per-capita) top = con.execute("SELECT country_id FROM rankings WHERE indicator_id='gdp-per-capita' AND year=2024 AND rank=1").fetchone()[0] assert top == COUNTRIES[-1] assert con.execute("SELECT count(*) FROM coverage").fetchone()[0] == len(con.execute("SELECT * FROM countries").fetchall()) assert con.execute("SELECT count(*) FROM search_index WHERE type='country'").fetchone()[0] >= 200 assert con.execute("SELECT count(*) FROM insights").fetchone()[0] > 0 assert con.execute("SELECT count(*) FROM country_dna").fetchone()[0] > 0 assert con.execute("SELECT count(*) FROM import_runs").fetchone()[0] == 8 assert r.counts["series_fresher_source"] == 1 # canonical frequency only: the USA monthly FRED series is in observations but not in latest/rankings assert con.execute("SELECT count(*) FROM observations WHERE frequency='M'").fetchone()[0] == 7 usa = con.execute("SELECT year, frequency, source_id, n_world FROM latest WHERE country_id='USA' AND indicator_id='inflation'").fetchone() assert usa == (2025, "A", "worldbank", len(COUNTRIES)) assert con.execute("SELECT count(*) FROM latest l JOIN indicators i ON i.id=l.indicator_id WHERE l.frequency<>i.frequency").fetchone()[0] == 0 assert con.execute("SELECT count(*) FROM rankings WHERE indicator_id='inflation' AND year=2026").fetchone()[0] == 0 # rank fallback: CAN's latest inflation is 2026 (only country) → ranked within 2025 among all 24 can = con.execute("SELECT year, rank_year, n_world FROM latest WHERE country_id='CAN' AND indicator_id='inflation'").fetchone() assert can[0] == 2026 and can[1] == 2025 and can[2] == len(COUNTRIES) meta = dict(con.execute("SELECT key, value FROM meta").fetchall()) assert meta["schema_version"] == "1" and meta["build_run_id"] == "20260101T000000Z" assert int(meta["observation_count"]) == n_obs ind = con.execute("SELECT n_countries, last_year, primary_source_id FROM indicators WHERE id='gdp-per-capita'").fetchone() assert ind == (len(COUNTRIES), 2024, "worldbank") finally: con.close() @pytest.mark.usefixtures("data_dir") def test_failed_build_keeps_live_db_and_revisions_carry_forward(monkeypatch: pytest.MonkeyPatch) -> None: _tiny_staging() first = build_mod.build(run_id="20260101T000000Z", strict=False) before = settings.db_path.stat() # 1) a failure in the derived step must not touch the live DB def boom(con): raise RuntimeError("synthetic failure") with pytest.MonkeyPatch.context() as mp: # scoped: must not undo the data_dir fixture patch mp.setattr(build_mod.derived, "build_latest", boom) with pytest.raises(RuntimeError, match="synthetic"): build_mod.build(run_id="20260102T000000Z", strict=False) after = settings.db_path.stat() assert (after.st_ino, after.st_mtime_ns, after.st_size) == (before.st_ino, before.st_mtime_ns, before.st_size) assert not list(settings.build_dir.glob("atlas-20260102*")) # 2) an integrity failure in strict mode also keeps the live DB with pytest.MonkeyPatch.context() as mp: mp.setattr(build_mod, "HEADLINE_MIN_COUNTRIES", 10_000) with pytest.raises(build_mod.IntegrityError, match="only 24 countries"): build_mod.build(run_id="20260103T000000Z", strict=True) assert settings.db_path.stat().st_ino == before.st_ino # 3) a changed value is recorded in observation_revisions on the next successful build spec = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD") p = spec_paths(spec)["parquet"] import polars as pl df = pl.read_parquet(p) df = df.with_columns(pl.when((pl.col("country_id") == "CAN") & (pl.col("year") == 2020)).then(pl.col("value") * 1.5) .otherwise(pl.col("value")).alias("value")) write_parquet_atomic(df, p) second = build_mod.build(run_id="20260104T000000Z", strict=False) assert second.counts["revisions_new"] == 1 con = duckdb.connect(str(settings.db_path), read_only=True) rev = con.execute("SELECT country_id, old_value, new_value, run_id FROM observation_revisions").fetchall() con.close() assert rev[0][0] == "CAN" and rev[0][2] == pytest.approx(rev[0][1] * 1.5) and rev[0][3] == "20260104T000000Z" snaps = sorted(settings.snapshots_dir.glob("atlas-*.duckdb")) assert [s.name for s in snaps] == [f"atlas-{first.run_id}.duckdb", f"atlas-{second.run_id}.duckdb"] @pytest.mark.usefixtures("data_dir") def test_export_helpers(tmp_path: Path) -> None: _tiny_staging() build_mod.build(run_id="20260101T000000Z", strict=False) from countryatlas.pipeline.export import export_country, export_indicator p = export_indicator("gdp-per-capita", "json", tmp_path) doc = json.loads(p.read_bytes()) # 22 WB series × 25 years + POL and CHL taken whole from IMF (27 rows each, incl. 2 forecasts) assert doc["meta"]["run_id"] == "20260101T000000Z" and len(doc["rows"]) == 22 * 25 + 2 * 27 c = export_country("CAN", "csv", tmp_path) assert c.exists() and c.read_text().count("\n") > 50