fix(risk): correct freeze countdown timezone (Asia/Shanghai)
Treat naive app datetimes as local time, normalize legacy UTC-ms rows, and resolve cooloff end from stored until or last_close+duration. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -1,8 +1,9 @@
|
||||
import os
|
||||
import sqlite3
|
||||
import unittest
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime
|
||||
from unittest import mock
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from account_risk_lib import (
|
||||
CLOSE_SOURCE_USER_HUB,
|
||||
@@ -14,6 +15,7 @@ from account_risk_lib import (
|
||||
STATUS_NORMAL,
|
||||
account_risk_blocks_trading,
|
||||
compute_account_risk_status,
|
||||
enrich_risk_status_countdown,
|
||||
ensure_account_risk_schema,
|
||||
on_journal_saved,
|
||||
on_manual_close,
|
||||
@@ -21,6 +23,8 @@ from account_risk_lib import (
|
||||
parse_mood_issues,
|
||||
)
|
||||
|
||||
APP_TZ = ZoneInfo("Asia/Shanghai")
|
||||
|
||||
|
||||
def _mem_conn():
|
||||
conn = sqlite3.connect(":memory:")
|
||||
@@ -29,6 +33,10 @@ def _mem_conn():
|
||||
return conn
|
||||
|
||||
|
||||
def _local_ms(dt_naive: datetime) -> int:
|
||||
return int(dt_naive.replace(tzinfo=APP_TZ).timestamp() * 1000)
|
||||
|
||||
|
||||
class AccountRiskLibTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.env_patch = mock.patch.dict(os.environ, {}, clear=False)
|
||||
@@ -38,6 +46,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
os.environ["RISK_COOLING_HOURS_MANUAL_JOURNAL"] = "1"
|
||||
os.environ["RISK_MANUAL_CLOSE_DAILY_LIMIT"] = "2"
|
||||
os.environ["RISK_MOOD_ISSUES_DAILY_FREEZE"] = "1"
|
||||
os.environ["APP_TIMEZONE"] = "Asia/Shanghai"
|
||||
|
||||
def tearDown(self):
|
||||
self.env_patch.stop()
|
||||
@@ -45,7 +54,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
close_ms = _local_ms(now)
|
||||
on_user_initiated_close(
|
||||
conn,
|
||||
source=CLOSE_SOURCE_USER_INSTANCE,
|
||||
@@ -57,6 +66,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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"])
|
||||
self.assertAlmostEqual(st["freeze_remaining_sec"], 4 * 3600, delta=2)
|
||||
|
||||
def test_invalid_source_ignored(self):
|
||||
conn = _mem_conn()
|
||||
@@ -73,7 +83,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
close_ms = _local_ms(now)
|
||||
on_user_initiated_close(
|
||||
conn, source=CLOSE_SOURCE_USER_HUB, closed_at_ms=close_ms, trading_day="2026-06-14", now=now
|
||||
)
|
||||
@@ -86,7 +96,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
close_ms = _local_ms(now)
|
||||
on_user_initiated_close(
|
||||
conn,
|
||||
source=CLOSE_SOURCE_USER_HUB,
|
||||
@@ -111,11 +121,12 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
self.assertAlmostEqual(st["freeze_remaining_sec"], 4 * 3600, delta=2)
|
||||
|
||||
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)
|
||||
close_ms = _local_ms(now)
|
||||
on_manual_close(conn, trade_record_id=9, closed_at_ms=close_ms, trading_day="2026-06-14", now=now)
|
||||
on_journal_saved(
|
||||
conn,
|
||||
@@ -127,11 +138,12 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
)
|
||||
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
||||
self.assertEqual(st["status"], STATUS_FREEZE_1H)
|
||||
self.assertAlmostEqual(st["freeze_remaining_sec"], 3600, delta=2)
|
||||
|
||||
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)
|
||||
close_ms = _local_ms(now)
|
||||
on_user_initiated_close(
|
||||
conn,
|
||||
source=CLOSE_SOURCE_USER_HUB,
|
||||
@@ -153,7 +165,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
now_ms = _local_ms(now)
|
||||
close_ms = now_ms - 3600 * 1000
|
||||
until_ms = close_ms + 4 * 3600 * 1000
|
||||
conn.execute(
|
||||
@@ -181,7 +193,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
close_ms = _local_ms(close_at)
|
||||
on_user_initiated_close(
|
||||
conn,
|
||||
source=CLOSE_SOURCE_USER_INSTANCE,
|
||||
@@ -200,10 +212,30 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
)
|
||||
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,
|
||||
self.assertEqual(st["cooloff_until_ms"], _local_ms(journal_at) + 3600 * 1000)
|
||||
|
||||
def test_legacy_naive_utc_ms_countdown_normalized(self):
|
||||
conn = _mem_conn()
|
||||
now = datetime(2026, 6, 14, 12, 0, 0)
|
||||
now_ms = _local_ms(now)
|
||||
offset_ms = 8 * 3600 * 1000
|
||||
legacy_close = now_ms + offset_ms
|
||||
legacy_until = legacy_close + 4 * 3600 * 1000
|
||||
conn.execute(
|
||||
"""UPDATE account_risk_state SET
|
||||
trading_day='2026-06-14',
|
||||
manual_close_count=1,
|
||||
cooloff_until_ms=?,
|
||||
cooloff_hours=4,
|
||||
last_close_at_ms=?,
|
||||
daily_frozen=0
|
||||
WHERE id=1""",
|
||||
(legacy_until, legacy_close),
|
||||
)
|
||||
st = compute_account_risk_status(conn, trading_day="2026-06-14", now=now)
|
||||
st = enrich_risk_status_countdown(st, now=now, daily_reset_hour=8)
|
||||
self.assertEqual(st["status"], STATUS_FREEZE_4H)
|
||||
self.assertAlmostEqual(st["freeze_remaining_sec"], 4 * 3600, delta=2)
|
||||
|
||||
def test_journal_mood_issues_daily_freeze(self):
|
||||
conn = _mem_conn()
|
||||
@@ -222,7 +254,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
close_ms = _local_ms(start)
|
||||
on_user_initiated_close(
|
||||
conn, source=CLOSE_SOURCE_USER_INSTANCE, closed_at_ms=close_ms, trading_day="2026-06-14", now=start
|
||||
)
|
||||
@@ -251,7 +283,7 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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)
|
||||
close_ms = _local_ms(now)
|
||||
on_user_initiated_close(
|
||||
conn,
|
||||
source=CLOSE_SOURCE_USER_INSTANCE,
|
||||
@@ -260,8 +292,6 @@ class AccountRiskLibTests(unittest.TestCase):
|
||||
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"])
|
||||
|
||||
Reference in New Issue
Block a user