element-call-Github/src/utils/test.ts

/*
Copyright 2023, 2024 New Vector Ltd.

SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
Please see LICENSE in the repository root for full details.
*/
import {
  combineLatest,
  map,
  NEVER,
  type Observable,
  of,
  type SchedulerLike,
} from "rxjs";
import { type RunHelpers, TestScheduler } from "rxjs/testing";
import {
  expect,
  type Mock,
  type MockedObject,
  type MockInstance,
  onTestFinished,
  vi,
  vitest,
} from "vitest";
import {
  EventType,
  MatrixEvent,
  type Room as MatrixRoom,
  type Room,
  type RoomMember,
  TypedEventEmitter,
} from "matrix-js-sdk";
import {
  CallMembership,
  type LivekitFocusSelection,
  type UnstableLivekitTransport,
  type MatrixRTCSession,
  MatrixRTCSessionEvent,
  type MatrixRTCSessionEventHandlerMap,
  MembershipManagerEvent,
  type SessionMembershipData,
  Status,
  type Transport,
} from "matrix-js-sdk/lib/matrixrtc";
import { type MembershipManagerEventHandlerMap } from "matrix-js-sdk/lib/matrixrtc/IMembershipManager";
import { type RemoteParticipant } from "livekit-client";
import { randomUUID } from "crypto";
import EventEmitter from "events";
import {
  type KeyTransportEvents,
  type KeyTransportEventsHandlerMap,
} from "matrix-js-sdk/lib/matrixrtc/IKeyTransport";

import { E2eeType } from "../e2ee/e2eeType";
import {
  DEFAULT_CONFIG,
  type ResolvedConfigOptions,
} from "../config/ConfigOptions";
import { Config } from "../config/Config";
import { type MediaDevices } from "../state/MediaDevices";
import {
  type AudioMediaTrack,
  type Behavior,
  constant,
  type DisconnectReason,
  type LocalMemberMedia,
  type LocalRTCMember,
  type MatrixRTCClient,
  type MatrixRTCError,
  type MemberMedia,
  type RemoteRTCMember,
  type TransportMetadata,
  type VideoMediaTrack,
} from "@element-hq/matrixrtc-sdk";
import { ObservableScope } from "../state/ObservableScope";
import { MuteStates } from "../state/MuteStates";
import { nullHostBridge } from "../HostBridge";
import {
  createLocalUserMedia,
  type LocalUserMediaViewModel,
} from "../state/media/LocalUserMediaViewModel";
import {
  createRemoteUserMedia,
  type RemoteUserMediaViewModel,
} from "../state/media/RemoteUserMediaViewModel";
import {
  createRemoteScreenShare,
  type RemoteScreenShareViewModel,
} from "../state/media/RemoteScreenShareViewModel";

export function withFakeTimers(continuation: () => void): void {
  vi.useFakeTimers();
  try {
    continuation();
  } finally {
    vi.useRealTimers();
  }
}

export async function flushPromises(): Promise<void> {
  await new Promise<void>((resolve) => window.setTimeout(resolve));
}

export interface OurRunHelpers extends RunHelpers {
  /**
   * Schedules a sequence of actions to happen, as described by a marble
   * diagram.
   */
  schedule: (marbles: string, actions: Record<string, () => void>) => void;
  behavior: <T>(
    marbles: string,
    values?: { [marble: string]: T },
    error?: unknown,
  ) => Behavior<T>;
  scope: ObservableScope;
}

interface TestRunnerGlobal {
  rxjsTestScheduler?: SchedulerLike;
}

/**
 * Create a new ObservableScope which ends when the current test ends.
 */
export function testScope(): ObservableScope {
  const scope = new ObservableScope();
  onTestFinished(() => scope.end());
  return scope;
}

/**
 * Run Observables with a scheduler that virtualizes time, for testing purposes.
 */
