From f2515970ae655b1d62267a5275ef0fcc6bfb9640 Mon Sep 17 00:00:00 2001 From: Chinmoy Chakraborty Date: Tue, 16 Jul 2024 16:22:21 +0530 Subject: [PATCH] Enable scheduling of bots. --- package.json | 2 ++ src/app.module.ts | 2 ++ src/modules/bot/bot.controller.spec.ts | 40 ++++++++++++++++++++++---- src/modules/bot/bot.controller.ts | 10 ++++++- src/modules/bot/bot.service.spec.ts | 37 ++++++++++++++++++++++++ src/modules/bot/bot.service.ts | 14 ++++++++- 6 files changed, 97 insertions(+), 8 deletions(-) diff --git a/package.json b/package.json index 64245d3..ccd23bc 100644 --- a/package.json +++ b/package.json @@ -39,6 +39,7 @@ "@nestjs/passport": "^8.2.1", "@nestjs/platform-express": "^8.0.0", "@nestjs/platform-fastify": "^8.2.6", + "@nestjs/schedule": "^4.1.0", "@nestjs/swagger": "^5.2.0", "@nestjs/terminus": "^9.2.2", "@prisma/client": "3", @@ -50,6 +51,7 @@ "cache-manager-redis-store": "^2.0.0", "class-transformer": "^0.5.1", "class-validator": "^0.13.2", + "cron": "^3.1.7", "expect-type": "^0.13.0", "fastify-compress": "3.7.0", "fastify-helmet": "^7.1.0", diff --git a/src/app.module.ts b/src/app.module.ts index ff298ca..134790f 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -27,6 +27,7 @@ import { HealthModule } from './health/health.module'; import { FusionAuthClientProvider } from './modules/user-segment/fusionauth/fusionauthClientProvider'; import { VaultClientProvider } from './modules/secrets/secrets.service.provider'; import { MonitoringModule } from './monitoring/monitoring.module'; +import { ScheduleModule } from '@nestjs/schedule'; import * as redisStore from 'cache-manager-redis-store'; @@ -58,6 +59,7 @@ import * as redisStore from 'cache-manager-redis-store'; max: 1000 }), MonitoringModule, + ScheduleModule.forRoot(), ], controllers: [AppController, ServiceController], providers: [ diff --git a/src/modules/bot/bot.controller.spec.ts b/src/modules/bot/bot.controller.spec.ts index 69a48f4..afd1022 100644 --- a/src/modules/bot/bot.controller.spec.ts +++ b/src/modules/bot/bot.controller.spec.ts @@ -14,6 +14,7 @@ import { BotStatus, Prisma } from '../../../prisma/generated/prisma-client-js'; import { FusionAuthClientProvider } from '../user-segment/fusionauth/fusionauthClientProvider'; import { BadRequestException, CacheModule, ServiceUnavailableException } from '@nestjs/common'; import { VaultClientProvider } from '../secrets/secrets.service.provider'; +import { SchedulerRegistry } from '@nestjs/schedule'; class MockPrismaService { bot = { @@ -94,6 +95,7 @@ const mockBotService = { getBroadcastReport: jest.fn(), start: jest.fn(), + scheduleNotification: jest.fn(), } const mockBotData: Prisma.BotGetPayload<{ @@ -254,7 +256,8 @@ describe('BotController', () => { BotService, { provide: BotService, useValue: mockBotService, - } + }, + SchedulerRegistry, ], }).compile(); @@ -262,21 +265,25 @@ describe('BotController', () => { configService = module.get(ConfigService); }); + afterEach(() => { + jest.clearAllMocks(); + }); + it('bot start returns bad request on non existent bot', async () => { - expect(botController.startOne('testBotIdNotExisting', {})).rejects.toThrowError(new BadRequestException('Bot does not exist')); + expect(botController.startOne('testBotIdNotExisting', '1', {})).rejects.toThrowError(new BadRequestException('Bot does not exist')); }); it('bot start returns bad request when bot does not have user data', async () => { - expect(botController.startOne('noUser', {})).rejects.toThrowError(new BadRequestException('Bot does not contain user segment data')); + expect(botController.startOne('noUser', '1', {})).rejects.toThrowError(new BadRequestException('Bot does not contain user segment data')); }); it('disabled bot returns unavailable error',async () => { - await expect(() => botController.startOne('disabled', {})).rejects.toThrowError(ServiceUnavailableException); + await expect(() => botController.startOne('disabled', '1', {})).rejects.toThrowError(ServiceUnavailableException); }); it('only disabled bot returns unavailable error',async () => { - expect(botController.startOne('pinned', {})).resolves; - expect(botController.startOne('enabled', {})).resolves; + expect(botController.startOne('pinned', '1', {})).resolves; + expect(botController.startOne('enabled', '1', {})).resolves; }); it('update only passes relevant bot data to bot service', async () => { @@ -367,4 +374,25 @@ describe('BotController', () => { expect(resp).toBeTruthy(); updateParametersPassed = []; }); + + it('bot start schedule for future time', async () => { + const futureTime = new Date(Date.now() + 100000).toUTCString(); + await botController.startOne( + 'enabled', + futureTime, + { 'conversation-authorization': 'testToken' } + ); + expect(mockBotService.scheduleNotification).toHaveBeenCalledTimes(1); + expect(mockBotService.start).toHaveBeenCalledTimes(0); + }); + + it('bot start triggers immediately when triggerTime is not passed', async () => { + await botController.startOne( + 'enabled', + undefined, + { 'conversation-authorization': 'testToken' } + ); + expect(mockBotService.scheduleNotification).toHaveBeenCalledTimes(0); + expect(mockBotService.start).toHaveBeenCalledTimes(1); + }); }); diff --git a/src/modules/bot/bot.controller.ts b/src/modules/bot/bot.controller.ts index eb5345f..b1ba7dc 100644 --- a/src/modules/bot/bot.controller.ts +++ b/src/modules/bot/bot.controller.ts @@ -200,7 +200,7 @@ export class BotController { AddOwnerInfoInterceptor, AddROToResponseInterceptor, ) - async startOne(@Param('id') id: string, @Headers() headers) { + async startOne(@Param('id') id: string, @Query('triggerTime') triggerTime: string | undefined, @Headers() headers) { const bot: Prisma.BotGetPayload<{ include: { users: { @@ -228,6 +228,14 @@ export class BotController { if (bot?.status == BotStatus.DISABLED) { throw new ServiceUnavailableException("Bot is not enabled!"); } + if (triggerTime) { + const currentTime = new Date(); + const scheduledTime = new Date(triggerTime); + if (scheduledTime.getTime() > currentTime.getTime()) { + await this.botService.scheduleNotification(id, scheduledTime, bot?.users[0].all?.config, headers['conversation-authorization']); + return; + } + } const res = await this.botService.start(id, bot?.users[0].all?.config, headers['conversation-authorization']); return res; } diff --git a/src/modules/bot/bot.service.spec.ts b/src/modules/bot/bot.service.spec.ts index e5d7f6f..e69b59d 100644 --- a/src/modules/bot/bot.service.spec.ts +++ b/src/modules/bot/bot.service.spec.ts @@ -1,3 +1,13 @@ +const MockCronJob = { + start: jest.fn(), +}; + +jest.mock('cron', () => { + return { + CronJob: jest.fn().mockImplementation(() => MockCronJob), + } +}); + import { Test, TestingModule } from '@nestjs/testing'; import { BotService } from './bot.service'; import { ConfigService } from '@nestjs/config'; @@ -11,6 +21,8 @@ import { BotStatus } from '../../../prisma/generated/prisma-client-js'; import { UserSegmentService } from '../user-segment/user-segment.service'; import { ConversationLogicService } from '../conversation-logic/conversation-logic.service'; import { assert } from 'console'; +import { SchedulerRegistry } from '@nestjs/schedule'; +import { CronJob } from 'cron'; const MockPrismaService = { @@ -146,6 +158,10 @@ const mockFile: Express.Multer.File = { }) }; +const MockSchedulerRegistry = { + addCronJob: jest.fn(), +} + const mockBotsDb = [{ "id": "testId", "createdAt": "2023-05-04T19:22:40.768Z", @@ -391,6 +407,10 @@ describe('BotService', () => { }, UserSegmentService, ConversationLogicService, + SchedulerRegistry, { + provide: SchedulerRegistry, + useValue: MockSchedulerRegistry, + }, ], }).compile(); @@ -835,4 +855,21 @@ describe('BotService', () => { } fetchMock.restore(); }); + + it('bot scheduling works as expected', async () => { + const futureDate = new Date(Date.now() + 100000); + jest.spyOn(MockSchedulerRegistry, 'addCronJob').mockImplementation((id: string, cron) => { + expect(id.startsWith('notification_')).toBe(true); + expect(cron).toStrictEqual(MockCronJob); + }); + await botService.scheduleNotification( + 'mockBotId', + futureDate, + { + 'myVar': 'myVal', + }, + 'mockToken', + ); + expect(MockCronJob.start).toHaveBeenCalledTimes(1); + }); }); diff --git a/src/modules/bot/bot.service.ts b/src/modules/bot/bot.service.ts index 38c3704..e7d4613 100644 --- a/src/modules/bot/bot.service.ts +++ b/src/modules/bot/bot.service.ts @@ -17,7 +17,9 @@ import { Cache } from 'cache-manager'; import { DeleteBotsDTO } from './dto/delete-bot-dto'; import { UserSegmentService } from '../user-segment/user-segment.service'; import { ConversationLogicService } from '../conversation-logic/conversation-logic.service'; -import { createHash } from 'crypto'; +import { createHash, randomUUID } from 'crypto'; +import { SchedulerRegistry } from '@nestjs/schedule'; +import { CronJob } from 'cron'; @Injectable() export class BotService { @@ -28,6 +30,7 @@ export class BotService { private configService: ConfigService, private userSegmentService: UserSegmentService, private conversationLogicService: ConversationLogicService, + private schedulerRegistry: SchedulerRegistry, //@ts-ignore @Inject(CACHE_MANAGER) public cacheManager: Cache, ) { @@ -152,6 +155,15 @@ export class BotService { }); } + // Example Trigger Time: '2021-03-21T00:00:00.000Z' (This is UTC time). + async scheduleNotification(id: string, scheduledTime: Date, config: any, token: string) { + const job = new CronJob(scheduledTime, () => { + this.start(id, config, token); + }); + this.schedulerRegistry.addCronJob(`notification_${randomUUID()}`, job); + job.start(); + } + // dateString = '2020-01-01' private getDateFromString(dateString: string) { return new Date(dateString);