cfw-autumn/server/tests/utils/advancedUsageUtils.ts

import { getPriceForOverage } from "@/internal/products/prices/priceUtils.js";
import assert from "assert";
import { expect } from "chai";
import { Decimal } from "decimal.js";
import { AutumnCli } from "tests/cli/AutumnCli.js";
import { creditSystems } from "tests/global.js";
import { timeout } from "./genUtils.js";
import { features } from "tests/global.js";
import { Feature } from "@autumn/shared";

const PRECISION = 10;
const CREDIT_MULTIPLIER = 100000;

export const getCreditsUsed = (
  creditSystem: Feature,
  meteredFeatureId: string,
  value: number,
) => {
  let schemaItem = creditSystem.config.schema.find(
    (item: any) => item.metered_feature_id === meteredFeatureId,
  );

  return new Decimal(value).mul(schemaItem.credit_amount).toNumber();
};

export const checkCreditBalance = async ({
  customerId,
  featureId,
  totalCreditsUsed,
  originalAllowance,
}: {
  customerId: string;
  featureId: string;
  totalCreditsUsed: number;
  originalAllowance: number;
}) => {
  // Check entitled
  const { allowed, balanceObj }: any = await AutumnCli.entitled(
    customerId,
    featureId,
    true,
  );

  try {
    assert.equal(allowed, true);
    assert.equal(
      balanceObj.balance,
      new Decimal(originalAllowance).minus(totalCreditsUsed).toNumber(),
    );
  } catch (error) {
    console.group();
    console.log("   - Total credits used: ", totalCreditsUsed);
    console.log("   - Original allowance: ", originalAllowance);
    console.log(
      "   - Expected balance: ",
      originalAllowance - totalCreditsUsed,
    );
    console.log("   - Actual balance: ", balanceObj.balance);
    console.groupEnd();
    throw error;
  }
};

export const checkUsageInvoiceAmount = async ({
  invoices,
  totalUsage,
  product,
  featureId,
  invoiceIndex,
  includeBase = true,
}: {
  invoices: any;
  totalUsage: number;
  product: any;
  featureId: string;
  invoiceIndex?: number;
  includeBase?: boolean;
}) => {
  const featureEntitlement: any = Object.values(product.entitlements).find(
    (entitlement: any) => entitlement.feature_id === featureId,
  );

  let meteredPrice = product.prices[product.prices.length - 1];
  let overage = new Decimal(totalUsage)
    .minus(featureEntitlement.allowance)
    .toNumber();
  const overagePrice = getPriceForOverage(meteredPrice, overage);

  let basePrice = 0;
  if (includeBase && product.prices.length > 1) {
    basePrice = product.prices[0].config.amount;
  }

  let totalPrice = new Decimal(overagePrice.toFixed(2))
    .plus(basePrice)
    .toNumber();

  try {
    for (let i = 0; i < invoices.length; i++) {
      let invoice = invoices[i];
      if (invoice.total == totalPrice) {
        invoiceIndex = i;
        assert.equal(invoice.product_ids[0], product.id);
        return;
      }
    }
    assert.fail("No invoice found with correct total price");
  } catch (error) {
    console.group();
    console.log("Check usage invoice amount failed");
    console.log("- Base price: ", basePrice);
    console.log("- Overage price: ", overagePrice);
    console.log(
      `Expected to find invoice with total of ${totalPrice} and product id ${product.id}`,
    );
    // console.log("Instead got: ", invoices[invoiceIndex || 0].total);
    console.log("Last 3 invoices", invoices.slice(-3));
    console.group();
    throw error;
  }
};

export const sendGPUEvents = async ({
  customerId,
  eventCount,
  groupObj = {},
}: {
  customerId: string;
  eventCount: number;
  groupObj?: any;
}) => {
  let totalCreditsUsed = 0;
  const batchEvents = [];
  for (let i = 0; i < eventCount; i++) {
    let randomVal = new Decimal(Math.random().toFixed(PRECISION))
      .mul(CREDIT_MULTIPLIER)
      .toNumber();
    let gpuId = i % 2 == 0 ? features.gpu1.id : features.gpu2.id;

    let creditsUsed = getCreditsUsed(
      creditSystems.gpuCredits,
      gpuId,
      randomVal,
    );

    totalCreditsUsed = new Decimal(totalCreditsUsed)
      .plus(creditsUsed)
      .toNumber();

    batchEvents.push(
      AutumnCli.sendEvent({
        customerId: customerId,
        eventName: gpuId,
        properties: { value: randomVal, ...groupObj },
      }),
    );
  }

  await Promise.all(batchEvents);
  await timeout(10000);

  return { creditsUsed: totalCreditsUsed };
};