Refetch plans if workspace tariff is missing from cache

Introduces logic to refetch tariff plans from the database if a workspace references a plan not present in the cached list. Adds tests to verify that plans are refetched when needed and that an error is thrown if the plan is still missing after refetching.

Author: talyguryn

workers/paymaster/src/index.ts

@@ -53,10 +53,15 @@ export default class PaymasterWorker extends Worker {
    */
   private workspaces: Collection<WorkspaceDBScheme>;
 
+  /**
+   * Collection with plans
+   */
+  private plansCollection: Collection<PlanDBScheme>;
+
   /**
    * List of tariff plans
    */
-  private plans: PlanDBScheme[];
+  private plans: PlanDBScheme[] = [];
 
   /**
    * Check if today is a payday for passed timestamp
@@ -111,15 +116,53 @@ export default class PaymasterWorker extends Worker {
     const connection = await this.db.connect();
 
     this.workspaces = connection.collection('workspaces');
-    const plansCollection = connection.collection<PlanDBScheme>('plans');
+    this.plansCollection = connection.collection<PlanDBScheme>('plans');
 
-    this.plans = await plansCollection.find({}).toArray();
+    await this.fetchPlans();
+
+    await super.start();
+  }
+
+  /**
+   * Fetches tariff plans from database and keeps them cached
+   */
+  private async fetchPlans(): Promise<void> {
+    if (!this.plansCollection) {
+      throw new Error('Plans collection is not initialized');
+    }
+
+    this.plans = await this.plansCollection.find({}).toArray();
 
     if (this.plans.length === 0) {
       throw new Error('Please add tariff plans to the database');
     }
+  }
 
-    await super.start();
+  /**
+   * Finds plan by id from cached plans
+   */
+  private findPlanById(planId: WorkspaceDBScheme['tariffPlanId']): PlanDBScheme | undefined {
+    return this.plans.find((plan) => plan._id.toString() === planId.toString());
+  }
+
+  /**
+   * Returns workspace plan, refreshes cache when plan is missing
+   */
+  private async getWorkspacePlan(workspace: WorkspaceDBScheme): Promise<PlanDBScheme> {
+    let currentPlan = this.findPlanById(workspace.tariffPlanId);
+
+    if (currentPlan) {
+      return currentPlan;
+    }
+
+    await this.fetchPlans();
+    currentPlan = this.findPlanById(workspace.tariffPlanId);
+
+    if (!currentPlan) {
+      throw new Error(`[Paymaster] Tariff plan ${workspace.tariffPlanId.toString()} not found for workspace ${workspace._id.toString()} (${workspace.name})`);
+    }
+
+    return currentPlan;
   }
 
   /**
@@ -180,9 +223,7 @@ export default class PaymasterWorker extends Worker {
    */
   private async processWorkspaceSubscriptionCheck(workspace: WorkspaceDBScheme): Promise<[WorkspaceDBScheme, boolean]> {
     const date = new Date();
-    const currentPlan = this.plans.find(
-      (plan) => plan._id.toString() === workspace.tariffPlanId.toString()
-    );
+    const currentPlan = await this.getWorkspacePlan(workspace);
 
     /** Define readable values */
 

workers/paymaster/tests/index.test.ts

@@ -695,6 +695,103 @@ describe('PaymasterWorker', () => {
     MockDate.reset();
   });
 
+  test('Should refetch plans if workspace tariff appears after worker start', async () => {
+    /**
+     * Arrange
+     */
+    const currentDate = new Date('2005-12-22');
+    const cachedPlan = createPlanMock({
+      monthlyCharge: 100,
+      isDefault: true,
+    });
+
+    await tariffCollection.insertOne(cachedPlan);
+
+    const worker = new PaymasterWorker();
+
+    await worker.start();
+
+    const newPlan = createPlanMock({
+      monthlyCharge: 0,
+      isDefault: false,
+    });
+
+    await tariffCollection.insertOne(newPlan);
+
+    const workspace = createWorkspaceMock({
+      plan: newPlan,
+      subscriptionId: null,
+      lastChargeDate: new Date('2005-11-22'),
+      isBlocked: true,
+      billingPeriodEventsCount: 10,
+    });
+
+    await workspacesCollection.insertOne(workspace);
+
+    MockDate.set(currentDate);
+
+    /**
+     * Act
+     */
+    await expect(worker.handle(WORKSPACE_SUBSCRIPTION_CHECK)).resolves.not.toThrow();
+
+    /**
+     * Assert
+     */
+    const updatedWorkspace = await workspacesCollection.findOne({ _id: workspace._id });
+
+    expect(updatedWorkspace.lastChargeDate).toEqual(currentDate);
+    expect(updatedWorkspace.billingPeriodEventsCount).toEqual(0);
+    expect(updatedWorkspace.isBlocked).toEqual(false);
+
+    await worker.finish();
+    MockDate.reset();
+  });
+
+  test('Should throw an error when workspace plan is still missing after refetch', async () => {
+    /**
+     * Arrange
+     */
+    const currentDate = new Date('2005-12-22');
+    const cachedPlan = createPlanMock({
+      monthlyCharge: 100,
+      isDefault: true,
+    });
+
+    await tariffCollection.insertOne(cachedPlan);
+
+    const worker = new PaymasterWorker();
+
+    await worker.start();
+
+    const missingPlan = createPlanMock({
+      monthlyCharge: 50,
+      isDefault: false,
+    });
+
+    const workspace = createWorkspaceMock({
+      plan: missingPlan,
+      subscriptionId: null,
+      lastChargeDate: new Date('2005-11-22'),
+      isBlocked: false,
+      billingPeriodEventsCount: 10,
+    });
+
+    await workspacesCollection.insertOne(workspace);
+
+    MockDate.set(currentDate);
+
+    /**
+     * Act + Assert
+     */
+    await expect(worker.handle(WORKSPACE_SUBSCRIPTION_CHECK)).rejects.toThrow(
+      `[Paymaster] Tariff plan ${missingPlan._id.toString()} not found for workspace ${workspace._id.toString()} (${workspace.name})`
+    );
+
+    await worker.finish();
+    MockDate.reset();
+  });
+
   afterAll(async () => {
     await connection.close();
     MockDate.reset();