trading.ai/test_strategy.py

#!/usr/bin/env python3
"""
K线形态策略测试脚本
"""

import sys
from pathlib import Path
import pandas as pd
import numpy as np

# 将src目录添加到Python路径
current_dir = Path(__file__).parent
src_dir = current_dir / "src"
sys.path.insert(0, str(src_dir))

from src.data.data_fetcher import ADataFetcher
from src.utils.notification import NotificationManager
from src.strategy.kline_pattern_strategy import KLinePatternStrategy


def create_test_kline_data():
    """创建测试K线数据 - 包含两阳线+阴线+阳线形态"""
    dates = pd.date_range('2023-01-01', periods=10, freq='D')

    # 模拟K线数据
    test_data = {
        'trade_date': dates,
        'open': [10.0, 10.5, 11.0, 12.0, 11.5, 11.0, 10.5, 11.0, 11.8, 12.5],
        'high': [10.8, 11.2, 11.8, 12.5, 12.0, 11.5, 11.2, 11.5, 12.2, 13.0],
        'low':  [9.8,  10.3, 10.8, 11.8, 10.8, 10.5, 10.2, 10.8, 11.6, 12.3],
        'close':[10.6, 11.0, 11.5, 12.2, 11.0, 10.8, 11.2, 11.3, 12.0, 12.8],
        'volume': [1000] * 10
    }

    df = pd.DataFrame(test_data)
    print("测试K线数据:")
    print(df)
    print()

    return df


def test_pattern_detection():
    """测试形态检测功能"""
    print("="*60)
    print("           K线形态检测功能测试")
    print("="*60)

    # 创建测试配置
    strategy_config = {
        'min_entity_ratio': 0.55,
        'timeframes': ['daily'],
        'scan_stocks_count': 10,
        'analysis_days': 60
    }

    notification_config = {
        'dingtalk': {
            'enabled': False,
            'webhook_url': ''
        }
    }

    # 初始化组件
    data_fetcher = ADataFetcher()
    notification_manager = NotificationManager(notification_config)
    strategy = KLinePatternStrategy(data_fetcher, notification_manager, strategy_config)

    print("1. 策略信息:")
    print(strategy.get_strategy_summary())

    print("\n2. 测试K线特征计算:")
    test_df = create_test_kline_data()
    df_with_features = strategy.calculate_kline_features(test_df)

    print("添加特征后的数据:")
    relevant_cols = ['trade_date', 'open', 'high', 'low', 'close', 'is_yang', 'is_yin', 'entity_ratio']
    print(df_with_features[relevant_cols])

    print("\n3. 测试形态检测:")
    signals = strategy.detect_pattern(df_with_features)

    if signals:
        print(f"发现 {len(signals)} 个形态信号:")
        for i, signal in enumerate(signals, 1):
            print(f"\n信号 {i}:")
            print(f"  日期: {signal['date']}")
            print(f"  形态: {signal['pattern_type']}")
            print(f"  突破价格: {signal['breakout_price']:.2f}")
            print(f"  突破幅度: {signal['breakout_pct']:.2f}%")
            print(f"  阳线1实体比例: {signal['yang1_entity_ratio']:.1%}")
            print(f"  阳线2实体比例: {signal['yang2_entity_ratio']:.1%}")
    else:
        print("未发现形态信号")

    print("\n4. 测试真实股票数据:")
    test_stocks = ["000001.SZ", "000002.SZ"]  # 平安银行、万科A

    for stock_code in test_stocks:
        print(f"\n分析股票: {stock_code}")
        try:
            results = strategy.analyze_stock(stock_code, days=30)  # 分析最近30天

            total_signals = sum(len(signals) for signals in results.values())
            print(f"总信号数: {total_signals}")

            for timeframe, signals in results.items():
                if signals:
                    print(f"{timeframe}: {len(signals)}个信号")
                    # 显示最新信号
                    latest = signals[-1]
                    print(f"  最新: {latest['date']} {latest['breakout_price']:.2f}元")
                else:
                    print(f"{timeframe}: 无信号")

        except Exception as e:
            print(f"分析失败: {e}")

    print("\n5. 测试通知功能:")
    try:
        # 测试日志通知
        notification_manager.send_strategy_signal(
            stock_code="TEST001",
            stock_name="测试股票",
            timeframe="daily",
            signal_type="测试信号",
            price=15.50,
            additional_info={
                "阳线1实体比例": "65%",
                "阳线2实体比例": "70%",
                "突破幅度": "2.5%"
            }
        )
        print("✅ 通知功能测试完成（日志记录）")

    except Exception as e:
        print(f"❌ 通知功能测试失败: {e}")

    print("\n" + "="*60)
    print("           策略测试完成")
    print("="*60)


def test_weekly_monthly_conversion():
    """测试周线月线转换功能"""
    print("\n测试周线/月线数据转换:")

    # 创建更多天数的测试数据
    dates = pd.date_range('2023-01-01', periods=50, freq='D')

    test_data = {
        'trade_date': dates,
        'open': np.random.uniform(10, 15, 50),
        'high': np.random.uniform(15, 20, 50),
        'low': np.random.uniform(8, 12, 50),
        'close': np.random.uniform(10, 15, 50),
        'volume': np.random.randint(1000, 5000, 50)
    }

    daily_df = pd.DataFrame(test_data)

    strategy_config = {'min_entity_ratio': 0.55, 'timeframes': ['daily']}
    notification_config = {'dingtalk': {'enabled': False}}

    data_fetcher = ADataFetcher()
    notification_manager = NotificationManager(notification_config)
    strategy = KLinePatternStrategy(data_fetcher, notification_manager, strategy_config)

    # 测试周线转换
    weekly_df = strategy._convert_to_weekly(daily_df)
    print(f"日线数据: {len(daily_df)} 条")
    print(f"周线数据: {len(weekly_df)} 条")

    # 测试月线转换
    monthly_df = strategy._convert_to_monthly(daily_df)
    print(f"月线数据: {len(monthly_df)} 条")

    if not weekly_df.empty:
        print("\n周线数据样本:")
        print(weekly_df[['trade_date', 'open', 'high', 'low', 'close']].head())


if __name__ == "__main__":
    test_pattern_detection()
    test_weekly_monthly_conversion()