Got pending balance updating correctly

1 files changed, 25 insertions, 76 deletions

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()
   }
 }