element-call/src/lotus/lotusDenoise.test.ts

/*
Copyright 2026 Lotus Guild

SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
Please see LICENSE in the repository root for full details.
*/

import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
import { ParticipantEvent, Track } from "livekit-client";

import { type CallViewModel } from "../state/CallViewModel/CallViewModel";
import { startLotusDenoise } from "./lotusDenoise";
import { LotusWidgetActions } from "./lotusActions";

// Track constructed instances so tests can assert exactly one processor was
// built across racing apply() calls, and can assert the pending one is
// destroyed on early teardown. `vi.hoisted` is required because `vi.mock`
// factories are hoisted above this file's other top-level statements.
const instances = vi.hoisted(
  () =>
    [] as {
      config: { model: string };
      destroy: ReturnType<typeof vi.fn>;
      setMicMuted: ReturnType<typeof vi.fn>;
    }[],
);
// Ordered log of the heavy-lifting calls, to assert prepare runs before
// setProcessor (#7).
const calls = vi.hoisted(() => [] as string[]);
vi.mock("./lotusDenoiseProcessor", () => ({
  LotusDenoiseProcessor: class {
    public destroy = vi.fn().mockResolvedValue(undefined);
    public setMicMuted = vi.fn();
    public constructor(public config: { model: string }) {
      instances.push(this);
    }
  },
  prepareDenoiseAssets: vi.fn(async () => {
    calls.push("prepare");
    await Promise.resolve();
  }),
  releasePreparedDenoiseAssets: vi.fn().mockResolvedValue(undefined),
}));
const hostSend = vi.hoisted(() => vi.fn().mockResolvedValue({}));
vi.mock("../widget", () => ({
  widget: { api: { transport: { send: hostSend } } },
}));

/** A promise plus externally-callable resolve, for controlling ordering. */
function deferred<T>(): {
  promise: Promise<T>;
  resolve: (v: T) => void;
} {
  let resolve!: (v: T) => void;
  const promise = new Promise<T>((r) => {
    resolve = r;
  });
  return { promise, resolve };
}

function makeRoomAndVm(mic: {
  getProcessor: () => unknown;
  setProcessor: ReturnType<typeof vi.fn>;
  isMuted?: boolean;
}): {
  vm: CallViewModel;
  room: { localParticipant: Record<string, unknown> };
  firePublished: () => void;
  fire: (event: string, arg?: unknown) => void;
} {
  const handlers = new Map<string, (arg?: unknown) => void>();
  const localParticipant = {
    getTrackPublication: (
      source: Track.Source,
    ): { track: typeof mic } | undefined =>
      source === Track.Source.Microphone ? { track: mic } : undefined,
    on: (
      event: string,
      cb: (arg?: unknown) => void,
    ): Map<string, (arg?: unknown) => void> => handlers.set(event, cb),
    off: (event: string): boolean => handlers.delete(event),
  };
  const room = { localParticipant };
  const vm = {
    allConnections$: {
      subscribe: (
        cb: (data: {
          getConnections: () => { livekitRoom: unknown }[];
        }) => void,
      ) => {
        cb({ getConnections: () => [{ livekitRoom: room }] });
        return { unsubscribe: (): void => undefined };
      },
    },
  } as unknown as CallViewModel;
  return {
    vm,
    room,
    firePublished: () => handlers.get(ParticipantEvent.LocalTrackPublished)?.(),
    fire: (event, arg) => handlers.get(event)?.(arg),
  };
}

/** Let the async apply() chain (attach → retry → host notify) settle. */
async function flush(): Promise<void> {
  for (let i = 0; i < 8; i++) await Promise.resolve();
}

beforeEach(() => {
  instances.length = 0;
  calls.length = 0;
  hostSend.mockClear();
  // `lotusParam`/`lotusFlag` cache the URL params at first read; seed the hash
  // before startLotusDenoise() so the dedicated `lotusDenoiseSource` flag reads on.
  window.location.hash = "#/room?lotusDenoiseSource=1";
});

afterEach(() => {
  vi.resetModules();
});

