OpenBB/openbb_platform/providers/imf/tests/test_dot_helpers.py

"""Tests for openbb_imf.utils.dot_helpers."""

# ruff: noqa: I001

from unittest.mock import MagicMock, patch

import pytest


COUNTRY_FIXTURE = [
    {"value": "USA", "label": "United States"},
    {"value": "GBR", "label": "United Kingdom"},
    {"value": "DEU", "label": "Germany"},
    {"value": "G001", "label": "World"},
    {"value": "G163", "label": "Euro Area"},
    {"value": "GX170", "label": "Europe"},
    {"value": "G7", "label": "Group of Seven"},
    {"value": "TX1", "label": "Total Exports"},
    {"value": "TX2", "label": "Other TX"},
]


@pytest.fixture
def _patch_metadata(monkeypatch):
    """Patch ``ImfMetadata().get_dataflow_parameters`` for IMTS country data."""

    class FakeMeta:
        def get_dataflow_parameters(self, _dataflow):
            return {"COUNTRY": list(COUNTRY_FIXTURE)}

    monkeypatch.setattr("openbb_imf.utils.metadata.ImfMetadata", FakeMeta)


class TestLoadCountryChoices:
    """Tests for load_country_choices."""

    def test_groups_split_and_sorted(self, _patch_metadata):
        """Plain countries, G-regions, and TX groups appear in that order."""
        from openbb_imf.utils.dot_helpers import load_country_choices

        result = load_country_choices()
        values = [item["value"] for item in result]

        # plain country block first
        plain_values = {"USA", "GBR", "DEU"}
        plain_block_idx = [i for i, v in enumerate(values) if v in plain_values]
        # G regions next
        g_block_idx = [
            i for i, v in enumerate(values) if v in {"G001", "G163", "GX170", "G7"}
        ]
        # TX groups last
        tx_block_idx = [i for i, v in enumerate(values) if v in {"TX1", "TX2"}]

        assert max(plain_block_idx) < min(g_block_idx)
        assert max(g_block_idx) < min(tx_block_idx)


class TestListCountryChoices:
    """Tests for list_country_choices."""

    def test_returns_codes(self, _patch_metadata):
        """Returns just the value codes."""
        from openbb_imf.utils.dot_helpers import list_country_choices

        result = list_country_choices()
        assert "USA" in result
        assert "G001" in result
        assert all(isinstance(c, str) for c in result)


class TestLabelToCodeMap:
    """Tests for get_label_to_code_map."""

    def test_normalized_labels(self, _patch_metadata):
        """Normalized label maps to its code."""
        from openbb_imf.utils.dot_helpers import get_label_to_code_map

        result = get_label_to_code_map()
        assert result["united_states"] == "USA"
        assert result["united_kingdom"] == "GBR"
        assert result["world"] == "G001"


class TestCodeToLabelMap:
    """Tests for get_code_to_label_map."""

    def test_codes_map_to_labels(self, _patch_metadata):
        """Codes map to normalized labels."""
        from openbb_imf.utils.dot_helpers import get_code_to_label_map

        result = get_code_to_label_map()
        assert result["USA"] == "united_states"
        assert result["G001"] == "world"


class TestResolveCountryInput:
    """Tests for resolve_country_input."""

    def test_empty_raises(self):
        """Empty input raises ValueError."""
        from openbb_imf.utils.dot_helpers import resolve_country_input

        with pytest.raises(ValueError, match="cannot be empty"):
            resolve_country_input("")

    def test_all_alias(self):
        """'all' / '*' return '*'."""
        from openbb_imf.utils.dot_helpers import resolve_country_input

        assert resolve_country_input("all") == "*"
        assert resolve_country_input("*") == "*"

    def test_common_aliases(self):
        """Common aliases resolve to fixed codes."""
        from openbb_imf.utils.dot_helpers import resolve_country_input

        assert resolve_country_input("world") == "G001"
        assert resolve_country_input("eurozone") == "G163"
        assert resolve_country_input("euro_area") == "G163"
        assert resolve_country_input("Euro Area") == "G163"
        assert resolve_country_input("EU") == "G998"
        assert resolve_country_input("european_union") == "G998"
        assert resolve_country_input("europe") == "GX170"

    def test_iso_code(self, _patch_metadata):
        """Valid ISO-like code passes through."""
        from openbb_imf.utils.dot_helpers import resolve_country_input

        assert resolve_country_input("USA") == "USA"
        assert resolve_country_input("usa") == "USA"

    def test_resolve_by_label(self, _patch_metadata):
        """Country label resolves to code."""
        from openbb_imf.utils.dot_helpers import resolve_country_input

        assert resolve_country_input("united_states") == "USA"
        assert resolve_country_input("United States") == "USA"

    def test_unknown_raises(self, _patch_metadata):
        """Unknown country raises ValueError with hint."""
        from openbb_imf.utils.dot_helpers import resolve_country_input

        with pytest.raises(ValueError, match="not a valid IMF country code"):
            resolve_country_input("Atlantis")


