python_market_data_service/scripts/test_extended_fields.py

"""测试扩展字段功能

测试内容包括:
1. 获取日K线数据（包含新字段）
2. 验证新字段数据是否正确

使用方法:
    python scripts/test_extended_fields.py
"""
import os
import sys
import asyncio

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from datetime import datetime, timedelta
from app.adapters.akshare_adapter import AKShareAdapter
from app.core.logger import info, error


async def test_fetch_daily_klines():
    """测试获取日K线数据（包含扩展字段）"""
    adapter = AKShareAdapter()
    await adapter.connect({})
    
    # 测试股票: 平安银行 000001.SZ
    symbol = "000001.SZ"
    
    # 获取最近5个交易日的数据
    end_date = datetime.now()
    start_date = end_date - timedelta(days=30)  # 获取最近30天的数据
    
    start_str = start_date.strftime("%Y%m%d")
    end_str = end_date.strftime("%Y%m%d")
    
    info(f"\n测试获取 {symbol} 的日K线数据 ({start_str} - {end_str})")
    
    try:
        klines = await adapter.fetch_klines(
            symbol=symbol,
            start=start_str,
            end=end_str,
            freq="1d"
        )
        
        if not klines:
            error(f"未获取到 {symbol} 的K线数据")
            return False
        
        info(f"成功获取 {len(klines)} 条K线数据\n")
        
        # 打印第一条数据的详细信息
        first_kline = klines[0]
        info("第一条K线数据:")
        info(f"  日期: {datetime.fromtimestamp(first_kline.time).strftime('%Y-%m-%d')}")
        info(f"  开盘价: {first_kline.open}")
        info(f"  最高价: {first_kline.high}")
        info(f"  最低价: {first_kline.low}")
        info(f"  收盘价: {first_kline.close}")
        info(f"  成交量: {first_kline.volume}")
        info(f"  成交额: {first_kline.amount}")
        info("")
        info("扩展字段:")
        info(f"  交易日: {first_kline.trade_date}")
        info(f"  是否涨停: {first_kline.is_limit_up}")
        info(f"  是否跌停: {first_kline.is_limit_down}")
        info(f"  总市值: {first_kline.total_market_cap:,.0f} 元" if first_kline.total_market_cap else "  总市值: N/A")
        info(f"  流通市值: {first_kline.float_market_cap:,.0f} 元" if first_kline.float_market_cap else "  流通市值: N/A")
        info(f"  机构持仓占比: {first_kline.inst_holding_ratio:.4f}%" if first_kline.inst_holding_ratio else "  机构持仓占比: N/A")
        info(f"  可交易日数: {first_kline.trading_days}")
        
        # 统计涨停跌停情况
        limit_up_count = sum(1 for k in klines if k.is_limit_up)
        limit_down_count = sum(1 for k in klines if k.is_limit_down)
        
        info(f"\n统计信息:")
        info(f"  涨停天数: {limit_up_count}")
        info(f"  跌停天数: {limit_down_count}")
        
        # 验证数据完整性
        all_have_market_cap = all(k.total_market_cap > 0 for k in klines)
        all_have_trading_days = all(k.trading_days > 0 for k in klines)
        
        info(f"\n数据完整性检查:")
        info(f"  所有记录都有市值数据: {all_have_market_cap}")
        info(f"  所有记录都有交易日数: {all_have_trading_days}")
        
        return True
        
    except Exception as e:
        error(f"测试失败: {e}")
        import traceback
        error(traceback.format_exc())
        return False
    finally:
        await adapter.close()


async def test_save_and_read():
    """测试保存和读取包含扩展字段的K线数据"""
    from sqlalchemy.orm import Session
    from app.repositories.database import SessionLocal
    from app.repositories.stock_repository import StockRepository
    from app.models.types import KLineItem, Frequency
    from datetime import datetime
    
    info("\n\n测试保存和读取包含扩展字段的K线数据...")
    
    db = SessionLocal()
    try:
        repo = StockRepository(db)
        
        # 创建测试数据
        test_items = [
            KLineItem(
                time=datetime(2024, 1, 15, 0, 0, 0),
                open=10.0,
                high=11.0,
                low=9.5,
                close=10.8,
                volume=100000,
                amount=1080000.0,
                trade_date="2024-01-15",
                is_limit_up=True,
                is_limit_down=False,
                total_market_cap=1500000000.0,
                float_market_cap=1200000000.0,
                inst_holding_ratio=25.5,
                trading_days=5000
            )
        ]
        
        # 使用 object.__setattr__ 绕过 Pydantic 验证添加 symbol 属性
        for item in test_items:
            object.__setattr__(item, 'symbol', "000001.SZ")
        
        # 保存数据
        repo.save_klines(Frequency.FREQ_1D, test_items)
        info("测试数据保存成功")
        
        # 读取数据
        items = repo.get_klines(
            symbol="000001.SZ",
            freq=Frequency.FREQ_1D,
            start=datetime(2024, 1, 1),
            end=datetime(2024, 1, 31)
        )
        
        if items:
            item = items[0]
            info(f"\n从数据库读取的数据:")
            info(f"  日期: {item.time}")
            info(f"  收盘价: {item.close}")
            info(f"  交易日: {item.trade_date}")
            info(f"  是否涨停: {item.is_limit_up}")
            info(f"  是否跌停: {item.is_limit_down}")
            info(f"  总市值: {item.total_market_cap}")
            info(f"  流通市值: {item.float_market_cap}")
            info(f"  机构持仓占比: {item.inst_holding_ratio}")
            info(f"  可交易日数: {item.trading_days}")
            return True
        else:
            error("未读取到数据")
            return False
            
    except Exception as e:
        error(f"测试失败: {e}")
        import traceback
        error(traceback.format_exc())
        return False
    finally:
        db.close()


async def main():
    """主函数"""
    info("=" * 60)
    info("开始测试扩展字段功能")
    info("=" * 60)
    
    # 测试1: 获取日K线数据
    success1 = await test_fetch_daily_klines()
    
    # 测试2: 保存和读取数据
    success2 = await test_save_and_read()
    
    info("\n" + "=" * 60)
    if success1 and success2:
        info("所有测试通过！")
    else:
        error("部分测试失败")
    info("=" * 60)


if __name__ == "__main__":
    asyncio.run(main())