describe("startLotusDenoise", () => {
  test("racing LocalTrackPublished events only construct one processor", async () => {
    const setProcessorDeferred = deferred<void>();
    let attached: unknown;
    const mic = {
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: unknown) => {
        calls.push("setProcessor");
        await setProcessorDeferred.promise;
        attached = p;
      }),
    };
    const { vm, firePublished } = makeRoomAndVm(mic);

    startLotusDenoise(vm);
    // [#7] The heavy assets were kicked off before setProcessor() was called.
    expect(calls).toEqual(["prepare", "setProcessor"]);
    // Simulate a second LocalTrackPublished (e.g. camera) firing before the
    // first setProcessor() has resolved.
    firePublished();

    expect(mic.setProcessor).toHaveBeenCalledTimes(1);
    expect(instances).toHaveLength(1);

    setProcessorDeferred.resolve();
    await Promise.resolve();
    await Promise.resolve();

    // Once attached, a further publish should be a no-op (mic.getProcessor()
    // is now set).
    firePublished();
    expect(mic.setProcessor).toHaveBeenCalledTimes(1);
  });

  test("destroys the in-flight processor if torn down before setProcessor resolves", async () => {
    const setProcessorDeferred = deferred<void>();
    const mic = {
      getProcessor: (): unknown => undefined,
      setProcessor: vi.fn(async () => setProcessorDeferred.promise),
    };
    const { vm } = makeRoomAndVm(mic);

    const teardown = startLotusDenoise(vm);
    expect(instances).toHaveLength(1);

    teardown();

    expect(instances[0].destroy).toHaveBeenCalledTimes(1);

    // Resolving afterwards must not throw/reject unhandled.
    setProcessorDeferred.resolve();
    await Promise.resolve();
  });

  test("reports the active model to the host once attached (#8)", async () => {
    let attached: unknown;
    const mic = {
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: unknown) => {
        await Promise.resolve();
        attached = p;
      }),
    };
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();
    expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, {
      active: true,
      model: "rnnoise",
    });
  });

  test("retries once with rnnoise when the requested model fails to init (#8)", async () => {
    window.location.hash =
      "#/room?lotusDenoiseSource=1&lotusModel=deepfilternet";
    let attached: unknown;
    const mic = {
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: { config: { model: string } }) => {
        await Promise.resolve();
        if (p.config.model === "deepfilternet")
          throw new Error("dfn assets missing");
        attached = p;
      }),
    };
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();

    expect(instances.map((i) => i.config.model)).toEqual([
      "deepfilternet",
      "rnnoise",
    ]);
    expect(mic.setProcessor).toHaveBeenCalledTimes(2);
    expect(hostSend).toHaveBeenCalledTimes(1);
    expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, {
      active: true,
      model: "rnnoise",
    });
  });

  test("notifies the host with the error when the rnnoise fallback also fails (#8)", async () => {
    window.location.hash = "#/room?lotusDenoiseSource=1&lotusModel=dtln";
    const mic = {
      getProcessor: (): unknown => undefined,
      setProcessor: vi.fn(async () => {
        await Promise.resolve();
        throw new Error("no worklet for you");
      }),
    };
    const { vm } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    await flush();

    expect(instances.map((i) => i.config.model)).toEqual(["dtln", "rnnoise"]);
    expect(hostSend).toHaveBeenCalledTimes(1);
    expect(hostSend).toHaveBeenCalledWith(LotusWidgetActions.DenoiseState, {
      active: false,
      error: "no worklet for you",
      model: "rnnoise",
    });
  });

  test("mirrors mic TrackMuted/TrackUnmuted onto the processor (#9)", async () => {
    let attached: unknown;
    const mic = {
      isMuted: true,
      getProcessor: (): unknown => attached,
      setProcessor: vi.fn(async (p: unknown) => {
        await Promise.resolve();
        attached = p;
      }),
    };
    const { vm, fire } = makeRoomAndVm(mic);
    startLotusDenoise(vm);
    // Seeded from the mic's current state before setProcessor().
    expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(true);
    await flush();

    const micPub = { source: Track.Source.Microphone };
    const camPub = { source: Track.Source.Camera };
    fire(ParticipantEvent.TrackUnmuted, micPub);
    expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(false);
    fire(ParticipantEvent.TrackMuted, micPub);
    expect(instances[0].setMicMuted).toHaveBeenLastCalledWith(true);

    // Camera mute must not touch the audio context.
    instances[0].setMicMuted.mockClear();
    fire(ParticipantEvent.TrackMuted, camPub);
    fire(ParticipantEvent.TrackUnmuted, camPub);
    expect(instances[0].setMicMuted).not.toHaveBeenCalled();
  });
});