From 8f7d6950c29105473eab346e49bb2279633c6fac Mon Sep 17 00:00:00 2001
From: Davide Casale <davi0kprogramsthings@gmail.com>
Date: Wed, 24 Apr 2024 17:33:32 +0200
Subject: [PATCH] Add support for event balance_update.

---
 bfxapi/types/__init__.py                             | 1 +
 bfxapi/types/dataclasses.py                          | 6 ++++++
 bfxapi/types/serializers.py                          | 6 ++++++
 bfxapi/websocket/_event_emitter/bfx_event_emitter.py | 1 +
 bfxapi/websocket/_handlers/auth_events_handler.py    | 2 ++
 5 files changed, 16 insertions(+)

diff --git a/bfxapi/types/__init__.py b/bfxapi/types/__init__.py
index 0b1df46..cf17b88 100644
--- a/bfxapi/types/__init__.py
+++ b/bfxapi/types/__init__.py
@@ -1,5 +1,6 @@
 from .dataclasses import (
     BalanceAvailable,
+    BalanceInfo,
     BaseMarginInfo,
     Candle,
     CurrencyConversion,
diff --git a/bfxapi/types/dataclasses.py b/bfxapi/types/dataclasses.py
index c3e4458..c1d40d6 100644
--- a/bfxapi/types/dataclasses.py
+++ b/bfxapi/types/dataclasses.py
@@ -608,6 +608,12 @@ class DerivativePositionCollateralLimits(_Type):
     max_collateral: float
 
 
+@dataclass
+class BalanceInfo(_Type):
+    aum: float
+    aum_net: float
+
+
 # endregion
 
 # region Dataclass definitions for types of merchant use
diff --git a/bfxapi/types/serializers.py b/bfxapi/types/serializers.py
index b9abaca..25054e4 100644
--- a/bfxapi/types/serializers.py
+++ b/bfxapi/types/serializers.py
@@ -904,4 +904,10 @@ DerivativePositionCollateralLimits = generate_labeler_serializer(
     labels=["min_collateral", "max_collateral"],
 )
 
+BalanceInfo = generate_labeler_serializer(
+    name="BalanceInfo",
+    klass=dataclasses.BalanceInfo,
+    labels=["aum", "aum_net"],
+)
+
 # endregion
diff --git a/bfxapi/websocket/_event_emitter/bfx_event_emitter.py b/bfxapi/websocket/_event_emitter/bfx_event_emitter.py
index 2e48a6b..21bbfd6 100644
--- a/bfxapi/websocket/_event_emitter/bfx_event_emitter.py
+++ b/bfxapi/websocket/_event_emitter/bfx_event_emitter.py
@@ -67,6 +67,7 @@ _COMMON = [
     "base_margin_info",
     "symbol_margin_info",
     "funding_info_update",
+    "balance_update",
     "notification",
     "on-req-notification",
     "ou-req-notification",
diff --git a/bfxapi/websocket/_handlers/auth_events_handler.py b/bfxapi/websocket/_handlers/auth_events_handler.py
index 0a2230a..486e9f7 100644
--- a/bfxapi/websocket/_handlers/auth_events_handler.py
+++ b/bfxapi/websocket/_handlers/auth_events_handler.py
@@ -34,6 +34,7 @@ class AuthEventsHandler:
         "ws": "wallet_snapshot",
         "wu": "wallet_update",
         "fiu": "funding_info_update",
+        "bu": "balance_update",
     }
 
     __SERIALIZERS: Dict[Tuple[str, ...], serializers._Serializer] = {
@@ -45,6 +46,7 @@ class AuthEventsHandler:
         ("fls", "fln", "flu", "flc"): serializers.FundingLoan,
         ("ws", "wu"): serializers.Wallet,
         ("fiu",): serializers.FundingInfo,
+        ("bu",): serializers.BalanceInfo,
     }
 
     def __init__(self, event_emitter: EventEmitter) -> None: