cyb50-quant/research/dragon/v2/tests/test_followthrough_profit_loop_review.py

from __future__ import annotations

import unittest
from pathlib import Path

from dragon_branch_configs import alpha_first_glued_followthrough_probe_config, alpha_first_glued_refined_hot_cap_config
from dragon_followthrough_profit_loop_review import collect_followthrough_reentries, collect_glued_block_candidates
from dragon_rc1_golden_baseline import _load_indicator_snapshot


class TestFollowthroughProfitLoopReview(unittest.TestCase):
    def setUp(self) -> None:
        self.base_dir = Path(__file__).resolve().parents[1]
        self.indexed, _ = _load_indicator_snapshot(self.base_dir)

    def test_collects_known_mid_zone_block_candidate(self) -> None:
        blocked = collect_glued_block_candidates(
            self.indexed,
            alpha_first_glued_refined_hot_cap_config(),
            branch_name="base_rc1",
        )
        target = blocked[
            (blocked["signal_date"] == "2020-11-30")
            & (blocked["subtype"] == "mid_zone_very_weak_b1")
        ].copy()
        self.assertEqual(len(target), 1)

    def test_maps_probe_mid_followthrough_origin(self) -> None:
        reentries = collect_followthrough_reentries(
            self.indexed,
            alpha_first_glued_followthrough_probe_config(),
            branch_name="probe_mid_zone",
        )
        self.assertEqual(len(reentries), 1)
        row = reentries.iloc[0]
        self.assertEqual(row["origin_date"], "2020-11-30")
        self.assertEqual(row["buy_date"], "2020-12-01")
        self.assertEqual(row["subtype"], "mid_zone_very_weak_b1")
        self.assertEqual(int(row["bars_waited"]), 1)


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