class TestImtsQuery:
    """Tests for imts_query."""

    def _make_qb(self, **overrides):
        """Build a mocked ImfQueryBuilder context manager."""
        qb = MagicMock()
        qb.metadata.get_dataflow_parameters.return_value = {
            "COUNTRY": [
                {"value": "USA"},
                {"value": "G001"},
                {"value": "GBR"},
            ],
            "COUNTERPART_COUNTRY": [
                {"value": "G001"},
                {"value": "USA"},
            ],
        }
        qb.fetch_data.return_value = {"data": [], "metadata": {}}
        for k, v in overrides.items():
            setattr(qb, k, v)
        return qb

    def test_empty_country_raises(self):
        """Empty country/counterpart raises ValueError."""
        from openbb_imf.utils.dot_helpers import imts_query

        with pytest.raises(ValueError, match="cannot be empty"):
            imts_query(country="", counterpart="G001", indicator="*")

    def test_invalid_frequency_raises(self):
        """Frequency not in A/Q/M raises."""
        from openbb_imf.utils.dot_helpers import imts_query

        with pytest.raises(ValueError, match="Frequency must be one of"):
            imts_query(country="USA", counterpart="G001", indicator="*", freq="X")

    def test_simple_query_dispatches(self):
        """Valid inputs dispatch to query_builder.fetch_data."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            out = dot_helpers.imts_query(
                country="USA",
                counterpart="G001",
                indicator="X",
                freq="annual",
            )
        assert out == {"data": [], "metadata": {}}
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["dataflow"] == "IMTS"
        assert kwargs["FREQUENCY"] == "A"
        assert kwargs["COUNTRY"] == "USA"
        assert kwargs["COUNTERPART_COUNTRY"] == "G001"
        assert kwargs["INDICATOR"] == "X"

    def test_comma_separated_country(self):
        """Comma-separated country string becomes a list."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(
                country="USA,GBR",
                counterpart="G001",
                indicator="*",
                freq="A",
            )
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["COUNTRY"] == ["USA", "GBR"]

    def test_wildcard_counterpart(self):
        """'*' counterpart passes through."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(
                country="USA",
                counterpart="*",
                indicator="*",
                freq="A",
            )
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["COUNTERPART_COUNTRY"] == "*"

    def test_invalid_country_raises(self):
        """Invalid country code raises ValueError listing offending codes."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            with pytest.raises(ValueError, match="Invalid country"):
                dot_helpers.imts_query(
                    country="ZZZ",
                    counterpart="G001",
                    indicator="*",
                )

    def test_invalid_counterpart_raises(self):
        """Invalid counterpart code raises ValueError."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            with pytest.raises(ValueError, match="Invalid counterpart"):
                dot_helpers.imts_query(
                    country="USA",
                    counterpart="ZZZ",
                    indicator="*",
                )

    def test_wildcard_in_country_list(self):
        """A '*' inside a list collapses to wildcard."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(
                country=["USA", "*"],
                counterpart=["G001"],
                indicator="*",
            )
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["COUNTRY"] == "*"

    def test_no_valid_values_passthrough(self):
        """When params lack the dimension, selection passes through unchanged."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        qb.metadata.get_dataflow_parameters.return_value = {}
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(country="ZZZ", counterpart="QQQ", indicator="*")
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["COUNTRY"] == "ZZZ"
        assert kwargs["COUNTERPART_COUNTRY"] == "QQQ"

    def test_indicator_comma_list(self):
        """Comma-separated indicator splits to list."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(
                country="USA",
                counterpart="G001",
                indicator="X, M",
            )
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["INDICATOR"] == ["X", "M"]

    def test_single_country_list(self):
        """Single-item list with one country returns just the string."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(
                country=["USA"],
                counterpart=["G001"],
                indicator="*",
            )
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["COUNTRY"] == "USA"
        assert kwargs["COUNTERPART_COUNTRY"] == "G001"

    def test_counterpart_country_fallback(self):
        """Missing COUNTERPART_COUNTRY falls back to COUNTRY param values."""
        from openbb_imf.utils import dot_helpers

        qb = self._make_qb()
        qb.metadata.get_dataflow_parameters.return_value = {
            "COUNTRY": [{"value": "USA"}, {"value": "G001"}],
        }
        with patch("openbb_imf.utils.query_builder.ImfQueryBuilder", return_value=qb):
            dot_helpers.imts_query(country="USA", counterpart="G001", indicator="*")
        kwargs = qb.fetch_data.call_args.kwargs
        assert kwargs["COUNTERPART_COUNTRY"] == "G001"