cherry-pick tests refactoring

workers/limiter/package.json

@@ -10,7 +10,6 @@
     "workerType": "cron-tasks/limiter",
     "dependencies": {
         "@types/redis": "^2.8.28",
-        "axios": "^0.21.2",
-        "redis": "^3.1.1"
+        "axios": "^0.21.2"
     }
 }

workers/limiter/src/index.ts

@@ -81,6 +81,9 @@ export default class LimiterWorker extends Worker {
 
     this.projectsCollection = accountDbConnection.collection<ProjectDBScheme>('projects');
     this.workspacesCollection = accountDbConnection.collection<WorkspaceDBScheme>('workspaces');
+
+    await this.redis.initialize();
+
     await super.start();
   }
 
@@ -91,6 +94,7 @@ export default class LimiterWorker extends Worker {
     await super.finish();
     await this.eventsDb.close();
     await this.accountsDb.close();
+    await this.redis.close();
   }
 
   /**

workers/limiter/src/redisHelper.ts

@@ -1,5 +1,5 @@
 import HawkCatcher from '@hawk.so/nodejs';
-import redis from 'redis';
+import { createClient, RedisClientType } from 'redis';
 import createLogger from '../../../lib/logger';
 
 /**
@@ -9,7 +9,8 @@ export default class RedisHelper {
   /**
    * Redis client for making queries
    */
-  private readonly redisClient = redis.createClient({ url: process.env.REDIS_URL });
+  // private readonly redisClient = redis.createClient({ url: process.env.REDIS_URL });
+  private readonly redisClient: RedisClientType;
 
   /**
    * Logger instance
@@ -22,24 +23,66 @@ export default class RedisHelper {
    */
   private readonly redisDisabledProjectsKey = 'DisabledProjectsSet';
 
