Implement set method for all limiters

animir · Jan 14, 2020 · 4fea264 · 4fea264
1 parent 9f7bed0
commit 4fea264
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 2 deletions.
diff --git a/lib/RateLimiterAbstract.js b/lib/RateLimiterAbstract.js
@@ -111,6 +111,10 @@ module.exports = class RateLimiterAbstract {
     throw new Error("You have to implement the method 'get'!");
   }
 
+  set() {
+    throw new Error("You have to implement the method 'set'!");
+  }
+
   block() {
     throw new Error("You have to implement the method 'block'!");
   }

diff --git a/lib/RateLimiterMemory.js b/lib/RateLimiterMemory.js
@@ -74,7 +74,18 @@ class RateLimiterMemory extends RateLimiterAbstract {
     const initPoints = this.points + 1;
 
     this._memoryStorage.set(this.getKey(key), initPoints, secDuration);
-    return Promise.resolve(new RateLimiterRes(0, msDuration, initPoints));
+    return Promise.resolve(
+      new RateLimiterRes(0, msDuration === 0 ? -1 : msDuration, initPoints)
+    );
+  }
+
+  set(key, points, secDuration) {
+    const msDuration = (secDuration >= 0 ? secDuration : this.duration) * 1000;
+
+    this._memoryStorage.set(this.getKey(key), points, secDuration);
+    return Promise.resolve(
+      new RateLimiterRes(0, msDuration === 0 ? -1 : msDuration, points)
+    );
   }
 
   get(key) {

diff --git a/lib/RateLimiterStoreAbstract.js b/lib/RateLimiterStoreAbstract.js
@@ -160,6 +160,21 @@ module.exports = class RateLimiterStoreAbstract extends RateLimiterAbstract {
     return this._block(this.getKey(key), this.points + 1, msDuration, options);
   }
 
+  /**
+   * Set points by key for any duration
+   *
+   * @param key
+   * @param points
+   * @param secDuration
+   * @param {Object} options
+   *
+   * @return Promise<RateLimiterRes>
+   */
+  set(key, points, secDuration, options = {}) {
+    const msDuration = (secDuration >= 0 ? secDuration : this.duration) * 1000;
+    return this._block(this.getKey(key), points, msDuration, options);
+  }
+
   /**
    *
    * @param key

diff --git a/lib/index.d.ts b/lib/index.d.ts
@@ -45,6 +45,8 @@ export class RateLimiterAbstract {
 
     get(key: string | number, options?: {[key: string]: any }): Promise<RateLimiterRes|null>;
 
+    set(key: string | number, points: number, secDuration: number, options?: {[key: string]: any }): Promise<RateLimiterRes>;
+
     delete(key: string | number, options?: {[key: string]: any }): Promise<boolean>;
 
     getKey(key: string | number): string;
@@ -140,6 +142,8 @@ export class RateLimiterMongo extends RateLimiterStoreAbstract {
 
     get(key: string | number, options?: IRateLimiterMongoFunctionOptions): Promise<RateLimiterRes|null>;
 
+    set(key: string | number, points: number, secDuration: number, options?: IRateLimiterMongoFunctionOptions): Promise<RateLimiterRes>;
+
     delete(key: string | number, options?: IRateLimiterMongoFunctionOptions): Promise<boolean>;
 }
 

diff --git a/test/RateLimiterMemory.test.js b/test/RateLimiterMemory.test.js
@@ -327,4 +327,39 @@ describe('RateLimiterMemory with fixed window', function RateLimiterMemoryTest()
         done(err);
       });
   });
+
+  it('set points by key', (done) => {
+    const testKey = 'set';
+    const rateLimiter = new RateLimiterMemory({ points: 10, duration: 1 });
+    rateLimiter.set(testKey, 12)
+      .then(() => {
+        rateLimiter.get(testKey)
+          .then((res) => {
+            expect(res.consumedPoints).to.equal(12);
+            done();
+          });
+      })
+      .catch((err) => {
+        done(err);
+      });
+  });
+
+  it('set points by key forever', (done) => {
+    const testKey = 'setforever';
+    const rateLimiter = new RateLimiterMemory({ points: 10, duration: 1 });
+    rateLimiter.set(testKey, 12, 0)
+      .then(() => {
+        setTimeout(() => {
+          rateLimiter.get(testKey)
+            .then((res) => {
+              expect(res.consumedPoints).to.equal(12);
+              expect(res.msBeforeNext).to.equal(-1);
+              done();
+            });
+        }, 1100);
+      })
+      .catch((err) => {
+        done(err);
+      });
+  });
 });
diff --git a/test/RateLimiterRedis.test.js b/test/RateLimiterRedis.test.js
@@ -764,7 +764,42 @@ describe('RateLimiterRedis with fixed window', function RateLimiterRedisTest() {
               expect(res.msBeforeNext).to.equal(-1);
               done();
             });
-        }, 1000);
+        }, 2000);
+      })
+      .catch((err) => {
+        done(err);
+      });
+  });
+
+  it('set points by key', (done) => {
+    const testKey = 'set';
+    const rateLimiter = new RateLimiterRedis({ storeClient: redisMockClient, points: 1, duration: 1 });
+    rateLimiter.set(testKey, 12)
+      .then(() => {
+        rateLimiter.get(testKey)
+          .then((res) => {
+            expect(res.consumedPoints).to.equal(12);
+            done();
+          });
+      })
+      .catch((err) => {
+        done(err);
+      });
+  });
+
+  it('set points by key forever', (done) => {
+    const testKey = 'setforever';
+    const rateLimiter = new RateLimiterRedis({ storeClient: redisMockClient, points: 1, duration: 1 });
+    rateLimiter.set(testKey, 12, 0)
+      .then(() => {
+        setTimeout(() => {
+          rateLimiter.get(testKey)
+            .then((res) => {
+              expect(res.consumedPoints).to.equal(12);
+              expect(res.msBeforeNext).to.equal(-1);
+              done();
+            });
+        }, 1100);
       })
       .catch((err) => {
         done(err);