Got pending balance updating correctly

3 files changed, 69 insertions, 81 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)
   }
 
 }
diff --git a/app/scripts/controllers/balances.js b/app/scripts/controllers/balances.js
index b0b366628..cf3c8a757 100644
--- a/app/scripts/controllers/balances.js
+++ b/app/scripts/controllers/balances.js
@@ -11,97 +11,46 @@ class BalancesController {
     this.txController = txController
 
     const initState = extend({
-      balances: [],
+      computedBalances: {},
     }, opts.initState)
     this.store = new ObservableStore(initState)
 
     this._initBalanceUpdating()
   }
 
-  // PUBLIC METHODS
-
-  setSelectedAddress (_address) {
-    return new Promise((resolve, reject) => {
-      const address = normalizeAddress(_address)
-      this.store.updateState({ selectedAddress: address })
-      resolve()
-    })
-  }
-
-  getSelectedAddress (_address) {
-    return this.store.getState().selectedAddress
+  _initBalanceUpdating () {
+    const store = this.ethStore.getState()
+    this.addAnyAccountsFromStore(store)
+    this.ethStore.subscribe(this.addAnyAccountsFromStore.bind(this))
   }
 
-  addToken (rawAddress, symbol, decimals) {
-    const address = normalizeAddress(rawAddress)
-    const newEntry = { address, symbol, decimals }
-
-    const tokens = this.store.getState().tokens
-    const previousIndex = tokens.find((token, index) => {
-      return token.address === address
-    })
+  addAnyAccountsFromStore(store) {
+    const balances = store.accounts
 
-    if (previousIndex) {
-      tokens[previousIndex] = newEntry
-    } else {
-      tokens.push(newEntry)
+    for (let address in balances) {
+      this.trackAddressIfNotAlready(address)
     }
-
-    this.store.updateState({ tokens })
-    return Promise.resolve()
-  }
-
-  getTokens () {
-    return this.store.getState().tokens
-  }
-
-  updateFrequentRpcList (_url) {
-    return this.addToFrequentRpcList(_url)
-      .then((rpcList) => {
-        this.store.updateState({ frequentRpcList: rpcList })
-        return Promise.resolve()
-      })
   }
 
-  setCurrentAccountTab (currentAccountTab) {
-    return new Promise((resolve, reject) => {
-      this.store.updateState({ currentAccountTab })
-      resolve()
-    })
-  }
-
-  addToFrequentRpcList (_url) {
-    const rpcList = this.getFrequentRpcList()
-    const index = rpcList.findIndex((element) => { return element === _url })
-    if (index !== -1) {
-      rpcList.splice(index, 1)
-    }
-    if (_url !== 'http://localhost:8545') {
-      rpcList.push(_url)
+  trackAddressIfNotAlready (address) {
+    const state = this.store.getState()
+    if (!(address in state.computedBalances)) {
+      this.trackAddress(address)
     }
-    if (rpcList.length > 2) {
-      rpcList.shift()
-    }
-    return Promise.resolve(rpcList)
-  }
-
-  getFrequentRpcList () {
-    return this.store.getState().frequentRpcList
   }
-  //
-  // PRIVATE METHODS
-  //
-  _initBalanceUpdating () {
-    const store = this.ethStore.getState()
-    const balances = store.accounts
 
-    for (let address in balances) {
-      let updater = new BalancesController({
-        address,
-        ethQuery: this.ethQuery,
-        txController: this.txController,
-      })
-    }
+  trackAddress (address) {
+    let updater = new BalanceController({
+      address,
+      ethStore: this.ethStore,
+      txController: this.txController,
+    })
+    updater.store.subscribe((accountBalance) => {
+      let newState = this.store.getState()
+      newState.computedBalances[address] = accountBalance
+      this.store.updateState(newState)
+    })
+    updater.updateBalance()
   }
 }
 
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 81e31a556..f2df45947 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -20,6 +20,7 @@ const BlacklistController = require('./controllers/blacklist')
 const MessageManager = require('./lib/message-manager')
 const PersonalMessageManager = require('./lib/personal-message-manager')
 const TransactionController = require('./controllers/transactions')
+const BalancesController = require('./controllers/balances')
 const ConfigManager = require('./lib/config-manager')
 const nodeify = require('./lib/nodeify')
 const accountImporter = require('./account-import-strategies')
@@ -174,6 +175,7 @@ module.exports = class MetamaskController extends EventEmitter {
     this.networkController.store.subscribe(this.sendUpdate.bind(this))
     this.ethStore.subscribe(this.sendUpdate.bind(this))
     this.txController.memStore.subscribe(this.sendUpdate.bind(this))
+    this.balancesController.store.subscribe(this.sendUpdate.bind(this))
     this.messageManager.memStore.subscribe(this.sendUpdate.bind(this))
     this.personalMessageManager.memStore.subscribe(this.sendUpdate.bind(this))
     this.keyringController.memStore.subscribe(this.sendUpdate.bind(this))
@@ -248,6 +250,7 @@ module.exports = class MetamaskController extends EventEmitter {
     const wallet = this.configManager.getWallet()
     const vault = this.keyringController.store.getState().vault
     const isInitialized = (!!wallet || !!vault)
+
     return extend(
       {
         isInitialized,
@@ -258,6 +261,7 @@ module.exports = class MetamaskController extends EventEmitter {
       this.messageManager.memStore.getState(),
       this.personalMessageManager.memStore.getState(),
       this.keyringController.memStore.getState(),
+      this.balancesController.store.getState(),
       this.preferencesController.store.getState(),
       this.addressBookController.store.getState(),
       this.currencyController.store.getState(),