export function withTestScheduler(
  continuation: (helpers: OurRunHelpers) => void,
): void {
  const scheduler = new TestScheduler((actual, expected) => {
    expect(actual).toStrictEqual(expected);
  });
  const scope = new ObservableScope();
  // we set the test scheduler as a global so that you can watch it in a debugger
  // and get the frame number. e.g. `rxjsTestScheduler?.now()`
  (global as unknown as TestRunnerGlobal).rxjsTestScheduler = scheduler;
  scheduler.run((helpers) =>
    continuation({
      ...helpers,
      scope,
      schedule(marbles, actions) {
        const actionsObservable$ = helpers
          .cold(marbles)
          .pipe(map((value) => actions[value]()));
        const results = Object.fromEntries(
          Object.keys(actions).map((value) => [value, undefined] as const),
        );
        // Run the actions and verify that none of them error
        helpers.expectObservable(actionsObservable$).toBe(marbles, results);
      },
      behavior<T>(
        marbles: string,
        values?: { [marble: string]: T },
        error?: unknown,
      ) {
        // Generate a hot Observable with helpers.hot and use it as a Behavior.
        // To do this, we need to ensure that the initial value emits
        // synchronously upon subscription. The issue is that helpers.hot emits
        // frame 0 of the marble diagram *asynchronously*, only once we return
        // from the continuation, so we need to splice out the initial marble
        // and turn it into a proper initial value.
        const initialMarbleIndex = marbles.search(/[^ ]/);
        if (initialMarbleIndex === -1)
          throw new Error("Behavior must have an initial value");
        const initialMarble = marbles[initialMarbleIndex];
        const initialValue =
          values === undefined ? (initialMarble as T) : values[initialMarble];
        // The remainder of the marble diagram should start on frame 1
        return scope.behavior(
          helpers.hot(
            `-${marbles.slice(initialMarbleIndex + 1)}`,
            values,
            error,
          ),
          initialValue,
        );
      },
    }),
  );
  scope.end();
}

interface EmitterMock<T> {
  on: (...args: unknown[]) => T;
  off: (...args: unknown[]) => T;
  addListener: (...args: unknown[]) => T;
  removeListener: (...args: unknown[]) => T;
  emit: (event: string | symbol, ...args: unknown[]) => boolean;
}

export function mockEmitter<T>(): EmitterMock<T> {
  const ee = new EventEmitter();
  return {
    on: ee.on.bind(ee) as unknown as (...args: unknown[]) => T,
    off: ee.off.bind(ee) as unknown as (...args: unknown[]) => T,
    addListener: ee.addListener.bind(ee) as unknown as (
      ...args: unknown[]
    ) => T,
    removeListener: ee.removeListener.bind(ee) as unknown as (
      ...args: unknown[]
    ) => T,
    emit: ee.emit.bind(ee),
  };
}

export const exampleTransport: UnstableLivekitTransport = {
  type: "livekit",
  livekit_service_url: "https://lk.example.org",
};

export function mockRtcMembership(
  user: string | RoomMember,
  deviceId: string,
  customOverwrites?: {
    rtcBackendIdentity?: string;
    /** Every identity the member may have in LiveKit, when more than one. */
    backendIdentities?: string[];
    callId?: string;
    fociPreferred?: Transport[];
    focusActive?: LivekitFocusSelection;
    membership?: Partial<SessionMembershipData>;
  },
): CallMembership {
  // setup defaults based on overwrites and fallback values.
  const { rtcBackendIdentity, callId, fociPreferred, focusActive, membership } =
    {
      fociPreferred: [exampleTransport],
      focusActive: {
        type: "livekit" as const,
        focus_selection: "multi_sfu" as const,
      },
      callId: "",
      membership: {},
      ...customOverwrites,
    };

  const data: SessionMembershipData = {
    application: "m.call",
    call_id: callId,
    device_id: deviceId,
    foci_preferred: fociPreferred,
    focus_active: focusActive,
    ...membership,
  };
  const userId = typeof user === "string" ? user : user.userId;
  const event = new MatrixEvent({
    sender: userId,
    event_id: `$-ev-${randomUUID()}:example.org`,
    type: EventType.GroupCallMemberPrefix,
    content: data,
  });

  const backendIdentity = rtcBackendIdentity ?? `${userId}:${deviceId}`;
  const cms = new CallMembership(
    event,
    CallMembership.membershipDataFromMatrixEvent(event),
    customOverwrites?.backendIdentities ?? [backendIdentity],
    backendIdentity,
  );
  vi.mocked(cms).getTransport = vi.fn().mockReturnValue(fociPreferred[0]);

  return cms;
}

