fx_quant_engine/tests/test_router_circuit_breaker.py

from datetime import datetime, timedelta, timezone

from fx_quant_engine.ingestion.base import BaseDataAdapter
from fx_quant_engine.ingestion.router import DataSourceRouter


class AlwaysFailingAdapter(BaseDataAdapter):
    def fetch_price_data(self, asset, start, end):
        raise RuntimeError("boom")

    def fetch_macro_data(self, key, start, end):
        raise RuntimeError("boom")

    def fetch_rate_data(self, asset, start, end):
        raise RuntimeError("boom")

    def health_check(self):
        return False


class AlwaysWorkingAdapter(BaseDataAdapter):
    def fetch_price_data(self, asset, start, end):
        import pandas as pd

        idx = pd.date_range(start=start, end=end, freq="B")
        return pd.DataFrame({"close": 1.0}, index=idx)

    def fetch_macro_data(self, key, start, end):
        import pandas as pd

        idx = pd.date_range(start=start, end=end, freq="B")
        return pd.DataFrame({key: 0.0}, index=idx)

    def fetch_rate_data(self, asset, start, end):
        import pandas as pd

        idx = pd.date_range(start=start, end=end, freq="B")
        return pd.DataFrame({f"{asset}_rate": 0.05}, index=idx)

    def health_check(self):
        return True


def test_circuit_breaker_skips_failing_source() -> None:
    cfg = {
        "enabled_sources": {"bad": True, "good": True},
        "source_priority": {"default": ["bad", "good"]},
        "asset_source_map": {"USDINR": ["bad", "good"]},
        "reliability_tags": {"bad": "low", "good": "high"},
        "latency_tags_ms": {"bad": 1, "good": 1},
        "circuit_breaker": {"failure_threshold": 1, "cooldown_seconds": 3600},
    }
    router = DataSourceRouter({"bad": AlwaysFailingAdapter(), "good": AlwaysWorkingAdapter()}, cfg)
    end = datetime.now(timezone.utc)
    start = end - timedelta(days=5)

    out1 = router.fetch_price_data("USDINR", start, end)
    assert out1.success is True
    assert out1.source == "good"

    out2 = router.fetch_price_data("USDINR", start, end)
    assert out2.success is True
    assert out2.source == "good"