예제 소스

test_compare_lora_heuristic.py

import unittest

import numpy as np
import pandas as pd

from compare_lora_heuristic import (
    _first_threshold_outcome,
    _simulate,
    buy_hold_actions,
    extreme_diagnostics,
    minmax_actions,
    trend_heuristic_actions,
)


class StrategySimulationTest(unittest.TestCase):
    def test_buy_and_hold_profit_in_rising_market(self):
        result = _simulate(
            [100.0, 110.0],
            buy_hold_actions(2),
            name="hold",
            initial_krw=10_000.0,
        )
        self.assertGreater(result.liquidation_return_pct, 9.0)
        self.assertTrue(result.open_position)


class TrendHeuristicTest(unittest.TestCase):
    def test_falling_market_does_not_trigger_buy(self):
        close = np.linspace(200.0, 100.0, 100)
        frame = pd.DataFrame(
            {
                "close": close,
                "macd": np.linspace(1.0, -1.0, 100),
                "rsi": np.full(100, 30.0),
                "ext_crash_signal": np.zeros(100, dtype=bool),
            }
        )
        actions = trend_heuristic_actions(frame)
        self.assertEqual(int((actions == 1).sum()), 0)

    def test_confirmed_recovery_can_trigger_buy(self):
        close = np.r_[np.linspace(100.0, 80.0, 60), np.linspace(80.0, 130.0, 80)]
        frame = pd.DataFrame(
            {
                "close": close,
                "macd": np.linspace(-2.0, 3.0, len(close)),
                "rsi": np.full(len(close), 55.0),
                "ext_crash_signal": np.zeros(len(close), dtype=bool),
            }
        )
        actions = trend_heuristic_actions(frame)
        self.assertGreater(int((actions == 1).sum()), 0)


class MinMaxHeuristicTest(unittest.TestCase):
    def test_minmax_uses_low_entry_and_high_exit(self):
        close = np.r_[
            np.linspace(100.0, 200.0, 50),
            np.linspace(200.0, 100.0, 20),
            np.linspace(100.0, 200.0, 20),
        ]
        actions = minmax_actions(pd.DataFrame({"close": close}), window=20)
        self.assertGreater(int((actions == 1).sum()), 0)
        self.assertGreater(int((actions == 0).sum()), 0)
        self.assertLess(np.flatnonzero(actions == 1)[0], np.flatnonzero(actions == 0)[0])


class ExtremeDiagnosticsTest(unittest.TestCase):
    def test_first_threshold_outcome(self):
        self.assertEqual(_first_threshold_outcome(np.array([-0.02, 0.03]), 0.01), "down_first")
        self.assertEqual(_first_threshold_outcome(np.array([0.02, -0.03]), 0.01), "up_first")
        self.assertEqual(_first_threshold_outcome(np.array([0.001]), 0.01), "neither")

    def test_crash_continuation_is_counted(self):
        frame = pd.DataFrame(
            {
                "close": [100.0, 99.0, 97.0, 96.0],
                "ext_crash_signal": [True, False, False, False],
                "ext_surge_signal": [False, False, False, False],
            }
        )
        result = extreme_diagnostics(frame, horizons=[3], threshold=0.01)
        row = result["ext_crash_signal"]["horizons"]["3"]
        self.assertEqual(row["down_first"], 1)
        self.assertEqual(row["continuation_rate_pct"], 100.0)


if __name__ == "__main__":
    unittest.main()