Draupnir/apps/draupnir/test/integration/throttleQueueTest.ts

// Copyright 2022 - 2024 Gnuxie <Gnuxie@protonmail.com>
// Copyright 2021 - 2022 The Matrix.org Foundation C.I.C.
//
// SPDX-License-Identifier: Apache-2.0
//
// SPDX-FileAttributionText: <text>
// This modified file incorporates work from mjolnir
// https://github.com/matrix-org/mjolnir
// </text>

import { Ok } from "@the-draupnir-project/matrix-protection-suite";
import { ThrottlingQueue } from "../../src/queues/ThrottlingQueue";
import { DraupnirTestContext } from "./mjolnirSetupUtils";

describe("Test: ThrottlingQueue", function () {
  it("Tasks enqueued with `push()` are executed exactly once and in the right order", async function (this: DraupnirTestContext) {
    this.timeout(20000);
    if (this.draupnir === undefined) {
      throw new TypeError("Test isn't setup correctly");
    }
    const queue = new ThrottlingQueue(this.draupnir.managementRoomOutput, 10);
    const state = new Map();
    const promises: Promise<void>[] = [];
    for (let counter = 0; counter < 10; ++counter) {
      const i = counter;
      promises.push(
        new Promise((resolve) => {
          queue.push(async () => {
            if (state.get(i)) {
              throw new Error(`We shouldn't have set state[${i}] yet`);
            }
            state.set(i, true);
            for (let j = 0; j < i; ++j) {
              if (!state.get(j)) {
                throw new Error(`We should have set state[${j}] already`);
              }
            }
            resolve();
            return Ok(undefined);
          });
        })
      );
    }
    await Promise.all(promises);
    for (let i = 0; i < 10; ++i) {
      if (!state.get(i)) {
        throw new Error(
          `This is the end of the test, we should have set state[${i}]`
        );
      }
    }

    // Give code a little bit more time to trip itself, in case `promises` are accidentally
    // resolved too early.
    await new Promise((resolve) => setTimeout(resolve, 1000));

    queue.dispose();
  });

  it("Tasks enqueued with `push()` are executed exactly once and in the right order, even if we call `block()` at some point", async function (this: DraupnirTestContext) {
    this.timeout(20000);
    if (this.draupnir === undefined) {
      throw new TypeError("Test isn't setup correctly");
    }
    const queue = new ThrottlingQueue(this.draupnir.managementRoomOutput, 10);
    const state = new Map();
    const promises: Promise<void>[] = [];
    for (let counter = 0; counter < 10; ++counter) {
      const i = counter;
      promises.push(
        new Promise((resolve) => {
          queue.push(async () => {
            if (state.get(i)) {
              throw new Error(`We shouldn't have set state[${i}] yet`);
            }
            state.set(i, true);
            for (let j = 0; j < i; ++j) {
              queue.block(100);
              if (!state.get(j)) {
                throw new Error(`We should have set state[${j}] already`);
              }
            }
            if (i % 2 === 0) {
              // Arbitrary call to `delay()`.
              queue.block(20);
            }
            resolve();
            return Ok(undefined);
          });
        })
      );
    }

    queue.block(100);

    await Promise.all(promises);
    for (let i = 0; i < 10; ++i) {
      if (!state.get(i)) {
        throw new Error(
          `This is the end of the test, we should have set state[${i}]`
        );
      }
    }

    // Give code a little bit more time to trip itself, in case `promises` are accidentally
    // resolved too early.
    await new Promise((resolve) => setTimeout(resolve, 1000));

    queue.dispose();
  });
});