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_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_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()