crypto_monitor/tests/test_account_risk_lib.py

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