ProStock/tests/test_daily.py

"""Test for daily market data API.

Tests the daily interface implementation against api.md requirements:
- A股日线行情所有输出字段
- tor 换手率
- vr 量比
"""

import pytest
import pandas as pd
from src.data.api_wrappers import get_daily


# Expected output fields according to api.md
EXPECTED_BASE_FIELDS = [
    "ts_code",  # 股票代码
    "trade_date",  # 交易日期
    "open",  # 开盘价
    "high",  # 最高价
    "low",  # 最低价
    "close",  # 收盘价
    "pre_close",  # 昨收价
    "change",  # 涨跌额
    "pct_chg",  # 涨跌幅
    "vol",  # 成交量
    "amount",  # 成交额
]

EXPECTED_FACTOR_FIELDS = [
    "turnover_rate",  # 换手率 (tor)
    "volume_ratio",  # 量比 (vr)
]


class TestGetDaily:
    """Test cases for get_daily function with real API calls."""

    def test_fetch_basic(self):
        """Test basic daily data fetch with real API."""
        result = get_daily("000001.SZ", start_date="20240101", end_date="20240131")

        assert isinstance(result, pd.DataFrame)
        assert len(result) >= 1
        assert result["ts_code"].iloc[0] == "000001.SZ"

    def test_fetch_with_factors(self):
        """Test fetch with tor and vr factors."""
        result = get_daily(
            "000001.SZ",
            start_date="20240101",
            end_date="20240131",
            factors=["tor", "vr"],
        )

        assert isinstance(result, pd.DataFrame)
        # Check all base fields are present
        for field in EXPECTED_BASE_FIELDS:
            assert field in result.columns, f"Missing base field: {field}"
        # Check factor fields are present
        for field in EXPECTED_FACTOR_FIELDS:
            assert field in result.columns, f"Missing factor field: {field}"

    def test_output_fields_completeness(self):
        """Verify all required output fields are returned."""
        result = get_daily("600000.SH")

        # Verify all base fields are present
        assert set(EXPECTED_BASE_FIELDS).issubset(result.columns.tolist()), (
            f"Missing fields: {set(EXPECTED_BASE_FIELDS) - set(result.columns)}"
        )

    def test_empty_result(self):
        """Test handling of empty results."""
        # 使用真实 API 测试无效股票代码的空结果
        result = get_daily("INVALID.SZ")
        assert isinstance(result, pd.DataFrame)
        assert result.empty

    def test_date_range_query(self):
        """Test query with date range."""
        result = get_daily(
            "000001.SZ",
            start_date="20240101",
            end_date="20240131",
        )

        assert isinstance(result, pd.DataFrame)
        assert len(result) >= 1

    def test_with_adj(self):
        """Test fetch with adjustment type."""
        result = get_daily("000001.SZ", adj="qfq")

        assert isinstance(result, pd.DataFrame)


def test_integration():
    """Integration test with real Tushare API (requires valid token)."""
    import os

    token = os.environ.get("TUSHARE_TOKEN")
    if not token:
        pytest.skip("TUSHARE_TOKEN not configured")

    result = get_daily(
        "000001.SZ", start_date="20240101", end_date="20240131", factors=["tor", "vr"]
    )

    # Verify structure
    assert isinstance(result, pd.DataFrame)
    if not result.empty:
        # Check base fields
        for field in EXPECTED_BASE_FIELDS:
            assert field in result.columns, f"Missing base field: {field}"
        # Check factor fields
        for field in EXPECTED_FACTOR_FIELDS:
            assert field in result.columns, f"Missing factor field: {field}"


if __name__ == "__main__":
    # 运行 pytest 单元测试（真实API调用）
    pytest.main([__file__, "-v"])