OpenBB/openbb_platform/providers/jodi/tests/test_models.py

"""Test the JODI model validators, views, and error branches."""

import asyncio
from datetime import date

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

from openbb_jodi.models.gas_balance import JodiGasBalanceFetcher
from openbb_jodi.models.gas_demand import JodiGasDemandFetcher, JodiGasDemandQueryParams
from openbb_jodi.models.gas_imports import JodiGasImportsFetcher
from openbb_jodi.models.oil_balance import (
    JodiOilBalanceFetcher,
    JodiOilBalanceQueryParams,
)
from openbb_jodi.models.oil_demand_by_product import JodiOilDemandByProductFetcher
from openbb_jodi.models.oil_production import (
    JodiOilProductionFetcher,
    JodiOilProductionQueryParams,
)

DATES = {"start_date": date(2024, 1, 1), "end_date": date(2024, 3, 31)}


@pytest.mark.parametrize(
    "value, expected",
    [
        (None, "united_states"),
        ("", "united_states"),
        ("SA", "saudi_arabia"),
        ("Saudi Arabia", "saudi_arabia"),
    ],
)
def test_single_country_validator(value, expected):
    assert JodiOilBalanceQueryParams(country=value).country == expected


def test_single_country_validator_rejects_multiple():
    with pytest.raises(OpenBBError, match="Only one country"):
        JodiOilBalanceQueryParams(country="united_states,saudi_arabia")


@pytest.mark.parametrize(
    "value, expected",
    [
        (None, None),
        ("", None),
        ("US", "united_states"),
        (["US", "Saudi Arabia"], "united_states,saudi_arabia"),
        ("united_states,us", "united_states"),
    ],
)
def test_multi_country_validator(value, expected):
    assert JodiOilProductionQueryParams(country=value).country == expected


def test_country_validator_invalid():
    with pytest.raises(OpenBBError, match="Invalid country: 'atlantis'"):
        JodiOilProductionQueryParams(country="atlantis")


def test_token_normalization():
    params = JodiOilBalanceQueryParams(product="Crude Oil", unit="KBD")
    assert params.product == "crude_oil"
    assert params.unit == "kbd"
    params = JodiGasDemandQueryParams(measure=None, unit=None)
    assert params.measure == "observed"
    assert params.unit == "m3"


def test_token_invalid():
    with pytest.raises(ValidationError):
        JodiOilBalanceQueryParams(unit="furlongs")
    with pytest.raises(ValidationError):
        JodiOilProductionQueryParams(product="gasoline")  # secondary-only product


def _fetch(fetcher, params: dict):
    return asyncio.run(fetcher.fetch_data(params, {}))


def test_start_after_end():
    with pytest.raises(OpenBBError, match="start_date must be"):
        _fetch(
            JodiOilBalanceFetcher,
            {"start_date": date(2024, 2, 1), "end_date": date(2024, 1, 1)},
        )


def test_oil_dates_before_data():
    with pytest.raises(OpenBBError, match="The data starts in 2002"):
        _fetch(
            JodiOilProductionFetcher,
            {"start_date": date(2000, 1, 1), "end_date": date(2001, 12, 31)},
        )


def test_gas_dates_before_data():
    with pytest.raises(OpenBBError, match="The data starts in 2009"):
        _fetch(
            JodiGasBalanceFetcher,
            {"start_date": date(2007, 1, 1), "end_date": date(2008, 12, 31)},
        )


@pytest.mark.usefixtures("mock_download")
def test_empty_result():
    with pytest.raises(EmptyDataError):
        _fetch(
            JodiOilProductionFetcher,
            {"start_date": date(2024, 4, 1), "end_date": date(2024, 12, 31)},
        )


@pytest.mark.usefixtures("mock_download")
def test_oil_balance_view():
    result = _fetch(
        JodiOilBalanceFetcher,
        {"country": "united_states", "product": "crude_oil", "unit": "kbbl", **DATES},
    )
    assert len(result.result) == 3
    row = result.result[0]
    assert row.production is not None
    assert row.closing_stocks is not None
    assert row.refinery_output is None  # secondary-only flow
    assert result.metadata["country"] == "United States"
    assert result.metadata["product"] == "Crude oil"
    assert result.metadata["unit"] == "Thousand barrels"
    assert result.metadata["assessments"]["production"] is not None


@pytest.mark.usefixtures("mock_download")
def test_oil_balance_secondary_product():
    result = _fetch(
        JodiOilBalanceFetcher,
        {"country": "united_states", "product": "gasoline", "unit": "kbd", **DATES},
    )
    row = result.result[0]
    assert row.refinery_output is not None
    assert row.demand is not None
    assert row.refinery_intake is None  # primary-only flow


@pytest.mark.usefixtures("mock_download")
def test_oil_production_comparison_view():
    result = _fetch(JodiOilProductionFetcher, {**DATES})
    row = result.result[0].model_dump(exclude_none=True)
    # Columns are countries, alphabetical, after the date.
    # Russia is in the sample but reports no kb/d production values.
    assert list(row) == ["date", "saudi_arabia", "united_states"]
    assert result.metadata["flow"] == "Production"
    assert result.metadata["assessments"]["saudi_arabia"] is not None


@pytest.mark.usefixtures("mock_download")
def test_gas_empty_result():
    with pytest.raises(EmptyDataError):
        _fetch(
            JodiGasBalanceFetcher,
            {
                "country": "norway",
                "start_date": date(2024, 4, 1),
                "end_date": date(2024, 12, 31),
            },
        )


@pytest.mark.usefixtures("mock_download")
def test_oil_demand_by_product_view():
    result = _fetch(
        JodiOilDemandByProductFetcher, {"country": "united_states", **DATES}
    )
    row = result.result[0]
    assert row.gasoline is not None
    assert row.total_products is not None
    assert result.metadata["flow"] == "Demand"


@pytest.mark.usefixtures("mock_download")
def test_gas_balance_view():
    result = _fetch(JodiGasBalanceFetcher, {"country": "norway", **DATES})
    row = result.result[0]
    assert row.production is not None
    assert row.pipeline_exports is not None
    assert result.metadata["product"] == "Natural gas"


@pytest.mark.usefixtures("mock_download")
def test_gas_demand_measures_differ():
    observed = _fetch(
        JodiGasDemandFetcher,
        {"country": "united_states", "measure": "observed", **DATES},
    )
    calculated = _fetch(
        JodiGasDemandFetcher,
        {"country": "united_states", "measure": "calculated", **DATES},
    )
    assert observed.metadata["flow"] == "Gross inland deliveries (observed)"
    assert calculated.metadata["flow"] == "Gross inland deliveries (calculated)"
    obs = observed.result[0].model_dump(exclude_none=True)["united_states"]
    calc = calculated.result[0].model_dump(exclude_none=True)["united_states"]
    assert obs != calc


@pytest.mark.usefixtures("mock_download")
def test_gas_imports_flow_selection():
    total = _fetch(JodiGasImportsFetcher, {"country": "united_states", **DATES})
    pipeline = _fetch(
        JodiGasImportsFetcher,
        {"country": "united_states", "flow": "pipeline", **DATES},
    )
    assert total.metadata["flow"] == "Imports"
    assert pipeline.metadata["flow"] == "Pipeline imports"
    total_value = total.result[0].model_dump(exclude_none=True)["united_states"]
    pipeline_value = pipeline.result[0].model_dump(exclude_none=True)["united_states"]
    assert pipeline_value <= total_value