galaxis-po/backend/tests/unit/test_collector_resilience.py

"""
Unit tests for ETFCollector and ValuationCollector resilience.

Tests verify that collectors handle KRX API failures gracefully:
- JSONDecodeError, ConnectionError, KeyError from pykrx
- Retry logic and fallback behavior
- Existing DB data is never deleted on failure
- Clear error messages when KRX login is required
"""
from datetime import date
from json import JSONDecodeError
from unittest.mock import patch, MagicMock, call

import pandas as pd
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool

from app.core.database import Base
from app.models.stock import ETF, Valuation
from app.services.collectors.etf_collector import ETFCollector, KRXDataError
from app.services.collectors.valuation_collector import ValuationCollector


@pytest.fixture
def db():
    """In-memory SQLite database for testing."""
    engine = create_engine(
        "sqlite:///:memory:",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(bind=engine)
    Session = sessionmaker(bind=engine)
    session = Session()
    yield session
    session.close()
    Base.metadata.drop_all(bind=engine)


# ── ETFCollector Tests (pykrx fallback path, no KRX_OPENAPI_KEY) ──


class TestETFCollectorResilience:

    @patch("app.services.collectors.etf_collector.get_krx_client", return_value=None)
    @patch("app.services.collectors.etf_collector.time.sleep")
    @patch("pykrx.website.krx.etx.core.ETF_전종목기본종목")
    def test_json_decode_error_retries_once(self, mock_etf_cls, mock_sleep, mock_client, db):
        """JSONDecodeError on both attempts raises KRXDataError with login hint."""
        mock_fetcher = MagicMock()
        mock_fetcher.fetch.side_effect = [
            JSONDecodeError("msg", "doc", 0),
            JSONDecodeError("msg", "doc", 0),
        ]
        mock_etf_cls.return_value = mock_fetcher

        collector = ETFCollector(db)
        with pytest.raises(KRXDataError, match="KRX_ID/KRX_PW"):
            collector.collect()

        assert mock_fetcher.fetch.call_count == 2
        mock_sleep.assert_called_once_with(3)

    @patch("app.services.collectors.etf_collector.get_krx_client", return_value=None)
    @patch("app.services.collectors.etf_collector.time.sleep")
    @patch("pykrx.website.krx.etx.core.ETF_전종목기본종목")
    def test_connection_error_retries_and_raises(self, mock_etf_cls, mock_sleep, mock_client, db):
        """ConnectionError on both attempts raises KRXDataError."""
        mock_fetcher = MagicMock()
        mock_fetcher.fetch.side_effect = ConnectionError("timeout")
        mock_etf_cls.return_value = mock_fetcher

        collector = ETFCollector(db)
        with pytest.raises(KRXDataError):
            collector.collect()

        assert mock_fetcher.fetch.call_count == 2

    @patch("app.services.collectors.etf_collector.get_krx_client", return_value=None)
    @patch("app.services.collectors.etf_collector.time.sleep")
    @patch("pykrx.website.krx.etx.core.ETF_전종목기본종목")
    def test_retry_succeeds_on_second_attempt(self, mock_etf_cls, mock_sleep, mock_client, db):
        """If first attempt fails but retry succeeds, data is processed normally."""
        good_df = pd.DataFrame([{
            "ISU_SRT_CD": "069500",
            "ISU_ABBRV": "KODEX 200",
            "IDX_ASST_CLSS_NM": "주식",
            "IDX_MKT_CLSS_NM": "코스피",
            "ETF_TOT_FEE": "0.15",
        }])
        mock_fetcher = MagicMock()
        mock_fetcher.fetch.side_effect = [
            JSONDecodeError("msg", "doc", 0),
            good_df,
        ]
        mock_etf_cls.return_value = mock_fetcher

        collector = ETFCollector(db)
        result = collector.collect()

        assert result == 1

    @patch("app.services.collectors.etf_collector.get_krx_client", return_value=None)
    @patch("app.services.collectors.etf_collector.time.sleep")
    @patch("pykrx.website.krx.etx.core.ETF_전종목기본종목")
    def test_failure_does_not_delete_existing_data(self, mock_etf_cls, mock_sleep, mock_client, db):
        """Existing ETF records are preserved when fetch fails."""
        db.execute(
            ETF.__table__.insert().values(
                ticker="069500", name="KODEX 200",
                asset_class="equity", market="코스피",
            )
        )
        db.commit()

        mock_fetcher = MagicMock()
        mock_fetcher.fetch.side_effect = JSONDecodeError("msg", "doc", 0)
        mock_etf_cls.return_value = mock_fetcher

        collector = ETFCollector(db)
        with pytest.raises(KRXDataError):
            collector.collect()

        existing = db.query(ETF).filter_by(ticker="069500").first()
        assert existing is not None
        assert existing.name == "KODEX 200"


# ── ETFCollector Open API Tests ──


class TestETFCollectorOpenAPI:

    @patch("app.services.collectors.etf_collector.get_krx_client")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260417")
    def test_openapi_success(self, mock_biz, mock_get_client, db):
        """Open API returns valid ETF data."""
        mock_client = MagicMock()
        mock_client.get_etf_daily.return_value = pd.DataFrame([{
            "ISU_SRT_CD": "069500",
            "ISU_ABBRV": "KODEX 200",
            "TDD_CLSPRC": 35000.0,
            "ACC_TRDVOL": 1000000,
        }])
        mock_get_client.return_value = mock_client

        collector = ETFCollector(db)
        result = collector.collect()

        assert result == 1
        etf = db.query(ETF).filter_by(ticker="069500").first()
        assert etf is not None
        assert etf.name == "KODEX 200"

    @patch("app.services.collectors.etf_collector.get_krx_client")
    @patch("app.services.collectors.etf_collector.time.sleep")
    @patch("pykrx.website.krx.etx.core.ETF_전종목기본종목")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260417")
    def test_openapi_failure_falls_back_to_pykrx(self, mock_biz, mock_etf_cls, mock_sleep, mock_get_client, db):
        """When Open API fails, falls back to pykrx scraping."""
        mock_client = MagicMock()
        mock_client.get_etf_daily.side_effect = Exception("API down")
        mock_get_client.return_value = mock_client

        good_df = pd.DataFrame([{
            "ISU_SRT_CD": "069500",
            "ISU_ABBRV": "KODEX 200",
            "IDX_ASST_CLSS_NM": "주식",
            "IDX_MKT_CLSS_NM": "코스피",
            "ETF_TOT_FEE": "0.15",
        }])
        mock_fetcher = MagicMock()
        mock_fetcher.fetch.return_value = good_df
        mock_etf_cls.return_value = mock_fetcher

        collector = ETFCollector(db)
        result = collector.collect()

        assert result == 1


# ── ValuationCollector Tests ──


class TestValuationCollectorResilience:

    @patch("app.services.collectors.valuation_collector.pykrx_stock")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260327")
    def test_key_error_falls_back_to_previous_days(self, mock_biz, mock_pykrx, db):
        """KeyError triggers fallback to previous business days."""
        good_df = pd.DataFrame(
            {"BPS": [50000], "PER": [10.0], "PBR": [1.5], "EPS": [5000], "DIV": [2.0], "DPS": [1000]},
            index=["005930"],
        )
        mock_pykrx.get_market_fundamental_by_ticker.side_effect = [
            KeyError("BPS"),       # day 0
            KeyError("PER"),       # day -1
            good_df,               # day -2 succeeds
        ]

        collector = ValuationCollector(db, biz_day="20260327")
        result = collector.collect()

        assert result == 1
        assert mock_pykrx.get_market_fundamental_by_ticker.call_count == 3

    @patch("app.services.collectors.valuation_collector.pykrx_stock")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260327")
    def test_all_days_fail_raises_with_login_hint(self, mock_biz, mock_pykrx, db):
        """When all 4 date attempts fail with KeyError, raises RuntimeError with login hint."""
        mock_pykrx.get_market_fundamental_by_ticker.side_effect = KeyError("BPS")

        collector = ValuationCollector(db, biz_day="20260327")
        with pytest.raises(RuntimeError, match="KRX_ID"):
            collector.collect()

        assert mock_pykrx.get_market_fundamental_by_ticker.call_count == 4

    @patch("app.services.collectors.valuation_collector.pykrx_stock")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260327")
    def test_json_decode_error_raises_with_login_hint(self, mock_biz, mock_pykrx, db):
        """JSONDecodeError on all attempts raises RuntimeError with login hint."""
        mock_pykrx.get_market_fundamental_by_ticker.side_effect = JSONDecodeError("msg", "doc", 0)

        collector = ValuationCollector(db, biz_day="20260327")
        with pytest.raises(RuntimeError, match="KRX requires login"):
            collector.collect()

    @patch("app.services.collectors.valuation_collector.pykrx_stock")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260327")
    def test_empty_df_triggers_fallback(self, mock_biz, mock_pykrx, db):
        """Empty DataFrame (not exception) also triggers fallback."""
        good_df = pd.DataFrame(
            {"BPS": [60000], "PER": [15.0], "PBR": [2.0], "EPS": [4000], "DIV": [1.5], "DPS": [800]},
            index=["005930"],
        )
        mock_pykrx.get_market_fundamental_by_ticker.side_effect = [
            pd.DataFrame(),   # empty
            good_df,          # day -1 has data
        ]

        collector = ValuationCollector(db, biz_day="20260327")
        result = collector.collect()

        assert result == 1

    @patch("app.services.collectors.valuation_collector.pykrx_stock")
    @patch("app.services.collectors.base.BaseCollector._get_latest_biz_day", return_value="20260327")
    def test_failure_preserves_existing_data(self, mock_biz, mock_pykrx, db):
        """Existing valuation records are preserved when fetch fails."""
        db.execute(
            Valuation.__table__.insert().values(
                ticker="005930", base_date=date(2026, 3, 26),
                per=10.0, pbr=1.5, dividend_yield=2.0,
            )
        )
        db.commit()

        mock_pykrx.get_market_fundamental_by_ticker.side_effect = KeyError("BPS")

        collector = ValuationCollector(db, biz_day="20260327")
        with pytest.raises(RuntimeError):
            collector.collect()

        existing = db.query(Valuation).filter_by(ticker="005930").first()
        assert existing is not None
        assert existing.per == 10.0