Merge pull request #418 from isocolsky/feat/cache

Optimize balance computation

lib/blockchainmonitor.js

@@ -92,7 +92,7 @@ BlockchainMonitor.prototype._handleTxId = function(data, processIt) {
 
   self.storage.fetchTxByHash(data.txid, function(err, txp) {
     if (err) {
-      log.error('Could not fetch tx the db');
+      log.error('Could not fetch tx from the db');
       return;
     }
     if (!txp || txp.status != 'accepted') return;
@@ -164,13 +164,25 @@ BlockchainMonitor.prototype._handleTxOuts = function(data) {
         },
         walletId: walletId,
       });
+      self._updateActiveAddresses(address, function() {
         self._storeAndBroadcastNotification(notification, next);
       });
+    });
   }, function(err) {
     return;
   });
 };
 
+BlockchainMonitor.prototype._updateActiveAddresses = function(address, cb) {
+  var self = this;
+
+  self.storage.storeActiveAddresses(address.walletId, address.address, function(err) {
+    if (err) {
+      log.warn('Could not update wallet cache', err);
+    }
+    return cb(err);
+  });
+};
 
 BlockchainMonitor.prototype._handleIncommingTx = function(data) {
   this._handleTxId(data);

lib/common/defaults.js

@@ -37,4 +37,7 @@ Defaults.FEE_LEVELS = [{
   defaultValue: 10000
 }];
 
+// Minimum nb of addresses a wallet must have to start using 2-step balance optimization
+Defaults.TWO_STEP_BALANCE_THRESHOLD = 100;
+
 module.exports = Defaults;

lib/expressapp.js

@@ -214,6 +214,7 @@ ExpressApp.prototype.start = function(opts, cb) {
     getServerWithAuth(req, res, function(server) {
       var opts = {};
       if (req.query.includeExtendedInfo == '1') opts.includeExtendedInfo = true;
+      if (req.query.twoStep == '1') opts.twoStep = true;
 
       server.getStatus(opts, function(err, status) {
         if (err) return returnError(err, res, req);
@@ -316,7 +317,9 @@ ExpressApp.prototype.start = function(opts, cb) {
 
   router.get('/v1/balance/', function(req, res) {
     getServerWithAuth(req, res, function(server) {
-      server.getBalance({}, function(err, balance) {
+      var opts = {};
+      if (req.query.twoStep == '1') opts.twoStep = true;
+      server.getBalance(opts, function(err, balance) {
         if (err) return returnError(err, res, req);
         res.json(balance);
       });

lib/server.js

@@ -264,6 +264,7 @@ WalletService.prototype.getWallet = function(opts, cb) {
 /**
  * Retrieves wallet status.
  * @param {Object} opts
+ * @param {Object} opts.twoStep[=false] - Optional: use 2-step balance computation for improved performance
  * @param {Object} opts.includeExtendedInfo - Include PKR info & address managers for wallet & copayers
  * @returns {Object} status
  */
@@ -297,7 +298,7 @@ WalletService.prototype.getStatus = function(opts, cb) {
       });
     },
     function(next) {
-      self.getBalance({}, function(err, balance) {
+      self.getBalance(opts, function(err, balance) {
         if (err) return next(err);
         status.balance = balance;
         next();
@@ -833,7 +834,7 @@ WalletService.prototype._getBlockchainExplorer = function(network) {
   return this.blockchainExplorer;
 };
 
-WalletService.prototype._getUtxosForAddresses = function(addresses, cb) {
+WalletService.prototype._getUtxos = function(addresses, cb) {
   var self = this;
 
   if (addresses.length == 0) return cb(null, []);
@@ -856,24 +857,36 @@ WalletService.prototype._getUtxosForAddresses = function(addresses, cb) {
   });
 };
 
-WalletService.prototype._getUtxosForCurrentWallet = function(cb) {
+WalletService.prototype._getUtxosForCurrentWallet = function(addresses, cb) {
   var self = this;
 
   function utxoKey(utxo) {
     return utxo.txid + '|' + utxo.vout
   };
 
-  // Get addresses for this wallet
-  self.storage.fetchAddresses(self.walletId, function(err, addresses) {
-    if (err) return cb(err);
+  async.waterfall([
+
+    function(next) {
+      if (_.isArray(addresses)) {
+        if (!_.isEmpty(addresses)) {
+          next(null, addresses);
+        } else {
+          next(null, []);
+        }
+      } else {
+        self.storage.fetchAddresses(self.walletId, next);
+      }
+    },
+    function(addresses, next) {
+      if (addresses.length == 0) return next(null, []);
 
       var addressStrs = _.pluck(addresses, 'address');
-    self._getUtxosForAddresses(addressStrs, function(err, utxos) {
-      if (err) return cb(err);
-      if (utxos.length == 0) return cb(null, []);
+      self._getUtxos(addressStrs, function(err, utxos) {
+        if (err) return next(err);
+        if (utxos.length == 0) return next(null, []);
 
         self.getPendingTxs({}, function(err, txps) {
-        if (err) return cb(err);
+          if (err) return next(err);
 
           var lockedInputs = _.map(_.flatten(_.pluck(txps, 'inputs')), utxoKey);
           var utxoIndex = _.indexBy(utxos, utxoKey);
@@ -890,10 +903,11 @@ WalletService.prototype._getUtxosForCurrentWallet = function(cb) {
             utxo.publicKeys = addressToPath[utxo.address].publicKeys;
           });
 
-        return cb(null, utxos);
-      });
+          return next(null, utxos);
         });
       });
+    },
+  ], cb);
 };
 
 /**
@@ -908,9 +922,9 @@ WalletService.prototype.getUtxos = function(opts, cb) {
   opts = opts || {};
 
   if (_.isUndefined(opts.addresses)) {
-    self._getUtxosForCurrentWallet(cb);
+    self._getUtxosForCurrentWallet(null, cb);
   } else {
-    self._getUtxosForAddresses(opts.addresses, cb);
+    self._getUtxos(opts.addresses, cb);
   }
 };
 
@@ -950,16 +964,12 @@ WalletService.prototype._computeBytesToSendMax = function(utxos, cb) {
   });
 };
 
-/**
- * Creates a new transaction proposal.
- * @param {Object} opts
- * @returns {Object} balance - Total amount & locked amount.
- */
-WalletService.prototype.getBalance = function(opts, cb) {
+WalletService.prototype._getBalanceFromAddresses = function(addresses, cb) {
   var self = this;
 
-  self.getUtxos({}, function(err, utxos) {
+  self._getUtxosForCurrentWallet(addresses, function(err, utxos) {
     if (err) return cb(err);
+
     var balance = self._totalizeUtxos(utxos);
 
     // Compute balance by address
@@ -980,14 +990,118 @@ WalletService.prototype.getBalance = function(opts, cb) {
 
     self._computeBytesToSendMax(utxos, function(err, size) {
       if (err) {
-        log.error('Could not compute fees needed to transfer max amount', err);
+        log.error('Could not compute size of send max transaction', err);
       }
-      balance.totalBytesToSendMax = size || 0;
+      balance.totalBytesToSendMax = _.isNumber(size) ? size : null;
       return cb(null, balance);
     });
   });
 };
 
+WalletService.prototype._getBalanceOneStep = function(opts, cb) {
+  var self = this;
+
+  self.storage.fetchAddresses(self.walletId, function(err, addresses) {
+    if (err) return cb(err);
+    self._getBalanceFromAddresses(addresses, function(err, balance) {
+      if (err) return cb(err);
+
+      // Update cache
+      async.series([
+
+        function(next) {
+          self.storage.cleanActiveAddresses(self.walletId, next);
+        },
+        function(next) {
+          var active = _.pluck(balance.byAddress, 'address')
+          self.storage.storeActiveAddresses(self.walletId, active, next);
+        },
+      ], function(err) {
+        if (err) {
+          log.warn('Could not update wallet cache', err);
+        }
+        return cb(null, balance);
+      });
+    });
+  });
+};
+
+
+WalletService.prototype._getActiveAddresses = function(cb) {
+  var self = this;
+
+  self.storage.fetchActiveAddresses(self.walletId, function(err, active) {
+    if (err) {
+      log.warn('Could not fetch active addresses from cache', err);
+      return cb();
+    }
+
+    if (!_.isArray(active)) return cb();
+
+    self.storage.fetchAddresses(self.walletId, function(err, allAddresses) {
+      if (err) return cb(err);
+
+      var now = Math.floor(Date.now() / 1000);
+      var recent = _.pluck(_.filter(allAddresses, function(address) {
+        return address.createdOn > (now - 24 * 3600);
+      }), 'address');
+
+      var result = _.union(active, recent);
+
+      var index = _.indexBy(allAddresses, 'address');
+      result = _.map(result, function(r) {
+        return index[r];
+      });
+      return cb(null, result);
+    });
+  });
+};
+
+/**
+ * Get wallet balance.
+ * @param {Object} opts
+ * @param {Boolean} opts.twoStep[=false] - Optional - Use 2 step balance computation for improved performance
+ * @returns {Object} balance - Total amount & locked amount.
+ */
+WalletService.prototype.getBalance = function(opts, cb) {
+  var self = this;
+
+  opts = opts || {};
+
+  if (!opts.twoStep)
+    return self._getBalanceOneStep(opts, cb);
+
+  self.storage.countAddresses(self.walletId, function(err, nbAddresses) {
+    if (err) return cb(err);
+    if (nbAddresses < Defaults.TWO_STEP_BALANCE_THRESHOLD) {
+      return self._getBalanceOneStep(opts, cb);
+    }
+
+    self._getActiveAddresses(function(err, activeAddresses) {
+      if (err) return cb(err);
+      if (!_.isArray(activeAddresses)) {
+        return self._getBalanceOneStep(opts, cb);
+      } else {
+        log.debug('Requesting partial balance for ' + activeAddresses.length + ' out of ' + nbAddresses + ' addresses');
+        self._getBalanceFromAddresses(activeAddresses, function(err, partialBalance) {
+          if (err) return cb(err);
+          cb(null, partialBalance);
+          setTimeout(function() {
+            self._getBalanceOneStep(opts, function(err, fullBalance) {
+              if (err) return;
+              if (!_.isEqual(partialBalance, fullBalance)) {
+                log.debug('Cache miss: balance in active addresses differs from final balance');
+                self._notify('BalanceUpdated', fullBalance);
+              }
+            });
+          }, 1);
+          return;
+        });
+      }
+    });
+  });
+};
+
 WalletService.prototype._sampleFeeLevels = function(network, points, cb) {
   var self = this;
 
@@ -1102,7 +1216,7 @@ WalletService.prototype._selectTxInputs = function(txp, utxosToExclude, cb) {
     return _.pluck(_.sortBy(list, 'order'), 'utxo');
   };
 
-  self.getUtxos({}, function(err, utxos) {
+  self._getUtxosForCurrentWallet(null, function(err, utxos) {
     if (err) return cb(err);
 
     var excludeIndex = _.reduce(utxosToExclude, function(res, val) {
@@ -1601,7 +1715,7 @@ WalletService.prototype._broadcastRawTx = function(network, raw, cb) {
   bc.broadcast(raw, function(err, txid) {
     if (err) return cb(err);
     return cb(null, txid);
-  })
+  });
 };
 
 /**

lib/storage.js

@@ -20,6 +20,7 @@ var collections = {
   COPAYERS_LOOKUP: 'copayers_lookup',
   PREFERENCES: 'preferences',
   EMAIL_QUEUE: 'email_queue',
+  CACHE: 'cache',
 };
 
 var Storage = function(opts) {
@@ -56,6 +57,11 @@ Storage.prototype._createIndexes = function() {
   this.db.collection(collections.EMAIL_QUEUE).createIndex({
     notificationId: 1,
   });
+  this.db.collection(collections.CACHE).createIndex({
+    walletId: 1,
+    type: 1,
+    key: 1,
+  });
 };
 
 Storage.prototype.connect = function(opts, cb) {
@@ -225,7 +231,7 @@ Storage.prototype.fetchLastTxs = function(walletId, creatorId, limit, cb) {
 Storage.prototype.fetchPendingTxs = function(walletId, cb) {
   var self = this;
 
-  this.db.collection(collections.TXS).find({
+  self.db.collection(collections.TXS).find({
     walletId: walletId,
     isPending: true,
   }).sort({
@@ -383,6 +389,12 @@ Storage.prototype.fetchAddresses = function(walletId, cb) {
   });
 };
 
+Storage.prototype.countAddresses = function(walletId, cb) {
+  this.db.collection(collections.ADDRESSES).find({
+    walletId: walletId,
+  }).count(cb);
+};
+
 Storage.prototype.storeAddress = function(address, cb) {
   var self = this;
 
@@ -501,6 +513,65 @@ Storage.prototype.fetchEmailByNotification = function(notificationId, cb) {
   });
 };
 
+Storage.prototype.cleanActiveAddresses = function(walletId, cb) {
+  var self = this;
+
+  async.series([
+
+    function(next) {
+      self.db.collection(collections.CACHE).remove({
+        walletId: walletId,
+        type: 'activeAddresses',
+      }, {
+        w: 1
+      }, next);
+    },
+    function(next) {
+      self.db.collection(collections.CACHE).insert({
+        walletId: walletId,
+        type: 'activeAddresses',
+        key: null
+      }, {
+        w: 1
+      }, next);
+    },
+  ], cb);
+};
+
+Storage.prototype.storeActiveAddresses = function(walletId, addresses, cb) {
+  var self = this;
+
+  async.each(addresses, function(address, next) {
+    var record = {
+      walletId: walletId,
+      type: 'activeAddresses',
+      key: address,
+    };
+    self.db.collection(collections.CACHE).update({
+      walletId: record.walletId,
+      type: record.type,
+      key: record.key,
+    }, record, {
+      w: 1,
+      upsert: true,
+    }, next);
+  }, cb);
+};
+
+Storage.prototype.fetchActiveAddresses = function(walletId, cb) {
+  var self = this;
+
+  self.db.collection(collections.CACHE).find({
+    walletId: walletId,
+    type: 'activeAddresses',
+  }).toArray(function(err, result) {
+    if (err) return cb(err);
+    if (_.isEmpty(result)) return cb();
+
+    return cb(null, _.compact(_.pluck(result, 'key')));
+  });
+};
+
 Storage.prototype._dump = function(cb, fn) {
   fn = fn || console.log;
   cb = cb || function() {};

test/expressapp.js

@@ -113,6 +113,44 @@ describe('ExpressApp', function() {
         });
       });
 
+      describe('Balance', function() {
+        it('should handle cache argument', function(done) {
+          var server = {
+            getBalance: sinon.stub().callsArgWith(1, null, {}),
+          };
+          var TestExpressApp = proxyquire('../lib/expressapp', {
+            './server': {
+              initialize: sinon.stub().callsArg(1),
+              getInstanceWithAuth: sinon.stub().callsArgWith(1, null, server),
+            }
+          });
+          start(TestExpressApp, function() {
+            var reqOpts = {
+              url: testHost + ':' + testPort + config.basePath + '/v1/balance',
+              headers: {
+                'x-identity': 'identity',
+                'x-signature': 'signature'
+              }
+            };
+            request(reqOpts, function(err, res, body) {
+              should.not.exist(err);
+              res.statusCode.should.equal(200);
+              var args = server.getBalance.getCalls()[0].args[0];
+              should.not.exist(args.twoStep);
+
+              reqOpts.url += '?twoStep=1';
+              request(reqOpts, function(err, res, body) {
+                should.not.exist(err);
+                res.statusCode.should.equal(200);
+                var args = server.getBalance.getCalls()[1].args[0];
+                args.twoStep.should.equal(true);
+                done();
+              });
+            });
+          });
+        });
+      });
+
       describe('/v1/notifications', function(done) {
         var server, TestExpressApp, clock;
         beforeEach(function() {

test/integration/helpers.js

@@ -211,12 +211,26 @@ helpers.toSatoshi = function(btc) {
   }
 };
 
-helpers.stubUtxos = function(server, wallet, amounts, cb) {
+helpers.stubUtxos = function(server, wallet, amounts, opts, cb) {
+  if (_.isFunction(opts)) {
+    cb = opts;
+    opts = {};
+  }
+  opts = opts || {};
+
+  if (!helpers._utxos) helpers._utxos = {};
+
+  async.waterfall([
+
+    function(next) {
+      if (opts.addresses) return next(null, [].concat(opts.addresses));
       async.mapSeries(_.range(0, amounts.length > 2 ? 2 : 1), function(i, next) {
         server.createAddress({}, next);
-  }, function(err, addresses) {
-    should.not.exist(err);
+      }, next);
+    },
+    function(addresses, next) {
       addresses.should.not.be.empty;
+
       var utxos = _.compact(_.map([].concat(amounts), function(amount, i) {
         var confirmations;
         if (_.isString(amount) && _.startsWith(amount, 'u')) {
@@ -249,14 +263,25 @@ helpers.stubUtxos = function(server, wallet, amounts, cb) {
           confirmations: confirmations
         };
       }));
+
+      if (opts.keepUtxos) {
+        helpers._utxos = helpers._utxos.concat(utxos);
+      } else {
+        helpers._utxos = utxos;
+      }
+
       blockchainExplorer.getUnspentUtxos = function(addresses, cb) {
-      var selected = _.filter(utxos, function(utxo) {
+        var selected = _.filter(helpers._utxos, function(utxo) {
           return _.contains(addresses, utxo.address);
         });
         return cb(null, selected);
       };
 
-    return cb(utxos);
+      return next();
+    },
+  ], function(err) {
+    should.not.exist(err);
+    return cb(helpers._utxos);
   });
 };
 
@@ -450,16 +475,16 @@ helpers.createProposalOpts = function(type, outputs, signingKey, moreOpts, input
   return opts;
 };
 helpers.createAddresses = function(server, wallet, main, change, cb) {
-  var clock = sinon.useFakeTimers(Date.now(), 'Date');
-  async.map(_.range(main + change), function(i, next) {
-    clock.tick(1000);
+  // var clock = sinon.useFakeTimers('Date');
+  async.mapSeries(_.range(main + change), function(i, next) {
+    // clock.tick(1000);
     var address = wallet.createAddress(i >= main);
     server.storage.storeAddressAndWallet(wallet, address, function(err) {
       next(err, address);
     });
   }, function(err, addresses) {
-    if (err) throw new Error('Could not generate addresses');
-    clock.restore();
+    should.not.exist(err);
+    // clock.restore();
     return cb(_.take(addresses, main), _.takeRight(addresses, change));
   });
 };

test/integration/server.js

@@ -1504,6 +1504,316 @@ describe('Wallet service', function() {
     });
   });
 
+  describe('#getBalance 2 steps', function() {
+    var server, wallet, clock;
+    var _threshold = Defaults.TWO_STEP_BALANCE_THRESHOLD;
+    beforeEach(function(done) {
+      clock = sinon.useFakeTimers(Date.now(), 'Date');
+      Defaults.TWO_STEP_BALANCE_THRESHOLD = 0;
+
+      helpers.createAndJoinWallet(1, 1, function(s, w) {
+        server = s;
+        wallet = w;
+        done();
+      });
+    });
+    afterEach(function() {
+      clock.restore();
+      Defaults.TWO_STEP_BALANCE_THRESHOLD = _threshold;
+    });
+
+    it('should get balance', function(done) {
+      helpers.stubUtxos(server, wallet, [1, 'u2', 3], function() {
+        server.getBalance({
+          twoStep: true
+        }, function(err, balance) {
+          should.not.exist(err);
+          should.exist(balance);
+          balance.totalAmount.should.equal(helpers.toSatoshi(6));
+          balance.lockedAmount.should.equal(0);
+          balance.availableAmount.should.equal(helpers.toSatoshi(6));
+          balance.totalBytesToSendMax.should.equal(578);
+
+          balance.totalConfirmedAmount.should.equal(helpers.toSatoshi(4));
+          balance.lockedConfirmedAmount.should.equal(0);
+          balance.availableConfirmedAmount.should.equal(helpers.toSatoshi(4));
+
+          should.exist(balance.byAddress);
+          balance.byAddress.length.should.equal(2);
+          balance.byAddress[0].amount.should.equal(helpers.toSatoshi(4));
+          balance.byAddress[1].amount.should.equal(helpers.toSatoshi(2));
+          setTimeout(done, 100);
+        });
+      });
+    });
+
+    it('should trigger notification when balance of non-prioritary addresses is updated', function(done) {
+      var oldAddrs, newAddrs;
+
+      async.series([
+
+        function(next) {
+          helpers.createAddresses(server, wallet, 2, 0, function(addrs) {
+            oldAddrs = addrs;
+            next();
+          });
+        },
+        function(next) {
+          clock.tick(7 * 24 * 3600 * 1000);
+          helpers.createAddresses(server, wallet, 2, 0, function(addrs) {
+            newAddrs = addrs;
+            server._getActiveAddresses(function(err, active) {
+              should.not.exist(err);
+              should.not.exist(active);
+              helpers.stubUtxos(server, wallet, [1, 2], {
+                addresses: [oldAddrs[0], newAddrs[0]],
+              }, function() {
+                next();
+              });
+            });
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3));
+            next();
+          });
+        },
+        function(next) {
+          setTimeout(next, 100);
+        },
+        function(next) {
+          server._getActiveAddresses(function(err, active) {
+            should.not.exist(err);
+            should.exist(active);
+            active.length.should.equal(3);
+            next();
+          });
+        },
+        function(next) {
+          helpers.stubUtxos(server, wallet, 0.5, {
+            addresses: oldAddrs[1],
+            keepUtxos: true,
+          }, function() {
+            next();
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3));
+            next();
+          });
+        },
+        function(next) {
+          setTimeout(next, 100);
+        },
+        function(next) {
+          server.getNotifications({}, function(err, notifications) {
+            should.not.exist(err);
+            var last = _.last(notifications);
+            last.type.should.equal('BalanceUpdated');
+            var balance = last.data;
+            balance.totalAmount.should.equal(helpers.toSatoshi(3.5));
+            next();
+          });
+        },
+      ], function(err) {
+        should.not.exist(err);
+        done();
+      });
+    });
+
+    it('should not trigger notification when only balance of prioritary addresses is updated', function(done) {
+      var oldAddrs, newAddrs;
+
+      async.series([
+
+        function(next) {
+          helpers.createAddresses(server, wallet, 2, 0, function(addrs) {
+            oldAddrs = addrs;
+            next();
+          });
+        },
+        function(next) {
+          clock.tick(7 * 24 * 3600 * 1000);
+          helpers.createAddresses(server, wallet, 2, 0, function(addrs) {
+            newAddrs = addrs;
+            helpers.stubUtxos(server, wallet, [1, 2], {
+              addresses: newAddrs,
+            }, function() {
+              next();
+            });
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3));
+            next();
+          });
+        },
+        function(next) {
+          setTimeout(next, 100);
+        },
+        function(next) {
+          helpers.stubUtxos(server, wallet, 0.5, {
+            addresses: newAddrs[0],
+            keepUtxos: true,
+          }, function() {
+            next();
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3.5));
+            next();
+          });
+        },
+        function(next) {
+          setTimeout(next, 100);
+        },
+        function(next) {
+          server.getNotifications({}, function(err, notifications) {
+            should.not.exist(err);
+            var last = _.last(notifications);
+            last.type.should.not.equal('BalanceUpdated');
+            next();
+          });
+        },
+      ], function(err) {
+        should.not.exist(err);
+        done();
+      });
+    });
+
+    it('should resolve balance of new addresses immediately', function(done) {
+      var addresses;
+
+      async.series([
+
+        function(next) {
+          helpers.createAddresses(server, wallet, 4, 0, function(addrs) {
+            addresses = addrs;
+            helpers.stubUtxos(server, wallet, [1, 2], {
+              addresses: _.take(addresses, 2),
+            }, function() {
+              next();
+            });
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3));
+            next();
+          });
+        },
+        function(next) {
+          server.createAddress({}, function(err, addr) {
+            helpers.stubUtxos(server, wallet, 0.5, {
+              addresses: addr,
+              keepUtxos: true,
+            }, function() {
+              next();
+            });
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3.5));
+            next();
+          });
+        },
+        function(next) {
+          setTimeout(next, 100);
+        },
+        function(next) {
+          server.getNotifications({}, function(err, notifications) {
+            should.not.exist(err);
+            var last = _.last(notifications);
+            last.type.should.not.equal('BalanceUpdated');
+            next();
+          });
+        },
+      ], function(err) {
+        should.not.exist(err);
+        done();
+      });
+    });
+
+    it('should not perform 2 steps when nb of addresses below threshold', function(done) {
+      var oldAddrs, newAddrs;
+      Defaults.TWO_STEP_BALANCE_THRESHOLD = 5;
+
+      async.series([
+
+        function(next) {
+          helpers.createAddresses(server, wallet, 2, 0, function(addrs) {
+            oldAddrs = addrs;
+            next();
+          });
+        },
+        function(next) {
+          clock.tick(7 * 24 * 3600 * 1000);
+          helpers.createAddresses(server, wallet, 2, 0, function(addrs) {
+            newAddrs = addrs;
+            helpers.stubUtxos(server, wallet, [1, 2], {
+              addresses: [oldAddrs[0], newAddrs[0]],
+            }, function() {
+              next();
+            });
+          });
+        },
+        function(next) {
+          server.getBalance({
+            twoStep: true
+          }, function(err, balance) {
+            should.not.exist(err);
+            should.exist(balance);
+            balance.totalAmount.should.equal(helpers.toSatoshi(3));
+            next();
+          });
+        },
+        function(next) {
+          setTimeout(next, 100);
+        },
+        function(next) {
+          server.getNotifications({}, function(err, notifications) {
+            should.not.exist(err);
+            var last = _.last(notifications);
+            last.type.should.not.equal('BalanceUpdated');
+            next();
+          });
+        },
+      ], function(err) {
+        should.not.exist(err);
+        done();
+      });
+    });
+  });
+
   describe('#getFeeLevels', function() {
     var server, wallet;
     beforeEach(function(done) {
@@ -2226,7 +2536,7 @@ describe('Wallet service', function() {
 
       it('should be able to create tx with inputs argument', function(done) {
         helpers.stubUtxos(server, wallet, [1, 3, 2], function(utxos) {
-          server._getUtxosForCurrentWallet(function(err, utxos) {
+          server.getUtxos({}, function(err, utxos) {
             should.not.exist(err);
             var inputs = [utxos[0], utxos[2]];
             var txOpts = helpers.createExternalProposalOpts('18PzpUFkFZE8zKWUPvfykkTxmB9oMR8qP7', 2.5,

test/storage.js

@@ -77,6 +77,7 @@ describe('Storage', function() {
       });
     });
   });
+
   describe('Copayer lookup', function() {
     it('should correctly store and fetch copayer lookup', function(done) {
       var wallet = Model.Wallet.create({