From 1ca73149dc4ef3156ee53c7d75e2864e962b0f16 Mon Sep 17 00:00:00 2001
From: tekwani <tekwani.hemant@gmail.com>
Date: Mon, 23 Oct 2023 18:39:30 +0530
Subject: [PATCH] pass wallets data to strategy exeuctors (#251)

---
 .../api/handlers/strategy/strategy_manager.js | 43 +++++++++++++++++++
 .../lib/ws_servers/api/strategy/manager.js    |  3 +-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/lib/ws_servers/api/handlers/strategy/strategy_manager.js b/lib/ws_servers/api/handlers/strategy/strategy_manager.js
index 9828e9f0..00d9397e 100644
--- a/lib/ws_servers/api/handlers/strategy/strategy_manager.js
+++ b/lib/ws_servers/api/handlers/strategy/strategy_manager.js
@@ -23,6 +23,8 @@ const debug = require('debug')('bfx:hf:server:strategy-manager')
 
 const CLOSE_CONNECTIONS_DELAY = 60 * 1000 // one minute
 const EXECUTION_RESULTS_OMIT_FIELDS = ['candles', 'trades']
+const WALLET_SNAPSHOT_EVENT = 'auth:ws'
+const WALLET_UPDATE_EVENT = 'auth:wu'
 
 class StrategyManager {
   constructor (settings, bcast, strategyExecutionDB, sendDataToMetricsServer) {
@@ -106,11 +108,51 @@ class StrategyManager {
     this.ws2Manager.onceWS('event:auth:success', {}, this._onAuthSuccess.bind(this, resolve))
     this.ws2Manager.onceWS('event:auth:error', {}, this._onAuthError.bind(this, reject))
     this.ws2Manager.onWS('close', {}, this._onWSClose.bind(this))
+
+    // user wallets data for startegy execution
+    this.ws2Manager.onWS(WALLET_SNAPSHOT_EVENT, {}, async (data) => {
+      this._onWalletsData(data, WALLET_SNAPSHOT_EVENT)
+    })
+
+    this.ws2Manager.onWS(WALLET_UPDATE_EVENT, {}, async (data) => {
+      this._onWalletsData(data, WALLET_UPDATE_EVENT)
+    })
+
     this.ws2Manager.openWS()
 
     await onConnected
   }
 
+  _onWalletsData (data, type) {
+    if (type === WALLET_SNAPSHOT_EVENT) {
+      this.userWallets = data._collection
+        ? data._collection.map(wallet => {
+            return {
+              currency: wallet.currency,
+              type: wallet.type,
+              balance: wallet.balance,
+              balanceAvailable: wallet.balanceAvailable
+            }
+          })
+        : []
+    } else if (type === WALLET_UPDATE_EVENT && this.userWallets) {
+      this.userWallets.forEach(wallet => {
+        if (wallet.currency === data.currency && wallet.type === data.type) {
+          wallet.balance = data.balance || wallet.balance
+          wallet.balanceAvailable = data.balanceAvailable || wallet.balanceAvailable
+        }
+      })
+    }
+
+    // pass wallets data to strategy exeuctors
+    for (const [, strategy] of this.strategy.entries()) {
+      const { liveStrategyExecutor } = strategy
+      if (liveStrategyExecutor) {
+        liveStrategyExecutor.setWallets(this.userWallets)
+      }
+    }
+  }
+
   /**
    * @private
    */
@@ -315,6 +357,7 @@ class StrategyManager {
     })
 
     this._registerStrategyExecutionListeners(liveStrategyExecutor, strategyMapKey, strategyOptions)
+    liveStrategyExecutor.setWallets(this.userWallets)
 
     await liveStrategyExecutor.execute()
 
diff --git a/test/unit/lib/ws_servers/api/strategy/manager.js b/test/unit/lib/ws_servers/api/strategy/manager.js
index 5c191ebf..90e1878f 100644
--- a/test/unit/lib/ws_servers/api/strategy/manager.js
+++ b/test/unit/lib/ws_servers/api/strategy/manager.js
@@ -39,7 +39,8 @@ const StrategyExecutionStub = {
   execute: sandbox.stub(),
   stopExecution: sandbox.stub(),
   generateResults: sandbox.stub(),
-  removeAllListeners: sandbox.stub()
+  removeAllListeners: sandbox.stub(),
+  setWallets: sandbox.stub()
 }
 
 const PriceFeedStub = {