Add computed balance to account detail view

6 files changed, 20 insertions, 10 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({
diff --git a/ui/app/account-detail.js b/ui/app/account-detail.js
index 02089ecd0..90724dc3f 100644
--- a/ui/app/account-detail.js
+++ b/ui/app/account-detail.js
@@ -32,6 +32,7 @@ function mapStateToProps (state) {
     currentCurrency: state.metamask.currentCurrency,
     currentAccountTab: state.metamask.currentAccountTab,
     tokens: state.metamask.tokens,
+    computedBalances: state.metamask.computedBalances,
   }
 }
 
@@ -45,7 +46,7 @@ AccountDetailScreen.prototype.render = function () {
   var selected = props.address || Object.keys(props.accounts)[0]
   var checksumAddress = selected && ethUtil.toChecksumAddress(selected)
   var identity = props.identities[selected]
-  var account = props.accounts[selected]
+  var account = props.computedBalances[selected]
   const { network, conversionRate, currentCurrency } = props
 
   return (
@@ -180,7 +181,7 @@ AccountDetailScreen.prototype.render = function () {
         }, [
 
           h(EthBalance, {
-            value: account && account.balance,
+            value: account && account.ethBalance,
             conversionRate,
             currentCurrency,
             style: {
diff --git a/ui/app/conf-tx.js b/ui/app/conf-tx.js
index 8b93305eb..15fb9a59f 100644
--- a/ui/app/conf-tx.js
+++ b/ui/app/conf-tx.js
@@ -49,7 +49,6 @@ ConfirmTxScreen.prototype.render = function () {
   var txParams = txData.params || {}
   var isNotification = isPopupOrNotification() === 'notification'
 
-
   log.info(`rendering a combined ${unconfTxList.length} unconf msg & txs`)
   if (unconfTxList.length === 0) return h(Loading, { isLoading: true })