// Maybe it'd be good to move this to matrix-js-sdk? Our testing needs are
// rather simple, but if one util to mock a member is good enough for us, maybe
// it's useful for matrix-js-sdk consumers in general.
export function mockMatrixRoomMember(
  rtcMembership: CallMembership,
  member: Partial<RoomMember> = {},
): RoomMember {
  return {
    ...mockEmitter(),
    userId: rtcMembership.userId,
    getMxcAvatarUrl(): string | undefined {
      return undefined;
    },
    rawDisplayName: rtcMembership.userId,
    ...member,
  } as RoomMember;
}

export function mockMatrixRoom(room: Partial<MatrixRoom>): MatrixRoom {
  return { ...mockEmitter(), ...room } as Partial<MatrixRoom> as MatrixRoom;
}

/** A track as the SDK hands it out, with spies where a test may look. */
export type MockMediaTrack = Omit<AudioMediaTrack, "kind"> &
  Omit<VideoMediaTrack, "kind"> & { kind: "audio" | "video" };

export function mockMediaTrack(
  track: Partial<MockMediaTrack> = {},
): MockMediaTrack {
  return {
    source: "camera",
    kind: "video",
    id: "track",
    muted$: constant(false),
    encrypted$: constant(true),
    stats$: constant(undefined),
    attach: vi.fn(),
    detach: vi.fn(),
    setAudioContext: vi.fn(),
    setVolume: vi.fn(),
    facingMode$: constant(undefined),
    ...track,
  };
}

export const mockVideoTrack = (
  track: Partial<MockMediaTrack> = {},
): VideoMediaTrack =>
  mockMediaTrack({ kind: "video", ...track }) as VideoMediaTrack;

export const mockAudioTrack = (
  track: Partial<MockMediaTrack> = {},
): AudioMediaTrack =>
  mockMediaTrack({
    kind: "audio",
    source: "microphone",
    ...track,
  }) as AudioMediaTrack;

/** A member's media as the SDK hands it out, with nothing published unless given. */
export function mockMemberMedia(
  media: Partial<Omit<LocalMemberMedia, "local">> & { local?: boolean } = {},
): LocalMemberMedia {
  return {
    local: false,
    speaking$: constant(false),
    screenShareEnabled$: constant(false),
    microphone$: constant(undefined),
    camera$: constant(undefined),
    screenShare$: constant(undefined),
    screenShareAudio$: constant(undefined),
    encryptionError$: NEVER,
    switchCamera: vi.fn().mockResolvedValue(undefined),
    ...media,
  } as LocalMemberMedia;
}

/** Whoever publishes media, known by identity only. */
export function mockRemoteParticipant(
  participant: Partial<RemoteParticipant>,
): RemoteParticipant {
  return {
    isLocal: false,
    ...mockEmitter(),
    ...participant,
  } as RemoteParticipant;
}

export function mockLocalMedia(
  rtcMember: CallMembership,
  roomMember: Partial<RoomMember>,
  media: LocalMemberMedia | null,
  mediaDevices: MediaDevices,
): LocalUserMediaViewModel {
  const member = mockMatrixRoomMember(rtcMember, roomMember);
  return createLocalUserMedia(testScope(), {
    id: "local",
    userId: member.userId,
    rtcBackendIdentity: rtcMember.rtcBackendIdentity,
    media$: constant(media),
    encryptionSystem: { kind: E2eeType.PER_PARTICIPANT },
    focusUrl$: constant("https://rtc-example.org"),
    mediaDevices,
    displayName$: constant(member.rawDisplayName ?? "nodisplayname"),
    mxcAvatarUrl$: constant(member.getMxcAvatarUrl()),
    handRaised$: constant(null),
    reaction$: constant(null),
  });
}

