95 lines
3.9 KiB
Python
95 lines
3.9 KiB
Python
import unittest
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from primethai.config import AppConfig, StrategyConfig
|
|
from primethai.scanner import scan_candidates
|
|
from primethai.setups import setup_for
|
|
|
|
|
|
class SetupScannerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.previous = pd.Series({
|
|
"open": 10.0, "high": 11.0, "low": 9.5, "close": 10.0,
|
|
"ema21_high": 10.0,
|
|
})
|
|
self.row = pd.Series({
|
|
"open": 10.1, "high": 10.8, "low": 9.7, "close": 10.5,
|
|
"ema21_high": 10.2, "ema21_close": 10.0, "ema21_low": 9.0,
|
|
"trend_state": "NEUTRAL", "higher_low": False,
|
|
"extension_atr": 0.5,
|
|
})
|
|
|
|
def test_scanner_setup_classifies_base_breakout_and_ema_reclaim(self):
|
|
setup, entry, stop = setup_for(self.row, self.previous, prior_high=10.4)
|
|
self.assertEqual(setup, "Base Breakout")
|
|
self.assertEqual((entry, stop), (10.5, 9.0))
|
|
|
|
# Keep the current close below yesterday's high so the EMA reclaim is distinct.
|
|
prior = self.previous.copy()
|
|
prior["high"] = 11.0
|
|
prior["close"] = 10.0
|
|
self.row["close"] = 10.5
|
|
setup, _, _ = setup_for(self.row, prior, prior_high=11.0)
|
|
self.assertEqual(setup, "21DMA High Reclaim")
|
|
|
|
def test_extension_filter_blocks_scanner_setups(self):
|
|
self.row["extension_atr"] = 1.2
|
|
setup, entry, stop = setup_for(self.row, self.previous, prior_high=10.4)
|
|
self.assertIsNone(setup)
|
|
self.assertIsNone(entry)
|
|
self.assertIsNone(stop)
|
|
|
|
def test_watch_only_candidate_has_no_position_size(self):
|
|
index = pd.bdate_range("2026-01-01", periods=25)
|
|
frame = pd.DataFrame({
|
|
"open": 10.0, "high": 10.6, "low": 9.6, "close": 10.0,
|
|
"ema21_high": 9.5, "ema21_close": 9.4, "ema21_low": 9.0,
|
|
"trend_state": "NEUTRAL", "higher_low": False,
|
|
"extension_atr": 0.5, "avg_value_turnover20": 50_000_000.0,
|
|
"rs_percentile": 80.0, "volume_contracting": False,
|
|
"tightness_score": 50.0, "weekly_trend_state": "NEUTRAL",
|
|
"rs_1m": 0.0, "rs_3m": 0.0, "rs_12m": 0.0,
|
|
"rs_1m_vs_set": 0.0, "rs_3m_vs_set": 0.0,
|
|
"rs_12m_vs_set": 0.0, "rs_1m_vs_set100ew": 0.0,
|
|
"rs_3m_vs_set100ew": 0.0, "rs_12m_vs_set100ew": 0.0,
|
|
"52w_proximity_pct": 90.0,
|
|
}, index=index)
|
|
market = pd.DataFrame({
|
|
"regime": ["CONFIRMED_UPTREND"], "new_risk_allowed": [True],
|
|
}, index=[index[-1]])
|
|
config = AppConfig(strategy=StrategyConfig())
|
|
|
|
result = scan_candidates({"TEST": frame}, market, config=config)
|
|
|
|
self.assertEqual(result.candidates.iloc[0]["entry_type"], "WATCH")
|
|
self.assertEqual(result.candidates.iloc[0]["shares"], 0)
|
|
self.assertEqual(result.candidates.iloc[0]["risk_pct"], 0.0)
|
|
|
|
frame.loc[index[-1], "close"] = 10.8
|
|
frame.loc[index[-1], "high"] = 11.0
|
|
market.loc[index[-1], "new_risk_allowed"] = False
|
|
market.loc[index[-1], "regime"] = "CORRECTION"
|
|
|
|
blocked = scan_candidates({"TEST": frame}, market, config=config)
|
|
|
|
candidate = blocked.candidates.iloc[0]
|
|
self.assertEqual(candidate["setup"], "Watch")
|
|
self.assertEqual(candidate["entry_type"], "WATCH")
|
|
self.assertEqual(candidate["shares"], 0)
|
|
self.assertEqual(candidate["risk_pct"], 0.0)
|
|
self.assertIn("market regime blocks new entries", candidate["reasons"])
|
|
|
|
market.attrs["members_by_date"] = {pd.Timestamp(index[-1]): ["OTHER"]}
|
|
nonmember = scan_candidates({"TEST": frame}, market, config=config)
|
|
self.assertTrue(nonmember.candidates.empty)
|
|
|
|
market.attrs["members_by_date"] = {pd.Timestamp(index[-1]): ["TEST"]}
|
|
stale = scan_candidates({"TEST": frame.iloc[:-1]}, market, config=config)
|
|
self.assertTrue(stale.candidates.empty)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|