Merge pull request limitd#3 from pmalouin/reset-all-buckets

Add `resetAll()` method to clear all keys of all buckets
dafortune · Sep 20, 2018 · dd785be · dd785be
2 parents 71997f6 + c50b25e
commit dd785be
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 8 deletions.
diff --git a/lib/client.js b/lib/client.js
@@ -110,6 +110,10 @@ class LimitdRedis extends EventEmitter {
     this.put(type, key, opts, cb);
   }
 
+  resetAll(cb) {
+    this.db.resetAll(cb);
+  }
+
   close(callback) {
     this.db.close((err) => {
       this.db.removeAllListeners();

diff --git a/lib/db.js b/lib/db.js
@@ -1,6 +1,7 @@
 const ms = require('ms');
 const fs = require('fs');
 const _ = require('lodash');
+const async = require('async');
 const utils = require('./utils');
 const Redis = require('ioredis');
 const EventEmitter = require('events').EventEmitter;
@@ -258,6 +259,19 @@ class LimitDBRedis extends EventEmitter {
         });
       });
   }
+
+  /**
+   * Resets/re-fills all keys in all buckets.
+   * @param {function(Error)} [callback].
+   */
+  resetAll(callback) {
+    callback = callback || _.noop;
+
+    const dbs = this.redis.nodes ? this.redis.nodes('master') : [this.redis];
+    async.each(dbs, (db, cb) => {
+      db.flushdb(cb);
+    }, callback);
+  }
 }
 
 

diff --git a/test/client.tests.js b/test/client.tests.js
@@ -162,6 +162,13 @@ describe('LimitdRedis', () => {
     });
   });
 
+  describe('#resetAll', () => {
+    it('should call db.resetAll', (done) => {
+      client.db.resetAll = (cb) => cb();
+      client.resetAll(done);
+    });
+  });
+
   describe('#close', () => {
     it('should call db.close', (done) => {
       client.db.close = (cb) => cb();

diff --git a/test/db.tests.js b/test/db.tests.js
@@ -1,5 +1,4 @@
 /* eslint-env node, mocha */
-/*eslint no-useless-escape: "disable"*/
 const ms       = require('ms');
 const async    = require('async');
 const _        = require('lodash');
@@ -16,7 +15,7 @@ const buckets = {
         per_second: 100
       },
       'local-lan': {
-        match: '192\.168\.',
+        match: '192\\.168\\.',
         per_second: 50
       },
       '10.0.0.123': {
@@ -41,7 +40,7 @@ const buckets = {
     per_second: 5,
     overrides: {
       'regexp': {
-        match: '^regexp\|',
+        match: '^regexp',
         size: 10
       }
     }
@@ -56,10 +55,7 @@ describe('LimitDBRedis', () => {
     db = new LimitDB({ uri: 'localhost', buckets, prefix: 'tests:' });
     db.once('error', done);
     db.once('ready', () => {
-      const dbs = db.redis.nodes ? db.redis.nodes('master') : [db.redis];
-      async.each(dbs, (db, cb) => {
-        db.flushall(cb);
-      }, done);
+      db.resetAll(done);
     });
   });
 
@@ -401,7 +397,7 @@ describe('LimitDBRedis', () => {
     it('should add to the bucket', (done) => {
       db.take({ type: 'ip', key: '8.8.8.8', count: 5 }, (err) => {
         if (err) {
-          return cb(err);
+          return done(err);
         }
 
         db.put({ type: 'ip', key: '8.8.8.8', count: 4 }, (err, result) => {
@@ -539,4 +535,42 @@ describe('LimitDBRedis', () => {
       });
     });
   });
+
+  describe('#resetAll', function () {
+    it('should reset all keys of all buckets', (done) => {
+      async.parallel([
+        // Empty those buckets...
+        (cb) => db.take({ type: 'ip', key: '1.1.1.1', count: buckets.ip.size }, cb),
+        (cb) => db.take({ type: 'ip', key: '2.2.2.2', count: buckets.ip.size }, cb),
+        (cb) => db.take({ type: 'user', key: 'some_user', count: buckets.user.size }, cb)
+      ], (err) => {
+        if (err) {
+          return done(err);
+        }
+
+        db.resetAll((err) => {
+          if (err) {
+            return done(err);
+          }
+          async.parallel([
+            (cb) => db.take({ type: 'ip', key: '1.1.1.1' }, cb),
+            (cb) => db.take({ type: 'ip', key: '2.2.2.2' }, cb),
+            (cb) => db.take({ type: 'user', key: 'some_user' }, cb)
+          ], (err, results) => {
+            if (err) {
+              return done(err);
+            }
+
+            assert.equal(results[0].remaining, buckets.ip.size - 1);
+            assert.equal(results[0].conformant, true);
+            assert.equal(results[1].remaining, buckets.ip.size - 1);
+            assert.equal(results[0].conformant, true);
+            assert.equal(results[2].remaining, buckets.user.size - 1);
+            assert.equal(results[2].conformant, true);
+            done();
+          });
+        });
+      });
+    });
+  });
 });