Got pending balance updating correctly

1 files changed, 40 insertions, 5 deletions

diff --git a/app/scripts/controllers/balance.js b/app/scripts/controllers/balance.js
index 5dfe266e3..0d4ab7d4f 100644
--- a/app/scripts/controllers/balance.js
+++ b/app/scripts/controllers/balance.js
@@ -2,12 +2,14 @@ const ObservableStore = require('obs-store')
 const normalizeAddress = require('eth-sig-util').normalize
 const extend = require('xtend')
 const PendingBalanceCalculator = require('../lib/pending-balance-calculator')
+const BN = require('ethereumjs-util').BN
 
 class BalanceController {
 
   constructor (opts = {}) {
-    const { address, ethQuery, txController } = opts
-    this.ethQuery = ethQuery
+    const { address, ethStore, txController } = opts
+    this.address = address
+    this.ethStore = ethStore
     this.txController = txController
 
     const initState = extend({
@@ -17,10 +19,43 @@ class BalanceController {
 
     const { getBalance, getPendingTransactions } = opts
     this.balanceCalc = new PendingBalanceCalculator({
-      getBalance,
-      getPendingTransactions,
+      getBalance: () => Promise.resolve(this._getBalance()),
+      getPendingTransactions: this._getPendingTransactions.bind(this),
     })
-    this.updateBalance()
+
+    this.registerUpdates()
+  }
+
+  async updateBalance () {
+    const balance = await this.balanceCalc.getBalance()
+    this.store.updateState({
+      ethBalance: balance,
+    })
+  }
+
+  registerUpdates () {
+    const update = this.updateBalance.bind(this)
+    this.txController.on('submitted', update)
+    this.txController.on('confirmed', update)
+    this.txController.on('failed', update)
+    this.txController.blockTracker.on('block', update)
+  }
+
+  _getBalance () {
+    const store = this.ethStore.getState()
+    const balances = store.accounts
+    const entry = balances[this.address]
+    const balance = entry.balance
+    return balance ? new BN(balance.substring(2), 16) : new BN(0)
+  }
+
+  _getPendingTransactions () {
+    const pending = this.txController.getFilteredTxList({
+      from: this.address,
+      status: 'submitted',
+      err: undefined,
+    })
+    return Promise.resolve(pending)
   }
 
 }