export function mockRemoteMedia(
  rtcMember: CallMembership,
  roomMember: Partial<RoomMember>,
  media: MemberMedia | null,
  { focusUrl }: { focusUrl?: string } = { focusUrl: "https://rtc-example.org" },
): RemoteUserMediaViewModel {
  const member = mockMatrixRoomMember(rtcMember, roomMember);
  return createRemoteUserMedia(testScope(), {
    id: "remote",
    userId: member.userId,
    rtcBackendIdentity: rtcMember.rtcBackendIdentity,
    media$: constant(media),
    encryptionSystem: { kind: E2eeType.PER_PARTICIPANT },
    focusUrl$: constant(focusUrl),
    pretendToBeDisconnected$: constant(false),
    displayName$: constant(member.rawDisplayName ?? "nodisplayname"),
    mxcAvatarUrl$: constant(member.getMxcAvatarUrl()),
    handRaised$: constant(null),
    reaction$: constant(null),
  });
}

export function mockRemoteScreenShare(
  rtcMember: CallMembership,
  roomMember: Partial<RoomMember>,
  media: MemberMedia | null,
): RemoteScreenShareViewModel {
  const member = mockMatrixRoomMember(rtcMember, roomMember);
  return createRemoteScreenShare(testScope(), {
    id: "screenshare",
    userId: member.userId,
    media$: constant(media),
    encryptionSystem: { kind: E2eeType.PER_PARTICIPANT },
    focusUrl$: constant("https://rtc-example.org"),
    pretendToBeDisconnected$: constant(false),
    displayName$: constant(member.rawDisplayName ?? "nodisplayname"),
    mxcAvatarUrl$: constant(member.getMxcAvatarUrl()),
  });
}

export interface MockMemberInputs {
  membership: CallMembership;
  roomMember?: Partial<RoomMember>;
  media$?: Behavior<LocalMemberMedia | null>;
  transportUrl?: string | undefined;
}

/** A member as the client hands it out. */
export function mockRTCMember(
  local: true,
  inputs: MockMemberInputs,
): LocalRTCMember;
export function mockRTCMember(
  local: false,
  inputs: MockMemberInputs,
): RemoteRTCMember;
export function mockRTCMember(
  local: boolean,
  { membership, roomMember, media$, transportUrl }: MockMemberInputs,
): LocalRTCMember | RemoteRTCMember {
  const member = mockMatrixRoomMember(membership, roomMember);
  const transport: TransportMetadata | undefined =
    transportUrl === undefined
      ? undefined
      : {
          type: "livekit",
          id: transportUrl,
          raw: { type: "livekit", livekit_service_url: transportUrl },
          resolved$: constant(undefined),
        };
  const base = {
    id: membership.rtcBackendIdentity,
    userId: membership.userId,
    deviceId: membership.deviceId,
    membership$: constant(membership),
    displayName$: constant(member.rawDisplayName ?? membership.userId),
    avatarUrl$: constant(member.getMxcAvatarUrl()),
    transport$: constant(transport),
    media$: media$ ?? constant(null),
  };
  if (!local) return { ...base, local: false };
  return {
    ...base,
    local: true,
    sharingScreen$: constant(false),
    toggleScreenSharing: null,
    screenShareError$: constant(null),
    dismissScreenShareError: vi.fn(),
  };
}

export interface MockClientInputs {
  localMember$?: Behavior<LocalRTCMember | null>;
  remoteMembers$?: Behavior<RemoteRTCMember[]>;
  connected$?: Behavior<boolean>;
  reconnecting$?: Behavior<boolean>;
  disconnectReason$?: Behavior<DisconnectReason | null>;
  fatalError$?: Behavior<MatrixRTCError | null>;
  keyRotationSuppressed$?: Behavior<boolean>;
  connectedTransports$?: Behavior<TransportMetadata[]>;
}

