Add computed balance to account detail view

4 files changed, 17 insertions, 7 deletions

diff --git a/app/scripts/controllers/balance.js b/app/scripts/controllers/balance.js
index 0d4ab7d4f..b4e72e751 100644
--- a/app/scripts/controllers/balance.js
+++ b/app/scripts/controllers/balance.js
@@ -1,6 +1,4 @@
 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
 
@@ -12,12 +10,11 @@ class BalanceController {
     this.ethStore = ethStore
     this.txController = txController
 
-    const initState = extend({
+    const initState = {
       ethBalance: undefined,
-    }, opts.initState)
+    }
     this.store = new ObservableStore(initState)
 
-    const { getBalance, getPendingTransactions } = opts
     this.balanceCalc = new PendingBalanceCalculator({
       getBalance: () => Promise.resolve(this._getBalance()),
       getPendingTransactions: this._getPendingTransactions.bind(this),
@@ -46,7 +43,7 @@ class BalanceController {
     const balances = store.accounts
     const entry = balances[this.address]
     const balance = entry.balance
-    return balance ? new BN(balance.substring(2), 16) : new BN(0)
+    return balance ? new BN(balance.substring(2), 16) : undefined
   }
 
   _getPendingTransactions () {
diff --git a/app/scripts/controllers/balances.js b/app/scripts/controllers/balances.js
index cf3c8a757..89c2ca95d 100644
--- a/app/scripts/controllers/balances.js
+++ b/app/scripts/controllers/balances.js
@@ -1,5 +1,4 @@
 const ObservableStore = require('obs-store')
-const normalizeAddress = require('eth-sig-util').normalize
 const extend = require('xtend')
 const BalanceController = require('./balance')
 
@@ -14,10 +13,17 @@ class BalancesController {
       computedBalances: {},
     }, opts.initState)
     this.store = new ObservableStore(initState)
+    this.balances = {}
 
     this._initBalanceUpdating()
   }
 
+  updateAllBalances () {
+    for (let address in this.balances) {
+      this.balances[address].updateBalance()
+    }
+  }
+
   _initBalanceUpdating () {
     const store = this.ethStore.getState()
     this.addAnyAccountsFromStore(store)
@@ -50,6 +56,7 @@ class BalancesController {
       newState.computedBalances[address] = accountBalance
       this.store.updateState(newState)
     })
+    this.balances[address] = updater
     updater.updateBalance()
   }
 }
diff --git a/app/scripts/lib/pending-balance-calculator.js b/app/scripts/lib/pending-balance-calculator.js
index 474ed3261..c66bffbbb 100644
--- a/app/scripts/lib/pending-balance-calculator.js
+++ b/app/scripts/lib/pending-balance-calculator.js
@@ -22,6 +22,8 @@ class PendingBalanceCalculator {
     const balance = results[0]
     const pending = results[1]
 
+    if (!balance) return undefined
+
     const pendingValue = pending.reduce((total, tx) => {
       return total.add(this.valueFor(tx))
     }, new BN(0))
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index f2df45947..02c06ead2 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -121,6 +121,10 @@ module.exports = class MetamaskController extends EventEmitter {
       ethStore: this.ethStore,
       txController: this.txController,
     })
+    this.networkController.on('networkDidChange', () => {
+      this.balancesController.updateAllBalances()
+    })
+    this.balancesController.updateAllBalances()
 
     // notices
     this.noticeController = new NoticeController({