dekun
97370926d6
Expose freeze_until_ms from risk API and tick hub/instance badges with remaining 1h/4h/daily time. Co-authored-by: Cursor <cursoragent@cursor.com>
299 lines
11 KiB
Python
299 lines
11 KiB
Python
import os
|
|
import sqlite3
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from unittest import mock
|
|
|
|
from account_risk_lib import (
|
|
CLOSE_SOURCE_USER_HUB,
|
|
CLOSE_SOURCE_USER_INSTANCE,
|
|
CLOSE_SOURCE_USER_TREND_STOP,
|
|
STATUS_DAILY,
|
|
STATUS_FREEZE_1H,
|
|
STATUS_FREEZE_4H,
|
|
STATUS_NORMAL,
|
|
account_risk_blocks_trading,
|
|
compute_account_risk_status,
|
|
ensure_account_risk_schema,
|
|
on_journal_saved,
|
|
on_manual_close,
|
|
on_user_initiated_close,
|
|
parse_mood_issues,
|
|
)
|
|
|
|
|
|
def _mem_conn():
|
|
conn = sqlite3.connect(":memory:")
|
|
conn.row_factory = sqlite3.Row
|
|
ensure_account_risk_schema(conn)
|
|
return conn
|
|
|
|
|
|
class AccountRiskLibTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.env_patch = mock.patch.dict(os.environ, {}, clear=False)
|
|
self.env_patch.start()
|
|
os.environ["RISK_CONTROL_ENABLED"] = "1"
|
|
os.environ["RISK_COOLING_HOURS_MANUAL"] = "4"
|
|
os.environ["RISK_COOLING_HOURS_MANUAL_JOURNAL"] = "1"
|
|
os.environ["RISK_MANUAL_CLOSE_DAILY_LIMIT"] = "2"
|
|
os.environ["RISK_MOOD_ISSUES_DAILY_FREEZE"] = "1"
|
|
|
|
def tearDown(self):
|
|
self.env_patch.stop()
|
|
|
|
def test_user_instance_sets_4h_cooloff(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source=CLOSE_SOURCE_USER_INSTANCE,
|
|
trade_record_id=101,
|
|
closed_at_ms=close_ms,
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["status"], STATUS_FREEZE_4H)
|
|
self.assertFalse(st["can_trade"])
|
|
|
|
def test_invalid_source_ignored(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source="exchange_tpsl",
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["status"], STATUS_NORMAL)
|
|
|
|
def test_second_user_close_daily_freeze(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn, source=CLOSE_SOURCE_USER_HUB, closed_at_ms=close_ms, trading_day="2026-06-14", now=now
|
|
)
|
|
on_user_initiated_close(
|
|
conn, source=CLOSE_SOURCE_USER_HUB, closed_at_ms=close_ms + 1000, trading_day="2026-06-14", now=now
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["status"], STATUS_DAILY)
|
|
|
|
def test_hub_close_all_count(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source=CLOSE_SOURCE_USER_HUB,
|
|
closed_at_ms=close_ms,
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
count=2,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["manual_close_count"], 2)
|
|
self.assertEqual(st["status"], STATUS_DAILY)
|
|
|
|
def test_trend_stop_counts_as_manual(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source=CLOSE_SOURCE_USER_TREND_STOP,
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["manual_close_count"], 1)
|
|
self.assertEqual(st["status"], STATUS_FREEZE_4H)
|
|
|
|
def test_journal_manual_with_note_reduces_to_1h(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_manual_close(conn, trade_record_id=9, closed_at_ms=close_ms, trading_day="2026-06-14", now=now)
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="手动平仓",
|
|
early_exit_note="违反计划提前离场",
|
|
mood_issues_raw="",
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["status"], STATUS_FREEZE_1H)
|
|
|
|
def test_journal_hub_close_without_pending_reduces_to_1h(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source=CLOSE_SOURCE_USER_HUB,
|
|
closed_at_ms=close_ms,
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="手动平仓",
|
|
early_exit_note="中控全平后复盘说明",
|
|
mood_issues_raw="",
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["status"], STATUS_FREEZE_1H)
|
|
|
|
def test_journal_reduces_when_manual_count_cleared_but_cooloff_active(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 15, 10, 0, 0)
|
|
now_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
close_ms = now_ms - 3600 * 1000
|
|
until_ms = close_ms + 4 * 3600 * 1000
|
|
conn.execute(
|
|
"""UPDATE account_risk_state SET
|
|
trading_day='2026-06-15',
|
|
manual_close_count=0,
|
|
cooloff_until_ms=?,
|
|
cooloff_hours=4,
|
|
last_close_at_ms=?,
|
|
daily_frozen=0
|
|
WHERE id=1""",
|
|
(until_ms, close_ms),
|
|
)
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="手动平仓",
|
|
early_exit_note="切日后补复盘",
|
|
mood_issues_raw="",
|
|
trading_day="2026-06-15",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-15", now=now)
|
|
self.assertEqual(st["status"], STATUS_FREEZE_1H)
|
|
|
|
def test_journal_late_save_still_gets_1h_from_now(self):
|
|
conn = _mem_conn()
|
|
close_at = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(close_at.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source=CLOSE_SOURCE_USER_INSTANCE,
|
|
closed_at_ms=close_ms,
|
|
trading_day="2026-06-14",
|
|
now=close_at,
|
|
)
|
|
journal_at = datetime(2026, 6, 14, 14, 0, 0)
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="手动平仓",
|
|
early_exit_note="补写复盘说明",
|
|
mood_issues_raw="",
|
|
trading_day="2026-06-14",
|
|
now=journal_at,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=journal_at)
|
|
self.assertEqual(st["status"], STATUS_FREEZE_1H)
|
|
self.assertEqual(
|
|
st["cooloff_until_ms"],
|
|
int(journal_at.replace(tzinfo=timezone.utc).timestamp() * 1000) + 3600 * 1000,
|
|
)
|
|
|
|
def test_journal_mood_issues_daily_freeze(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="止损",
|
|
early_exit_note="",
|
|
mood_issues_raw=["报复开仓"],
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertEqual(st["status"], STATUS_DAILY)
|
|
|
|
def test_cooloff_expired_returns_normal(self):
|
|
conn = _mem_conn()
|
|
start = datetime(2026, 6, 14, 8, 0, 0)
|
|
close_ms = int(start.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn, source=CLOSE_SOURCE_USER_INSTANCE, closed_at_ms=close_ms, trading_day="2026-06-14", now=start
|
|
)
|
|
later = datetime(2026, 6, 14, 13, 0, 0)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=later)
|
|
self.assertEqual(st["status"], STATUS_NORMAL)
|
|
|
|
def test_trading_day_reset_clears_daily_frozen(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="止损",
|
|
early_exit_note="",
|
|
mood_issues_raw="扛单",
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
next_day = datetime(2026, 6, 15, 8, 0, 0)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-15", now=next_day)
|
|
self.assertEqual(st["status"], STATUS_NORMAL)
|
|
|
|
def test_parse_mood_issues_filters_unknown(self):
|
|
self.assertEqual(parse_mood_issues("怕踏空,未知标签,扛单"), ["怕踏空", "扛单"])
|
|
|
|
def test_enrich_countdown_for_daily_and_cooloff(self):
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
close_ms = int(now.replace(tzinfo=timezone.utc).timestamp() * 1000)
|
|
on_user_initiated_close(
|
|
conn,
|
|
source=CLOSE_SOURCE_USER_INSTANCE,
|
|
closed_at_ms=close_ms,
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
from account_risk_lib import enrich_risk_status_countdown
|
|
|
|
st = enrich_risk_status_countdown(st, now=now, daily_reset_hour=8)
|
|
self.assertGreater(st["freeze_remaining_sec"], 0)
|
|
self.assertEqual(st["freeze_until_ms"], st["cooloff_until_ms"])
|
|
|
|
on_journal_saved(
|
|
conn,
|
|
early_exit_trigger="止损",
|
|
early_exit_note="",
|
|
mood_issues_raw="扛单",
|
|
trading_day="2026-06-14",
|
|
now=now,
|
|
)
|
|
st2 = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
st2 = enrich_risk_status_countdown(st2, now=now, daily_reset_hour=8)
|
|
self.assertTrue(st2["daily_frozen"])
|
|
self.assertGreater(st2["freeze_remaining_sec"], 0)
|
|
self.assertIsNotNone(st2["freeze_until_ms"])
|
|
|
|
def test_disabled_risk_control(self):
|
|
os.environ["RISK_CONTROL_ENABLED"] = "0"
|
|
conn = _mem_conn()
|
|
now = datetime(2026, 6, 14, 12, 0, 0)
|
|
on_user_initiated_close(
|
|
conn, source=CLOSE_SOURCE_USER_INSTANCE, trading_day="2026-06-14", now=now
|
|
)
|
|
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
|
self.assertFalse(st["enabled"])
|
|
self.assertTrue(st["can_trade"])
|
|
ok, _ = account_risk_blocks_trading(conn, trading_day="2026-06-14", now=now)
|
|
self.assertTrue(ok)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|