+  /**
+   * Constructor of the Redis helper class
+   * Initializes the Redis client and sets up error handling
+   */
+  constructor() {
+    this.redisClient = createClient({ url: process.env.REDIS_URL });
+
+    this.redisClient.on('error', (error) => {
+      this.logger.error(error);
+      HawkCatcher.send(error);
+    });
+  }
+
+  /**
+   * Connect to redis client
+   */
+  public async initialize(): Promise<void> {
+    await this.redisClient.connect();
+  }
+
+  /**
+   * Close redis client
+   */
+  public async close(): Promise<void> {
+    if (this.redisClient.isOpen) {
+      await this.redisClient.quit();
+    }
+  }
   /**
    * Saves banned project ids to redis
    * If there is no projects, then previous data in Redis will be erased
    *
    * @param projectIdsToBan - ids to ban
    */
   public saveBannedProjectsSet(projectIdsToBan: string[]): Promise<void> {
-    return new Promise((resolve, reject) => {
+    return new Promise(async (resolve, reject) => {
       const callback = this.createCallback(resolve, reject);
 
       if (projectIdsToBan.length) {
-        this.redisClient.multi()
-          .del(this.redisDisabledProjectsKey)
-          .sadd(this.redisDisabledProjectsKey, projectIdsToBan)
-          .exec(callback);
+        const pipeline = this.redisClient.multi();
+
+        pipeline.del(this.redisDisabledProjectsKey);
+
+        pipeline.sAdd(this.redisDisabledProjectsKey, projectIdsToBan);
+
+        try {
+          await pipeline.exec();
+          callback(null);
+        } catch (err) {
+          callback(err);
+        }
       } else {
-        this.redisClient.del(this.redisDisabledProjectsKey, callback);
-      }
+        this.redisClient.del(this.redisDisabledProjectsKey)
+          .then(() => {
+            callback(null);
+          })
+          .catch((err) => {
+            callback(err);
+          });
+      }      
     });
   }
 
@@ -53,7 +96,9 @@ export default class RedisHelper {
       const callback = this.createCallback(resolve, reject);
 
       if (projectIds.length) {
-        this.redisClient.sadd(this.redisDisabledProjectsKey, projectIds, callback);
+        this.redisClient.sAdd(this.redisDisabledProjectsKey, projectIds)
+          .then(() => callback(null))
+          .catch((err) => callback(err));
       } else {
         resolve();
       }
@@ -70,7 +115,9 @@ export default class RedisHelper {
       const callback = this.createCallback(resolve, reject);
 
       if (projectIds.length) {
-        this.redisClient.srem(this.redisDisabledProjectsKey, projectIds, callback);
+        this.redisClient.sRem(this.redisDisabledProjectsKey, projectIds)
+          .then(() => callback(null))
+          .catch((err) => callback(err));
       } else {
         resolve();
       }

workers/limiter/tests/index.test.ts

@@ -2,7 +2,7 @@ import { Collection, Db, MongoClient, ObjectId } from 'mongodb';
 import '../../../env-test';
 import { GroupedEventDBScheme, PlanDBScheme, ProjectDBScheme, WorkspaceDBScheme } from '@hawk.so/types';
 import LimiterWorker from '../src';
-import redis from 'redis';
+import { createClient } from 'redis';
 import { mockedPlans } from './plans.mock';
 import axios from 'axios';
 import { mocked } from 'ts-jest/utils';
@@ -128,7 +128,8 @@ describe('Limiter worker', () => {
     projectCollection = db.collection<ProjectDBScheme>('projects');
     workspaceCollection = db.collection<WorkspaceDBScheme>('workspaces');
     planCollection = db.collection('plans');
-    redisClient = redis.createClient({ url: process.env.REDIS_URL });
+    redisClient = createClient({ url: process.env.REDIS_URL });
+    await redisClient.connect();
 
     /**
      * Insert mocked plans for using in tests
@@ -225,11 +226,9 @@ describe('Limiter worker', () => {
        *
        * Gets all members of set with key 'DisabledProjectsSet' from Redis
        */
-      redisClient.smembers('DisabledProjectsSet', (err, result) => {
-        expect(err).toBeNull();
-        expect(result).toContain(project._id.toString());
-        done();
-      });
+      const result = await redisClient.sMembers('DisabledProjectsSet');
+      expect(result).toContain(project._id.toString());
+      done();
     });
 
     test('Should not ban project if it does not reach the limit', async (done) => {
@@ -265,15 +264,9 @@ describe('Limiter worker', () => {
        *
        * Gets all members of set with key 'DisabledProjectsSet' from Redis
        */
-      redisClient.smembers('DisabledProjectsSet', (err, result) => {
-        expect(err).toBeNull();
-
-        /**
-         * Redis shouldn't contain id of project 'Test project #2' from 'Test workspace #2'
-         */
-        expect(result).not.toContain(project._id.toString());
-        done();
-      });
+      const result = await redisClient.sMembers('DisabledProjectsSet');
+      expect(result).not.toContain(project._id.toString());
+      done();
     });
 
     test('Should send a report with collected data', async () => {
@@ -344,11 +337,10 @@ describe('Limiter worker', () => {
        *
        * Gets all members of set with key 'DisabledProjectsSet' from Redis
        */
-      redisClient.smembers('DisabledProjectsSet', (err, result) => {
-        expect(err).toBeNull();
-        expect(result).toContain(project._id.toString());
-        done();
-      });
+      const result = await redisClient.sMembers('DisabledProjectsSet');
+
+      expect(result).toContain(project._id.toString());
+      done();
     });
   });
 
@@ -405,15 +397,13 @@ describe('Limiter worker', () => {
       /**
        * Gets all members of set with key 'DisabledProjectsSet' from Redis
        */
-      redisClient.smembers('DisabledProjectsSet', (err, result) => {
-        expect(err).toBeNull();
-
-        /**
-         * Redis shouldn't contain id of project 'Test project #2' from 'Test workspace #2'
-         */
-        expect(result).not.toContain(project._id.toString());
-        done();
-      });
+      const result = await redisClient.sMembers('DisabledProjectsSet');
+
+      /**
+       * Redis shouldn't contain id of project 'Test project #2' from 'Test workspace #2'
+       */
+      expect(result).not.toContain(project._id.toString());
+      done();
     });
 
     test('Should block workspace if the number of events exceed the limit', async (done) => {
@@ -468,15 +458,10 @@ describe('Limiter worker', () => {
       /**
        * Gets all members of set with key 'DisabledProjectsSet' from Redis
        */
-      redisClient.smembers('DisabledProjectsSet', (err, result) => {
-        expect(err).toBeNull();
-
-        /**
-         * Redis shouldn't contain id of the mocked project
-         */
-        expect(result).toContain(project._id.toString());
-        done();
-      });
+
+      const result = await redisClient.sMembers('DisabledProjectsSet');
+      expect(result).toContain(project._id.toString());
+      done();
     });
 
     test('Should correctly work if projects count equals 0', async () => {
@@ -504,5 +489,6 @@ describe('Limiter worker', () => {
 
   afterAll(async () => {
     await connection.close();
+    await redisClient.quit();
   });
 });

workers/notifier/src/redisHelper.ts

@@ -0,0 +1,79 @@
+import { createClient, RedisClientType } from 'redis';
+import { Rule } from '../types/rule';
+import { NotifierEvent } from '../types/notifier-task';
+
+/**
+ * Class with helper functions for working with Redis
+ */
+export default class RedisHelper {
+  /**
+   * Redis client for making queries
+   */
+  private readonly redisClient: RedisClientType;
+
+  /**
+   * Constructor of the Redis helper class
+   * Initializes the Redis client and sets up error handling
+   */
+  constructor() {
+    this.redisClient = createClient({ url: process.env.REDIS_URL });
+
+    this.redisClient.on('error', (error) => {
+      console.error(error);
+    });
+  }
+
+  /**
+   * Connect to redis client
+   */
+  public async initialize(): Promise<void> {
+    await this.redisClient.connect();
+  }
+
+  /**
+   * Close redis client
+   */
+  public async close(): Promise<void> {
+    if (this.redisClient.isOpen) {
+      await this.redisClient.quit();
+    }
+  }
+
+  /**
+   * Method that updates the event count respectfully to the threshold reset period
+   *
+   * @param ruleId - id of the rule used as a part of structure key
+   * @param groupHash - event group hash used as a part of structure key
+   * @param thresholdPeriod - period of time used to reset the event count
+   * @returns {number} current event count
+   */
+  public async computeEventCountForPeriod(ruleId: string, groupHash: NotifierEvent['groupHash'], thresholdPeriod: Rule['eventThresholdPeriod']): Promise<number> {
+    const script = `
+    local key = KEYS[1]
+    local currentTimestamp = tonumber(ARGV[1])
+    local thresholdExpirationPeriod = tonumber(ARGV[2])
+
+    local startPeriodTimestamp = tonumber(redis.call("HGET", key, "timestamp"))
+
+    if ((startPeriodTimestamp == nil) or (currentTimestamp >= startPeriodTimestamp + thresholdExpirationPeriod)) then
+        redis.call("HSET", key, "timestamp", currentTimestamp)
+        redis.call("HSET", key, "eventsCount", 0)
+    end
+    
+    local newCounter = redis.call("HINCRBY", key, "eventsCount", 1)
+    return newCounter
+    `;
+
+    const key = `${ruleId}:${groupHash}:${thresholdPeriod}:times`;
+
+    const currentTimestamp = Date.now();
+
+    const currentEventCount = await this.redisClient.eval(script, {
+      keys: [ key ],
+      arguments: [currentTimestamp.toString(), (currentTimestamp + thresholdPeriod).toString()],
+    }) as number;
+
+    return (currentEventCount !== null) ? currentEventCount : 0;
+  }
+}
+