spb/countryatlas

1from __future__ import annotations23from datetime import UTC, datetime4from pathlib import Path56from countryatlas.connectors.owid import OWIDConnector7from countryatlas.models import IndicatorSourceSpec, RawPayload8910def _conn() -> OWIDConnector:11    c = OWIDConnector.__new__(OWIDConnector)12    c._cache, c._frames = {}, {}13    import threading1415    c._lock = threading.Lock()16    return c171819def _raw(body: bytes, dataset: str, code: str) -> RawPayload:20    return RawPayload(connector="owid", dataset=dataset, code=code, url="http://test", retrieved_at=datetime.now(UTC),21                      status_code=200, content_type="text/csv", body=body)222324def test_co2_subset_normalizes_per_column(fixtures: Path) -> None:25    body = (fixtures / "owid_co2_subset.csv").read_bytes()26    raw = _raw(body, "co2", "owid-co2-data")27    conn = _conn()28    rows = conn.normalize(raw, IndicatorSourceSpec(indicator_id="co2-per-capita", connector="owid", dataset="co2", code="co2_per_capita"))29    isos = {r.country_id for r in rows}30    assert isos == {"CAN", "FRA", "NGA", "KWT"}  # World (OWID_WRL) and Asia (no code) dropped31    assert all(r.unit == "tonnes per person" for r in rows)32    assert all(r.frequency == "A" and r.period.month == 1 for r in rows)33    # second spec on the same payload reuses the parsed frame34    rows2 = conn.normalize(raw, IndicatorSourceSpec(indicator_id="co2-emissions", connector="owid", dataset="co2", code="co2"))35    assert len(rows2) >= len(rows) - 536    assert len(conn._frames) == 1373839def test_grapher_projection_column_becomes_forecast(fixtures: Path) -> None:40    body = (fixtures / "owid_grapher_median-age.csv").read_bytes()41    conn = _conn()42    rows = conn.normalize(_raw(body, "grapher", "median-age"),43                          IndicatorSourceSpec(indicator_id="median-age", connector="owid", dataset="grapher", code="median-age"))44    can = {r.year: r for r in rows if r.country_id == "CAN"}45    assert can[2020].is_forecast is False46    assert any(r.is_forecast for r in can.values())47    assert max(can) <= datetime.now(UTC).year + 648    assert "WORLD" not in {r.country_id for r in rows} and "OWID_WRL" not in {r.country_id for r in rows}49