OpenBB/openbb_platform/providers/eia/tests/test_api_query.py

"""Tests for the shared EIA APIv2 query, pagination, and transform machinery."""

from datetime import date
from urllib.parse import parse_qs, urlparse

import pytest
from openbb_core.app.model.abstract.error import OpenBBError
from openbb_core.provider.utils.errors import EmptyDataError

from openbb_us_eia.models.crude_oil_imports import EiaCrudeOilImportsQueryParams
from openbb_us_eia.models.ieo import EiaIeoQueryParams
from openbb_us_eia.models.registry import DATASET_FETCHERS
from openbb_us_eia.utils import catalog
from openbb_us_eia.utils.api_query import (
    MAX_AUTO_ROWS,
    empty_result_message,
    maybe_float,
    paginate,
    to_float,
    transform_rows,
)
from openbb_us_eia.utils.helpers import build_data_url

CoalReceiptsQueryParams = DATASET_FETCHERS["EiaCoalReceipts"].query_params_type


class TestQueryParamProperties:
    """Dataset resolution and facet collection on the base query params."""

    def test_dataset_key_from_dataset_class_attribute(self):
        query = CoalReceiptsQueryParams()
        assert query.dataset_key == "receipts"

    def test_dataset_key_from_release_field(self):
        query = EiaIeoQueryParams(release="2021")
        assert query.dataset_key == "2021"

    def test_dataset_key_single_dataset_group(self):
        query = EiaCrudeOilImportsQueryParams()
        assert query.dataset_key == "crude_oil_imports"

    def test_dataset_key_falls_back_to_first_of_group(self):
        from openbb_us_eia.utils.api_query import EiaApiQueryParams

        class BareSedsQuery(EiaApiQueryParams):
            """Synthetic query with neither __dataset__ nor a release field."""

            __group__ = "seds"

        assert BareSedsQuery().dataset_key == "seds"

    def test_facet_values_exclude_standard_fields(self):
        query = CoalReceiptsQueryParams(coal_rank="MET", limit=10, sort="desc")
        facets = query.facet_values
        assert facets["coal_rank"] == "MET"
        assert "limit" not in facets
        assert "sort" not in facets
        assert "dataset" not in facets


class TestBuildDataUrl:
    """URL rendering for one page of an EIA data request."""

    def test_full_request(self):
        request = {
            "path": "coal/receipts",
            "api_key": "KEY",
            "frequency": "quarterly",
            "data_columns": ["quantity", "price"],
            "facets": {"coalRankId": ["MET", "STM"]},
            "start": "2020-Q1",
            "end": "2021-Q4",
            "sort": "asc",
        }
        url = build_data_url(request, offset=5000, length=2500)
        parsed = urlparse(url)
        assert parsed.path == "/v2/coal/receipts/data/"
        params = parse_qs(parsed.query)
        assert params["api_key"] == ["KEY"]
        assert params["frequency"] == ["quarterly"]
        assert params["data[0]"] == ["quantity"]
        assert params["data[1]"] == ["price"]
        assert params["facets[coalRankId][]"] == ["MET", "STM"]
        assert params["start"] == ["2020-Q1"]
        assert params["end"] == ["2021-Q4"]
        assert params["sort[0][column]"] == ["period"]
        assert params["sort[0][direction]"] == ["asc"]
        assert params["offset"] == ["5000"]
        assert params["length"] == ["2500"]

    def test_optional_parts_omitted(self):
        request = {"path": "steo", "api_key": "KEY"}
        params = parse_qs(urlparse(build_data_url(request, 0, 5000)).query)
        assert "frequency" not in params
        assert "start" not in params
        assert "end" not in params
        assert params["sort[0][direction]"] == ["desc"]


def _page(rows, total, date_format="YYYY"):
    return {"response": {"data": rows, "total": str(total), "dateFormat": date_format}}


def _stub_http(pages):
    """Build amake_request/amake_requests stubs serving canned pages by offset."""
    calls = {"urls": []}

    async def amake_request(url, response_callback=None, **kwargs):
        calls["urls"].append(url)
        offset = int(parse_qs(urlparse(url).query)["offset"][0])
        return pages[offset]

    async def amake_requests(urls, response_callback=None, **kwargs):
        calls["urls"].extend(urls)
        return [pages[int(parse_qs(urlparse(url).query)["offset"][0])] for url in urls]

    return amake_request, amake_requests, calls


async def _noop_callback(response, session):
    return {}


