diff --git a/signal_v2/rate_limit.py b/signal_v2/rate_limit.py new file mode 100644 index 0000000..5482cda --- /dev/null +++ b/signal_v2/rate_limit.py @@ -0,0 +1,73 @@ +"""SignalDedup — SQLite-backed 24h duplicate signal blocker.""" +from __future__ import annotations +import sqlite3 +from contextlib import contextmanager +from datetime import datetime, timedelta +from pathlib import Path +from zoneinfo import ZoneInfo + +KST = ZoneInfo("Asia/Seoul") + + +def _now_iso() -> str: + """Test seam — overridable via monkeypatch.""" + return datetime.now(KST).isoformat() + + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS signal_dedup ( + ticker TEXT NOT NULL, + action TEXT NOT NULL, + last_sent TEXT NOT NULL, + confidence REAL NOT NULL, + PRIMARY KEY (ticker, action) +); +CREATE INDEX IF NOT EXISTS idx_signal_dedup_last_sent + ON signal_dedup(last_sent); +""" + + +class SignalDedup: + """24h dedup interface. WAL + busy_timeout=120000.""" + + def __init__(self, db_path: Path): + self._db_path = Path(db_path) + self._db_path.parent.mkdir(parents=True, exist_ok=True) + self._init_schema() + + @contextmanager + def _conn(self): + conn = sqlite3.connect(self._db_path, timeout=120.0) + try: + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA busy_timeout=120000") + yield conn + finally: + conn.close() + + def _init_schema(self) -> None: + with self._conn() as conn: + conn.executescript(_SCHEMA) + conn.commit() + + def is_recent(self, ticker: str, action: str, within_hours: int = 24) -> bool: + threshold_dt = datetime.fromisoformat(_now_iso()) - timedelta(hours=within_hours) + threshold_iso = threshold_dt.isoformat() + with self._conn() as conn: + row = conn.execute( + "SELECT last_sent FROM signal_dedup WHERE ticker = ? AND action = ?", + (ticker, action), + ).fetchone() + return row is not None and row[0] >= threshold_iso + + def record(self, ticker: str, action: str, confidence: float) -> None: + with self._conn() as conn: + conn.execute( + """INSERT INTO signal_dedup (ticker, action, last_sent, confidence) + VALUES (?, ?, ?, ?) + ON CONFLICT (ticker, action) DO UPDATE + SET last_sent = excluded.last_sent, + confidence = excluded.confidence""", + (ticker, action, _now_iso(), confidence), + ) + conn.commit() diff --git a/signal_v2/tests/test_rate_limit.py b/signal_v2/tests/test_rate_limit.py new file mode 100644 index 0000000..9cea492 --- /dev/null +++ b/signal_v2/tests/test_rate_limit.py @@ -0,0 +1,34 @@ +"""Tests for SignalDedup.""" +from datetime import datetime, timedelta +from zoneinfo import ZoneInfo + +from signal_v2.rate_limit import SignalDedup + +KST = ZoneInfo("Asia/Seoul") + + +def test_is_recent_returns_false_for_new_ticker_action(tmp_dedup_db): + dedup = SignalDedup(tmp_dedup_db) + assert dedup.is_recent("005930", "buy") is False + + +def test_is_recent_returns_true_within_24h(tmp_dedup_db): + dedup = SignalDedup(tmp_dedup_db) + dedup.record("005930", "buy", confidence=0.82) + assert dedup.is_recent("005930", "buy") is True + + +def test_is_recent_returns_false_after_24h(tmp_dedup_db, monkeypatch): + dedup = SignalDedup(tmp_dedup_db) + # Record with a timestamp 25 hours ago + now = datetime.now(KST) + fake_now = now - timedelta(hours=25) + monkeypatch.setattr( + "signal_v2.rate_limit._now_iso", lambda: fake_now.isoformat() + ) + dedup.record("005930", "buy", confidence=0.82) + # Reset to real now for is_recent check + monkeypatch.setattr( + "signal_v2.rate_limit._now_iso", lambda: now.isoformat() + ) + assert dedup.is_recent("005930", "buy", within_hours=24) is False