From 655f6c2292fe806a06d2607c5606a5a382861f92 Mon Sep 17 00:00:00 2001 From: Matt Hauff Date: Fri, 27 Jan 2023 16:20:01 -0700 Subject: [PATCH] Add a persitent id table to notification store (#14418) * Add a persitent id table to notification store * rename table * Add a test * lint --- chia/wallet/notification_store.py | 10 +++++++++- tests/wallet/test_notifications.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/chia/wallet/notification_store.py b/chia/wallet/notification_store.py index f296bd627650..02554ef55c77 100644 --- a/chia/wallet/notification_store.py +++ b/chia/wallet/notification_store.py @@ -46,6 +46,8 @@ class NotificationStore: "CREATE TABLE IF NOT EXISTS notifications(" "coin_id blob PRIMARY KEY," "msg blob," "amount blob" ")" ) + await conn.execute("CREATE TABLE IF NOT EXISTS all_notification_ids(coin_id blob PRIMARY KEY)") + try: await conn.execute("ALTER TABLE notifications ADD COLUMN height bigint DEFAULT 0") except sqlite3.OperationalError as e: @@ -74,6 +76,10 @@ class NotificationStore: notification.height, ), ) + cursor = await conn.execute( + "INSERT OR REPLACE INTO all_notification_ids (coin_id) VALUES(?)", + (notification.coin_id,), + ) await cursor.close() async def get_notifications(self, coin_ids: List[bytes32]) -> List[Notification]: @@ -155,7 +161,9 @@ class NotificationStore: async def notification_exists(self, id: bytes32) -> bool: async with self.db_wrapper.reader_no_transaction() as conn: - async with conn.execute("SELECT EXISTS (SELECT 1 from notifications WHERE coin_id=?)", (id,)) as cursor: + async with conn.execute( + "SELECT EXISTS (SELECT 1 from all_notification_ids WHERE coin_id=?)", (id,) + ) as cursor: row = await cursor.fetchone() assert row is not None exists: bool = row[0] > 0 diff --git a/tests/wallet/test_notifications.py b/tests/wallet/test_notifications.py index 7e4dc97b0291..e256a7893868 100644 --- a/tests/wallet/test_notifications.py +++ b/tests/wallet/test_notifications.py @@ -97,6 +97,16 @@ async def test_notifications(self_hostname: str, two_wallet_nodes: Any, trusted: notification_manager_1 = wsm_1.notification_manager notification_manager_2 = wsm_2.notification_manager + func = notification_manager_2.potentially_add_new_notification + notification_manager_2.most_recent_args = tuple() + + async def track_coin_state(*args: Any) -> bool: + notification_manager_2.most_recent_args = args + result: bool = await func(*args) + return result + + notification_manager_2.potentially_add_new_notification = track_coin_state + for case in ("block all", "block too low", "allow", "allow_larger", "block_too_large"): msg: bytes = bytes(case, "utf8") if case == "block all": @@ -165,3 +175,7 @@ async def test_notifications(self_hostname: str, two_wallet_nodes: Any, trusted: await notification_manager_2.notification_store.add_notification(notifications[0]) await notification_manager_2.notification_store.delete_notifications([n.coin_id for n in notifications]) assert len(await notification_manager_2.notification_store.get_all_notifications()) == 0 + + assert not await func(*notification_manager_2.most_recent_args) + await notification_manager_2.notification_store.delete_all_notifications() + assert not await func(*notification_manager_2.most_recent_args)