spb/countryatlas

1from __future__ import annotations23import json4from datetime import UTC, datetime5from pathlib import Path67import pytest89from countryatlas.connectors._util import ConnectorError10from countryatlas.connectors.worldbank import WorldBankConnector, _check_header11from countryatlas.models import IndicatorSourceSpec, RawPayload121314def _raw(body: bytes, code: str = "NY.GDP.PCAP.CD") -> RawPayload:15    return RawPayload(connector="worldbank", dataset="WDI", code=code, url="http://test", retrieved_at=datetime.now(UTC),16                      status_code=200, content_type="application/json", body=body,17                      source_updated_at=datetime(2026, 7, 13, tzinfo=UTC), meta={"lastupdated": "2026-07-13", "page": 1})181920def test_normalize_fixture_drops_aggregates_and_nulls(fixtures: Path) -> None:21    body = (fixtures / "wb_NY.GDP.PCAP.CD.json").read_bytes()22    doc = json.loads(body)23    n_country_rows = sum(1 for r in doc[1] if r["countryiso3code"] in {"CAN", "USA", "FRA", "DEU", "JPN", "NGA", "BRA", "IND"}24                         and r["value"] is not None)25    spec = IndicatorSourceSpec(indicator_id="gdp-per-capita", connector="worldbank", dataset="WDI", code="NY.GDP.PCAP.CD")26    conn = WorldBankConnector.__new__(WorldBankConnector)  # no HTTP client needed for normalize27    rows = conn.normalize([_raw(body)], spec)28    assert len(rows) == n_country_rows29    assert {r.country_id for r in rows} <= {"CAN", "USA", "FRA", "DEU", "JPN", "NGA", "BRA", "IND"}30    can = sorted((r for r in rows if r.country_id == "CAN"), key=lambda r: r.period)31    assert can[-1].year >= 2023 and can[-1].period.month == 1 and can[-1].frequency == "A"32    assert can[-1].unit == "current US$"33    assert can[-1].source_updated_at == datetime(2026, 7, 13, tzinfo=UTC)34    assert not any(r.is_forecast for r in rows)35    report = WorldBankConnector.validate(conn, rows)36    assert report.errors == 0 and not report.quarantine_dataset373839def test_transform_applied() -> None:40    body = json.dumps([{"page": 1, "pages": 1, "total": 1, "lastupdated": "2026-07-13"},41                       [{"countryiso3code": "CAN", "country": {"id": "CA"}, "date": "2020", "value": 2.0, "obs_status": "E"}]]).encode()42    spec = IndicatorSourceSpec(indicator_id="gdp", connector="worldbank", dataset="WDI", code="X", transform="x*1e9")43    conn = WorldBankConnector.__new__(WorldBankConnector)44    rows = conn.normalize(_raw(body, "X"), spec)45    assert rows[0].value == 2e9 and rows[0].is_estimate is True464748def test_retired_code_is_an_error(fixtures: Path) -> None:49    doc = json.loads((fixtures / "wb_retired.json").read_text())50    with pytest.raises(ConnectorError, match="retired or invalid"):51        _check_header(doc, "XX.OLD.CODE")525354def test_period_parsing() -> None:55    from countryatlas.connectors._util import parse_period5657    assert parse_period("2023")[1:] == (2023, "A")58    assert parse_period("2023Q3")[0].month == 759    assert parse_period("2023M11")[0].month == 1160    assert parse_period("garbage") is None61