class TestPaginate:
    """Automatic pagination against the response total."""

    @pytest.mark.asyncio
    async def test_single_page(self):
        rows = [{"period": "2020", "value": "1"}]
        amake_request, amake_requests, calls = _stub_http({0: _page(rows, 1)})
        result = await paginate(
            build_data_url,
            {"path": "seds", "api_key": "K"},
            limit=None,
            amake_request=amake_request,
            amake_requests=amake_requests,
            response_callback=_noop_callback,
        )
        assert result["rows"] == rows
        assert result["total"] == 1
        assert result["date_format"] == "YYYY"
        assert len(calls["urls"]) == 1

    @pytest.mark.asyncio
    async def test_multiple_pages(self):
        pages = {
            0: _page([{"period": str(y)} for y in range(5000)], 12000),
            5000: _page([{"period": str(y)} for y in range(5000)], 12000),
            10000: _page([{"period": str(y)} for y in range(2000)], 12000),
        }
        amake_request, amake_requests, calls = _stub_http(pages)
        result = await paginate(
            build_data_url,
            {"path": "seds", "api_key": "K"},
            limit=None,
            amake_request=amake_request,
            amake_requests=amake_requests,
            response_callback=_noop_callback,
        )
        assert len(result["rows"]) == 12000
        assert len(calls["urls"]) == 3

    @pytest.mark.asyncio
    async def test_limit_caps_requests_and_rows(self):
        pages = {
            0: _page([{"period": str(y)} for y in range(5000)], 50000),
            5000: _page([{"period": str(y)} for y in range(2000)], 50000),
        }
        amake_request, amake_requests, calls = _stub_http(pages)
        result = await paginate(
            build_data_url,
            {"path": "seds", "api_key": "K"},
            limit=7000,
            amake_request=amake_request,
            amake_requests=amake_requests,
            response_callback=_noop_callback,
        )
        assert len(result["rows"]) == 7000
        assert len(calls["urls"]) == 2
        last = parse_qs(urlparse(calls["urls"][-1]).query)
        assert last["length"] == ["2000"]

    @pytest.mark.asyncio
    async def test_empty_response_raises(self):
        amake_request, amake_requests, _ = _stub_http({0: _page([], 0)})
        with pytest.raises(EmptyDataError):
            await paginate(
                build_data_url,
                {"path": "seds", "api_key": "K"},
                limit=None,
                amake_request=amake_request,
                amake_requests=amake_requests,
                response_callback=_noop_callback,
            )

    @pytest.mark.asyncio
    async def test_auto_ceiling_raises_without_limit(self):
        pages = {0: _page([{"period": "2020"}] * 5000, MAX_AUTO_ROWS + 1)}
        amake_request, amake_requests, _ = _stub_http(pages)
        with pytest.raises(OpenBBError, match="pagination ceiling"):
            await paginate(
                build_data_url,
                {"path": "seds", "api_key": "K"},
                limit=None,
                amake_request=amake_request,
                amake_requests=amake_requests,
                response_callback=_noop_callback,
            )

    @pytest.mark.asyncio
    async def test_explicit_limit_bypasses_ceiling(self):
        pages = {0: _page([{"period": "2020"}] * 100, MAX_AUTO_ROWS + 1)}
        amake_request, amake_requests, _ = _stub_http(pages)
        result = await paginate(
            build_data_url,
            {"path": "seds", "api_key": "K"},
            limit=100,
            amake_request=amake_request,
            amake_requests=amake_requests,
            response_callback=_noop_callback,
        )
        assert len(result["rows"]) == 100


