From 2ad950d397904ed63e9b55b0b1ec88bb711f96d0 Mon Sep 17 00:00:00 2001
From: Davide Casale <davi0kprogramsthings@gmail.com>
Date: Thu, 14 Dec 2023 19:18:57 +0100
Subject: [PATCH] Fix bug in order_book.py and raw_order_book.py examples.

---
 examples/websocket/public/order_book.py     | 34 +++++++++++++--------
 examples/websocket/public/raw_order_book.py | 34 +++++++++++++--------
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/examples/websocket/public/order_book.py b/examples/websocket/public/order_book.py
index c79ade7..ff2cc7d 100644
--- a/examples/websocket/public/order_book.py
+++ b/examples/websocket/public/order_book.py
@@ -2,7 +2,7 @@
 
 import zlib
 from collections import OrderedDict
-from typing import Dict, List
+from typing import Any, Dict, List, cast
 
 from bfxapi import Client
 from bfxapi.types import TradingPairBook
@@ -15,8 +15,6 @@ class OrderBook:
             symbol: {"bids": OrderedDict(), "asks": OrderedDict()} for symbol in symbols
         }
 
-        self.cooldown: Dict[str, bool] = {symbol: False for symbol in symbols}
-
     def update(self, symbol: str, data: TradingPairBook) -> None:
         price, count, amount = data.price, data.count, data.amount
 
@@ -66,6 +64,15 @@ class OrderBook:
 
         return crc32 == checksum
 
+    def is_verifiable(self, symbol: str) -> bool:
+        return (
+            len(self.__order_book[symbol]["bids"]) >= 25
+            and len(self.__order_book[symbol]["asks"]) >= 25
+        )
+
+    def clear(self, symbol: str) -> None:
+        self.__order_book[symbol] = {"bids": OrderedDict(), "asks": OrderedDict()}
+
 
 SYMBOLS = ["tLTCBTC", "tETHUSD", "tETHBTC"]
 
@@ -100,17 +107,20 @@ def on_t_book_update(subscription: Book, data: TradingPairBook):
 async def on_checksum(subscription: Book, value: int):
     symbol = subscription["symbol"]
 
-    if order_book.verify(symbol, value):
-        order_book.cooldown[symbol] = False
-    elif not order_book.cooldown[symbol]:
-        print(
-            "Mismatch between local and remote checksums: "
-            f"restarting book for symbol <{symbol}>..."
-        )
+    if order_book.is_verifiable(symbol):
+        if not order_book.verify(symbol, value):
+            print(
+                "Mismatch between local and remote checksums: "
+                f"restarting book for symbol <{symbol}>..."
+            )
 
-        await bfx.wss.resubscribe(sub_id=subscription["sub_id"])
+            _subscription = cast(Dict[str, Any], subscription.copy())
 
-        order_book.cooldown[symbol] = True
+            await bfx.wss.unsubscribe(sub_id=_subscription.pop("sub_id"))
+
+            await bfx.wss.subscribe(**_subscription)
+
+            order_book.clear(symbol)
 
 
 bfx.wss.run()
diff --git a/examples/websocket/public/raw_order_book.py b/examples/websocket/public/raw_order_book.py
index fd6ebcc..b6c9fb2 100644
--- a/examples/websocket/public/raw_order_book.py
+++ b/examples/websocket/public/raw_order_book.py
@@ -2,7 +2,7 @@
 
 import zlib
 from collections import OrderedDict
-from typing import Dict, List
+from typing import Any, Dict, List, cast
 
 from bfxapi import Client
 from bfxapi.types import TradingPairRawBook
@@ -15,8 +15,6 @@ class RawOrderBook:
             symbol: {"bids": OrderedDict(), "asks": OrderedDict()} for symbol in symbols
         }
 
-        self.cooldown: Dict[str, bool] = {symbol: False for symbol in symbols}
-
     def update(self, symbol: str, data: TradingPairRawBook) -> None:
         order_id, price, amount = data.order_id, data.price, data.amount
 
@@ -66,6 +64,15 @@ class RawOrderBook:
 
         return crc32 == checksum
 
+    def is_verifiable(self, symbol: str) -> bool:
+        return (
+            len(self.__raw_order_book[symbol]["bids"]) >= 25
+            and len(self.__raw_order_book[symbol]["asks"]) >= 25
+        )
+
+    def clear(self, symbol: str) -> None:
+        self.__raw_order_book[symbol] = {"bids": OrderedDict(), "asks": OrderedDict()}
+
 
 SYMBOLS = ["tLTCBTC", "tETHUSD", "tETHBTC"]
 
@@ -100,17 +107,20 @@ def on_t_raw_book_update(subscription: Book, data: TradingPairRawBook):
 async def on_checksum(subscription: Book, value: int):
     symbol = subscription["symbol"]
 
-    if raw_order_book.verify(symbol, value):
-        raw_order_book.cooldown[symbol] = False
-    elif not raw_order_book.cooldown[symbol]:
-        print(
-            "Mismatch between local and remote checksums: "
-            f"restarting book for symbol <{symbol}>..."
-        )
+    if raw_order_book.is_verifiable(symbol):
+        if not raw_order_book.verify(symbol, value):
+            print(
+                "Mismatch between local and remote checksums: "
+                f"restarting book for symbol <{symbol}>..."
+            )
 
-        await bfx.wss.resubscribe(sub_id=subscription["sub_id"])
+            _subscription = cast(Dict[str, Any], subscription.copy())
 
-        raw_order_book.cooldown[symbol] = True
+            await bfx.wss.unsubscribe(sub_id=_subscription.pop("sub_id"))
+
+            await bfx.wss.subscribe(**_subscription)
+
+            raw_order_book.clear(symbol)
 
 
 bfx.wss.run()