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.daily 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'])