sanguo_vnpy_v2/tests/backtest/test_result_store.py

"""Tests for sanguo_backtest.result_store module."""
import pytest
import pandas as pd
from datetime import datetime
from sanguo_backtest.result_store import (
    BacktestResult,
    save_result,
    load_result,
    list_results
)


def test_save_and_load_result(temp_db_path):
    """Test saving and loading a backtest result without parquet files."""
    # Create a result without equity_curve and trades (file_dir=None by default)
    result = BacktestResult(
        task_id="test_task_001",
        type="cta",
        status="completed",
        strategy="DualThrustStrategy",
        symbol="IF2401",
        params={"window": 20, "entry_threshold": 0.5},
        start="2024-01-01",
        end="2024-03-31",
        statistics={
            "total_return": 0.15,
            "sharpe_ratio": 1.2,
            "max_drawdown": -0.08,
            "win_rate": 0.55
        }
    )

    # Save the result
    result_id = save_result(result, temp_db_path)
    assert result_id > 0
    assert isinstance(result_id, int)

    # Load the result
    loaded_result = load_result(result_id, temp_db_path)

    # Verify all fields match
    assert loaded_result.task_id == result.task_id
    assert loaded_result.type == result.type
    assert loaded_result.status == result.status
    assert loaded_result.strategy == result.strategy
    assert loaded_result.symbol == result.symbol
    assert loaded_result.params == result.params
    assert loaded_result.start == result.start
    assert loaded_result.end == result.end
    assert loaded_result.statistics == result.statistics
    assert loaded_result.equity_curve is None
    assert loaded_result.trades is None
    assert loaded_result.error_msg is None


def test_list_results_filter(temp_db_path):
    """Test listing results with type filter."""
    # Create multiple results of different types
    cta_result = BacktestResult(
        task_id="cta_task_001",
        type="cta",
        status="completed",
        strategy="DualThrustStrategy",
        symbol="IF2401",
        params={"window": 20},
        start="2024-01-01",
        end="2024-03-31",
        statistics={"total_return": 0.15}
    )

    factor_result = BacktestResult(
        task_id="factor_task_001",
        type="factor",
        status="completed",
        strategy="Alpha158Strategy",
        symbol="000001.SZ",
        params={"factors": ["alpha001", "alpha002"]},
        start="2024-01-01",
        end="2024-03-31",
        statistics={"total_return": 0.20}
    )

    # Save both results
    cta_id = save_result(cta_result, temp_db_path)
    factor_id = save_result(factor_result, temp_db_path)

    # List all results
    all_results = list_results(db_path=temp_db_path)
    assert len(all_results) == 2

    # List only CTA results
    cta_results = list_results(type_filter="cta", db_path=temp_db_path)
    assert len(cta_results) == 1
    assert cta_results[0].type == "cta"
    assert cta_results[0].task_id == "cta_task_001"

    # List only factor results
    factor_results = list_results(type_filter="factor", db_path=temp_db_path)
    assert len(factor_results) == 1
    assert factor_results[0].type == "factor"
    assert factor_results[0].task_id == "factor_task_001"


def test_failed_result_stores_error_msg(temp_db_path):
    """Test that failed results store error messages."""
    # Create a failed result with error message
    failed_result = BacktestResult(
        task_id="failed_task_001",
        type="cta",
        status="failed",
        strategy="DualThrustStrategy",
        symbol="IF2401",
        params={"window": 20},
        start="2024-01-01",
        end="2024-03-31",
        statistics={},  # Empty statistics for failed result
        error_msg="Data loading failed: insufficient historical data"
    )

    # Save the failed result
    result_id = save_result(failed_result, temp_db_path)

    # Load and verify error message is preserved
    loaded_result = load_result(result_id, temp_db_path)
    assert loaded_result.status == "failed"
    assert loaded_result.error_msg == "Data loading failed: insufficient historical data"
    assert loaded_result.statistics == {}


def test_save_sets_result_id_attribute(temp_db_path, tmp_path):
    """S1.1: save_result must set result.id to the DB row id (orchestrator uses it)."""
    result = BacktestResult(
        task_id="cta_id_test", type="cta", status="done", strategy="S", symbol="600000",
        params={"a": 1}, start="2024-01-01", end="2024-06-30", statistics={"sharpe": 1.2},
    )
    save_result(result, db_path=temp_db_path)
    assert result.id is not None
    assert isinstance(result.id, int)


def test_save_load_roundtrip_with_equity_curve(temp_db_path, tmp_path):
    """S1.1: equity_curve persists to parquet and reloads via result.id."""
    fdir = str(tmp_path / "files")
    result = BacktestResult(
        task_id="cta_eq_test", type="cta", status="done", strategy="S", symbol="600000",
        params={"a": 1}, start="2024-01-01", end="2024-06-30", statistics={"sharpe": 1.2},
        equity_curve=pd.DataFrame([
            {"date": "2024-01-01", "balance": 1_000_000},
            {"date": "2024-01-02", "balance": 1_010_000},
        ]),
    )
    save_result(result, db_path=temp_db_path, file_dir=fdir)
    assert result.id is not None

    loaded = load_result(result.id, temp_db_path)
    assert loaded.statistics == {"sharpe": 1.2}
    assert loaded.equity_curve is not None
    assert len(loaded.equity_curve) == 2
    assert loaded.equity_curve.iloc[1]["balance"] == 1_010_000