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

"""Tests for the EIA Data Browser model."""

from datetime import date

import pytest
from openbb_core.app.model.abstract.error import OpenBBError

from openbb_us_eia.models.data_browser import (
    EiaDataBrowserFetcher,
    EiaDataBrowserQueryParams,
)

METADATA = {
    "response": {
        "id": "spt",
        "frequency": [
            {"id": "daily", "format": "YYYY-MM-DD"},
            {"id": "monthly", "format": "YYYY-MM"},
        ],
        "defaultFrequency": "daily",
        "data": {"value": {"units": "$/BBL"}},
        "facets": [{"id": "duoarea"}, {"id": "product"}],
    }
}

PAGE = {
    "response": {
        "total": "2",
        "dateFormat": "YYYY-MM",
        "data": [
            {"period": "2024-02", "product": "EPCBRENT", "value": "83.5"},
            {"period": "2024-01", "product": "EPCBRENT", "value": "80.1"},
        ],
    }
}


def _stub(monkeypatch, metadata=METADATA, page=PAGE):
    from openbb_core.provider.utils import helpers as core_helpers

    calls = {"urls": []}

    async def fake_amake_request(url, response_callback=None, **kwargs):
        calls["urls"].append(url)
        return metadata if "/data/" not in url else page

    monkeypatch.setattr(core_helpers, "amake_request", fake_amake_request)
    return calls


class TestDataBrowserQueryParams:
    """Route cleaning and facet string parsing."""

    def test_route_cleaned(self):
        query = EiaDataBrowserQueryParams(route=" /petroleum/pri/spt/data/ ")
        assert query.route == "petroleum/pri/spt"

    def test_facet_query_parses_pairs(self):
        query = EiaDataBrowserQueryParams(
            route="x", facets="duoarea:NUS; product:EPCBRENT,EPCWTI"
        )
        assert query.facet_query == {
            "duoarea": ["NUS"],
            "product": ["EPCBRENT", "EPCWTI"],
        }

    def test_facet_query_empty(self):
        assert EiaDataBrowserQueryParams(route="x").facet_query == {}

    def test_facet_query_invalid_pair_raises(self):
        query = EiaDataBrowserQueryParams(route="x", facets="duoarea=NUS")
        with pytest.raises(OpenBBError, match="Invalid facet filter"):
            _ = query.facet_query

    def test_transform_query_validates_facets(self):
        with pytest.raises(OpenBBError, match="Invalid facet filter"):
            EiaDataBrowserFetcher.transform_query({"route": "x", "facets": "bad"})


class TestDataBrowserFetcher:
    """Live-metadata resolution and the request pipeline."""

    @pytest.mark.asyncio
    async def test_fetch_data_end_to_end(self, monkeypatch):
        calls = _stub(monkeypatch)
        rows = await EiaDataBrowserFetcher.fetch_data(
            {
                "route": "petroleum/pri/spt",
                "frequency": "monthly",
                "facets": "product:EPCBRENT",
                "start_date": date(2024, 1, 1),
                "end_date": date(2024, 2, 29),
            },
            {"eia_api_key": "MOCK_KEY"},
        )
        assert len(rows) == 2
        assert rows[0].date == date(2024, 1, 1)
        assert rows[0].model_dump()["value"] == 80.1
        data_url = calls["urls"][-1]
        assert "start=2024-01" in data_url
        assert "end=2024-02" in data_url
        assert "facets%5Bproduct%5D%5B%5D=EPCBRENT" in data_url

    @pytest.mark.asyncio
    async def test_default_frequency_used(self, monkeypatch):
        calls = _stub(monkeypatch)
        await EiaDataBrowserFetcher.fetch_data(
            {"route": "petroleum/pri/spt"}, {"eia_api_key": "MOCK_KEY"}
        )
        assert "frequency=daily" in calls["urls"][-1]

    @pytest.mark.asyncio
    async def test_category_route_raises(self, monkeypatch):
        metadata = {"response": {"routes": [{"id": "spt"}, {"id": "gnd"}]}}
        _stub(monkeypatch, metadata=metadata)
        with pytest.raises(OpenBBError, match="not a data route"):
            await EiaDataBrowserFetcher.fetch_data(
                {"route": "petroleum/pri"}, {"eia_api_key": "MOCK_KEY"}
            )

    @pytest.mark.asyncio
    async def test_unknown_route_raises(self, monkeypatch):
        _stub(monkeypatch, metadata={"response": {}})
        with pytest.raises(OpenBBError, match="not a valid EIA data route"):
            await EiaDataBrowserFetcher.fetch_data(
                {"route": "nonsense"}, {"eia_api_key": "MOCK_KEY"}
            )

    @pytest.mark.asyncio
    async def test_invalid_frequency_raises(self, monkeypatch):
        _stub(monkeypatch)
        with pytest.raises(OpenBBError, match="Frequency 'hourly'"):
            await EiaDataBrowserFetcher.fetch_data(
                {"route": "petroleum/pri/spt", "frequency": "hourly"},
                {"eia_api_key": "MOCK_KEY"},
            )

    @pytest.mark.asyncio
    async def test_data_type_subset_requested(self, monkeypatch):
        calls = _stub(monkeypatch)
        await EiaDataBrowserFetcher.fetch_data(
            {"route": "petroleum/pri/spt", "data_type": "value"},
            {"eia_api_key": "MOCK_KEY"},
        )
        assert "data%5B0%5D=value" in calls["urls"][-1]

    @pytest.mark.asyncio
    async def test_invalid_data_type_raises(self, monkeypatch):
        _stub(monkeypatch)
        with pytest.raises(OpenBBError, match="Invalid data_type"):
            await EiaDataBrowserFetcher.fetch_data(
                {"route": "petroleum/pri/spt", "data_type": "quantity"},
                {"eia_api_key": "MOCK_KEY"},
            )

    @pytest.mark.asyncio
    async def test_unknown_facet_raises(self, monkeypatch):
        _stub(monkeypatch)
        with pytest.raises(OpenBBError, match="Invalid facet id"):
            await EiaDataBrowserFetcher.fetch_data(
                {"route": "petroleum/pri/spt", "facets": "stateId:TX"},
                {"eia_api_key": "MOCK_KEY"},
            )

    @pytest.mark.asyncio
    async def test_desc_sort_and_limit(self, monkeypatch):
        _stub(monkeypatch)
        rows = await EiaDataBrowserFetcher.fetch_data(
            {"route": "petroleum/pri/spt", "sort": "desc", "limit": 1},
            {"eia_api_key": "MOCK_KEY"},
        )
        assert len(rows) == 1
        assert rows[0].date == date(2024, 2, 1)