/** A client that does nothing but hold the state a test hands it. */
export function mockMatrixRTCClient(
  scope: ObservableScope,
  {
    localMember$ = constant(null),
    remoteMembers$ = constant([]),
    connected$ = constant(true),
    reconnecting$ = constant(false),
    disconnectReason$ = constant(null),
    fatalError$ = constant(null),
    keyRotationSuppressed$ = constant(false),
    connectedTransports$ = constant([]),
  }: MockClientInputs = {},
): MatrixRTCClient {
  return {
    join: vi.fn(),
    leave: vi.fn(),
    status$: scope.behavior(
      connected$.pipe(map((c) => (c ? "connected" : "connecting"))),
    ),
    connected$,
    reconnecting$,
    disconnectReason$,
    setMicrophoneEnabled: vi.fn(async (enabled: boolean) =>
      Promise.resolve(enabled),
    ),
    setCameraEnabled: vi.fn(async (enabled: boolean) =>
      Promise.resolve(enabled),
    ),
    fatalError$,
    localMember$,
    remoteMembers$,
    memberCount$: scope.behavior(
      combineLatest(
        [localMember$, remoteMembers$],
        (local, remote) => remote.length + (local === null ? 0 : 1),
      ),
    ),
    keyRotationSuppressed$,
    connectedTransports$,
    sendData: vi.fn(async () => Promise.resolve()),
    data$: NEVER,
  };
}

export function mockConfig(
  config: Partial<ResolvedConfigOptions> = {},
): MockInstance<() => ResolvedConfigOptions> {
  const spy = vi.spyOn(Config, "get").mockReturnValue({
    ...DEFAULT_CONFIG,
    ...config,
    default_server_config: {
      ["m.homeserver"]: {
        base_url: "http://localhost:8008",
        server_name: "localhost",
      },
    },
  });
  // simulate loading the config
  vi.spyOn(Config, "init").mockResolvedValue(void 0);
  return spy;
}

export class MockRTCSession extends TypedEventEmitter<
  MatrixRTCSessionEvent | MembershipManagerEvent | KeyTransportEvents,
  KeyTransportEventsHandlerMap &
    MatrixRTCSessionEventHandlerMap &
    MembershipManagerEventHandlerMap
> {
  public asMockedSession(): MockedObject<MatrixRTCSession> {
    const session = this as unknown as MockedObject<MatrixRTCSession>;

    session.reemitEncryptionKeys = vi
      .fn<() => void>()
      .mockReturnValue(undefined);

    return session;
  }

  public readonly statistics = {
    counters: {},
  };

  public leave: ReturnType<typeof vitest.fn> = vitest
    .fn()
    .mockResolvedValue(undefined);

  public constructor(
    public readonly room: Room,
    public memberships: CallMembership[] = [],
  ) {
    super();
  }

  public joined = true;
  public isJoined(): boolean {
    return this.joined;
  }

  public isKeyRotationSuppressed = false;

  public withMemberships(
    rtcMembers$: Behavior<Partial<CallMembership>[]>,
  ): MockRTCSession {
    rtcMembers$.subscribe((m) => {
      const old = this.memberships;
      this.memberships = m as CallMembership[];
      this.emit(
        MatrixRTCSessionEvent.MembershipsChanged,
        old,
        this.memberships,
      );
    });

    return this;
  }

  public updateCallIntent: ReturnType<typeof vitest.fn> = vitest
    .fn()
    .mockImplementation(async () => Promise.resolve());

  private _membershipStatus = Status.Connected;
  public get membershipStatus(): Status {
    return this._membershipStatus;
  }
  public set membershipStatus(value: Status) {
    const prev = this._membershipStatus;
    this._membershipStatus = value;
    if (value !== prev)
      this.emit(MembershipManagerEvent.StatusChanged, prev, value);
  }

  private _probablyLeft = false;
  public get probablyLeft(): boolean {
    return this._probablyLeft;
  }
  public set probablyLeft(value: boolean) {
    const prev = this._probablyLeft;
    this._probablyLeft = value;
    if (value !== prev) this.emit(MembershipManagerEvent.ProbablyLeft, value);
  }

  public async join(): Promise<void> {
    return Promise.resolve();
  }
}