class TestEmptyResultMessage:
    """Empty results are explained against the live facet vocabularies."""

    @staticmethod
    def _facet_server(vocabularies):
        async def amake_request(url, response_callback=None, **kwargs):
            facet_id = url.split("/facet/")[1].split("?")[0]
            if vocabularies.get(facet_id) is None:
                raise ValueError("facet endpoint unavailable")
            return {
                "response": {
                    "facets": [
                        {"id": value, "name": label}
                        for value, label in vocabularies[facet_id].items()
                    ]
                }
            }

        return amake_request

    @pytest.mark.asyncio
    async def test_invalid_value_lists_valid_choices(self):
        request = {
            "path": "petroleum/pri/spt",
            "api_key": "K",
            "facets": {"series": ["RWTC"], "process": ["VAM"]},
            "facet_params": {"series": "series", "process": "process"},
        }
        amake_request = self._facet_server(
            {
                "series": {"RWTC": "Cushing, OK WTI Spot Price FOB"},
                "process": {"PF4": "Spot Price FOB"},
            }
        )
        message = await empty_result_message(request, amake_request, _noop_callback)
        assert "process: VAM is not valid for this dataset" in message
        assert "PF4 = Spot Price FOB" in message

    @pytest.mark.asyncio
    async def test_valid_combination_explains_no_intersection(self):
        request = {
            "path": "petroleum/pri/spt",
            "api_key": "K",
            "facets": {"duoarea": ["NUS"], "series": ["RWTC"]},
            "facet_params": {"duoarea": "region", "series": "series"},
            "coverage": ("1986-01-02", "2026-06-30"),
            "start": "2026-01-01",
        }
        amake_request = self._facet_server(
            {
                "duoarea": {"NUS": "U.S."},
                "series": {"RWTC": "Cushing, OK WTI Spot Price FOB"},
            }
        )
        message = await empty_result_message(request, amake_request, _noop_callback)
        assert "combination does not match" in message
        assert "region: NUS (U.S.)" in message
        assert "series: RWTC (Cushing, OK WTI Spot Price FOB)" in message
        assert "The dataset covers 1986-01-02 to 2026-06-30" in message

    @pytest.mark.asyncio
    async def test_facet_endpoint_failure_falls_back_to_codes(self):
        request = {
            "path": "petroleum/pri/spt",
            "api_key": "K",
            "facets": {"series": ["RWTC"]},
            "facet_params": {"series": "series"},
        }
        amake_request = self._facet_server({"series": None})
        message = await empty_result_message(request, amake_request, _noop_callback)
        assert "series: RWTC" in message

    @pytest.mark.asyncio
    async def test_no_facets_reports_date_range(self):
        request = {
            "path": "seds",
            "api_key": "K",
            "facets": {},
            "coverage": ("1960", "2023"),
            "start": "2030",
            "end": "2031",
        }
        message = await empty_result_message(request, None, _noop_callback)
        assert message.startswith("The request returned no data.")
        assert "Requested period: 2030 to 2031" in message
        assert "The dataset covers 1960 to 2023" in message


class TestTransformRows:
    """Row normalization for catalog datasets."""

    def test_facet_and_data_column_renames(self):
        spec = catalog.get_dataset("coal", "exports_imports_quantity_price")
        rows = [
            {
                "period": "2023-Q1",
                "exportImportType": "Exports",
                "coalRankId": "MET",
                "coalRankDescription": "Metallurgical",
                "countryId": "CA",
                "countryDescription": "Canada",
                "quantity": "100.5",
                "price": "not-a-number",
                "quantity-units": "short tons",
            }
        ]
        out = transform_rows(spec, rows, 'YYYY-"Q"Q')
        record = out[0]
        assert record["date"] == date(2023, 1, 1)
        assert record["export_import_type"] == "Exports"
        assert record["coal_rank"] == "MET"
        assert record["coal_rank_name"] == "Metallurgical"
        assert record["country_name"] == "Canada"
        assert record["quantity"] == 100.5
        assert record["price"] is None
        assert "quantity_units" not in record

    def test_unknown_columns_snake_cased(self):
        spec = catalog.get_dataset("seds", "seds")
        rows = [
            {
                "period": "2020",
                "seriesId": "TETPB",
                "seriesDescription": "Total energy",
                "stateId": "TX",
                "stateDescription": "Texas",
                "value": "285.1",
                "unit": "million Btu",
                "surpriseColumn": "x",
            }
        ]
        record = transform_rows(spec, rows, "YYYY")[0]
        assert record["series"] == "TETPB"
        assert record["series_name"] == "Total energy"
        assert record["state"] == "TX"
        assert record["state_name"] == "Texas"
        assert record["value"] == 285.1
        assert record["surprise_column"] == "x"


class TestNumericCasting:
    """Numeric columns are floats; withheld flags become null."""

    def test_to_float_casts_numeric_strings(self):
        assert to_float("1.5") == 1.5

    def test_to_float_nulls_withheld_flags(self):
        assert to_float("W") is None
        assert to_float("NA") is None

    def test_to_float_passes_numbers_and_none(self):
        assert to_float(None) is None
        assert to_float(2) == 2

    def test_maybe_float_for_raw_browser_columns(self):
        assert maybe_float("1.5") == 1.5
        assert maybe_float("W") == "W"
        assert maybe_float(None) is None