Enable scheduling of bots. (#244)

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",

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: [

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,29 +256,34 @@ describe('BotController', () => {
         BotService, {
           provide: BotService,
           useValue: mockBotService,
-        }
+        },
+        SchedulerRegistry,
       ],
     }).compile();
 
     botController = module.get<BotController>(BotController);
     configService = module.get<ConfigService>(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);
+  });
 });

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;
   }

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);
+  });
 });

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);