poly_strategy/tests/test_alerts.py at main · WW-shan/poly_strategy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import json
import tempfile
import unittest
from pathlib import Path

from poly_strategy.alerts import latest_alert_market_ids, latest_monitor_alerts, read_opportunity_alerts, write_alerts


class AlertTests(unittest.TestCase):
    def test_latest_monitor_alerts_extracts_stable_paper_trades(self):
        with tempfile.TemporaryDirectory() as tmp:
            report = Path(tmp) / "report.jsonl"
            out = Path(tmp) / "alerts.ndjson"
            report.write_text(
                json.dumps({"type": "paper_monitor_summary"}) + "\n"
                + json.dumps(
                    {
                        "type": "realtime_monitor_iteration",
                        "ts": "2026-05-09T00:00:00Z",
                        "iteration": 3,
                        "last_snapshot_ts": "2026-05-09T00:00:00Z",
                        "stable_paper_trades": [
                            _trade("high", paper_roi=0.05, paper_edge=0.2),
                            _trade("low", paper_roi=0.005, paper_edge=0.1),
                        ],
                    }
                )
                + "\n"
            )

            alerts = latest_monitor_alerts(report, min_paper_roi=0.01)
            count = write_alerts(alerts, out)
            written = [json.loads(line) for line in out.read_text().splitlines()]

        self.assertEqual(count, 1)
        self.assertEqual(alerts[0]["type"], "opportunity_alert")
        self.assertEqual(alerts[0]["key"], "high")
        self.assertEqual(alerts[0]["market_ids"], ["m1", "m2"])
        self.assertEqual(written[0]["alert_kind"], "stable_paper_trade")

    def test_latest_monitor_alerts_can_include_current_opportunities(self):
        with tempfile.TemporaryDirectory() as tmp:
            report = Path(tmp) / "report.jsonl"
            report.write_text(
                json.dumps(
                    {
                        "type": "paper_monitor_iteration",
                        "iteration": 1,
                        "stable_paper_trades": [],
                        "stable_opportunities": [_opportunity("stable")],
                        "current_opportunities": [_opportunity("current")],
                    }
                )
                + "\n"
            )

            alerts = latest_monitor_alerts(report, include_current=True)

        self.assertEqual([alert["alert_kind"] for alert in alerts], ["stable_opportunity", "current_opportunity"])

    def test_write_alerts_suppresses_duplicates_with_cooldown_state(self):
        with tempfile.TemporaryDirectory() as tmp:
            out = Path(tmp) / "alerts.ndjson"
            state = Path(tmp) / "state.json"
            rows = [
                {
                    "type": "opportunity_alert",
                    "alert_kind": "stable_paper_trade",
                    "kind": "yes_no_bundle",
                    "key": "same",
                    "market_ids": ["m1"],
                }
            ]

            first_count = write_alerts(rows, out, state_path=state, cooldown_seconds=60)
            second_count = write_alerts(rows, out, state_path=state, cooldown_seconds=60)
            written = out.read_text().splitlines()

        self.assertEqual(first_count, 1)
        self.assertEqual(second_count, 0)
        self.assertEqual(len(written), 1)

    def test_latest_alert_market_ids_reads_recent_alerts(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "alerts.ndjson"
            path.write_text(
                json.dumps({"type": "opportunity_alert", "market_ids": ["old"]})
                + "\n"
                + json.dumps({"type": "opportunity_alert", "market_ids": ["m1", "m2"]})
                + "\n"
                + json.dumps({"type": "opportunity_alert", "market_ids": ["m2", "m3"]})
                + "\n"
            )

            rows = read_opportunity_alerts(path)
            market_ids = latest_alert_market_ids(path, max_alerts=2)

        self.assertEqual(len(rows), 3)
        self.assertEqual(market_ids, ["m1", "m2", "m3"])


def _trade(key: str, paper_roi: float, paper_edge: float) -> dict:
    row = _opportunity(key)
    row.update({"paper_roi": paper_roi, "paper_edge": paper_edge})
    return row


def _opportunity(key: str) -> dict:
    return {
        "key": key,
        "kind": "yes_no_bundle",
        "net_edge_per_share": 0.02,
        "total_edge": 0.1,
        "legs": [{"market_id": "m1"}, {"market_id": "m2"}],
    }


if __name__ == "__main__":
    unittest.main()