export const deviceStub = {
  available$: of(new Map<never, never>()),
  selected$: of(undefined),
  select(): void {},
};

export function mockMediaDevices(data: Partial<MediaDevices>): MediaDevices {
  return {
    audioInput: deviceStub,
    audioOutput: deviceStub,
    videoInput: deviceStub,
    ...data,
  } as MediaDevices;
}

export function mockMuteStates(
  joined$: Observable<boolean> = of(true),
): MuteStates {
  const observableScope = new ObservableScope();
  return new MuteStates(
    observableScope,
    mockMediaDevices({}),
    { audioEnabled: false, videoEnabled: false },
    nullHostBridge,
  );
}

export interface StubbedCapture {
  getUserMedia: Mock;
  /** Grants the microphone, as the browser does once permission is given. */
  grant: () => void;
  track: { stop: Mock };
  contexts: { close: Mock }[];
  /** Runs the pending animation frames, in order. */
  drawFrames: (count: number) => void;
  /** Sets the microphone's loudness from the next frame, 0 to 1. */
  speak: (amplitude: number) => void;
}

/**
 * Stubs the capture and Web Audio APIs, with the grant and the animation
 * frames driven by the test. Undo with {@link restoreAudioCapture}.
 */
export function stubAudioCapture(): StubbedCapture {
  const track = { stop: vi.fn() };
  const contexts: { close: Mock }[] = [];
  let grant = (): void => {};
  const granted = new Promise<MediaStream>((resolve) => {
    grant = (): void =>
      resolve({ getTracks: () => [track] } as unknown as MediaStream);
  });

  const frames: FrameRequestCallback[] = [];
  let amplitude = 0;
  vi.stubGlobal("requestAnimationFrame", (callback: FrameRequestCallback) => {
    frames.push(callback);
    return frames.length;
  });
  vi.stubGlobal("cancelAnimationFrame", vi.fn());
  vi.stubGlobal(
    "AudioContext",
    class {
      public readonly state = "running";
      public readonly close = vi.fn();
      public constructor() {
        contexts.push(this);
      }
      public createAnalyser(): object {
        return {
          fftSize: 1024,
          getByteTimeDomainData: (samples: Uint8Array): void => {
            // Silence is the midpoint of the range; a zeroed buffer reads as full scale.
            if (amplitude <= 0) {
              samples.fill(128);
              return;
            }
            const peak = Math.round(Math.min(1, amplitude) * 127);
            for (let i = 0; i < samples.length; i++)
              samples[i] = 128 + (i % 2 ? peak : -peak);
          },
        };
      }
      public createMediaStreamSource(): object {
        return { connect: (): void => {} };
      }
    },
  );
  // Only this property: replacing navigator loses getters like userAgent.
  Object.defineProperty(navigator, "mediaDevices", {
    configurable: true,
    value: { getUserMedia: vi.fn().mockReturnValue(granted) },
  });

  return {
    getUserMedia: navigator.mediaDevices.getUserMedia as unknown as Mock,
    grant: () => grant(),
    track,
    contexts,
    drawFrames: (count): void => {
      for (let i = 0; i < count; i++) frames.shift()?.(i);
    },
    speak: (next): void => {
      amplitude = next;
    },
  };
}

const realMediaDevices = Object.getOwnPropertyDescriptor(
  navigator,
  "mediaDevices",
);

/** Undoes {@link stubAudioCapture}. */
export function restoreAudioCapture(): void {
  vi.unstubAllGlobals();
  if (realMediaDevices === undefined) {
    Reflect.deleteProperty(navigator, "mediaDevices");
  } else {
    Object.defineProperty(navigator, "mediaDevices", realMediaDevices);
  }
}