From 86cd4e4fedbea9639de33827733b4b85ef988bee Mon Sep 17 00:00:00 2001
From: Dan Finlay <dan@danfinlay.com>
Date: Wed, 13 Sep 2017 14:20:19 -0700
Subject: Got pending balance updating correctly

---
 app/scripts/controllers/balance.js  |  45 ++++++++++++++--
 app/scripts/controllers/balances.js | 101 +++++++++---------------------------
 2 files changed, 65 insertions(+), 81 deletions(-)

(limited to 'app/scripts/controllers')

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