Run marble tests on vitest's fake clock

`withTestScheduler` kept rxjs's virtual time and vitest's fake timers
apart. As a result, an rxjs timer advanced with the marbles while a
plain timer or `Date` did not. This meant a behavior could only be
marble tested while it was built from rxjs operators.

The `TestScheduler` now drives vitest's fake clock. Its frame is read
from the clock, so an rxjs action started from a raw timer's callback
is scheduled from when that callback ran, and its flush runs the queued
actions and the raw timers in time order, never letting the clock jump
past a queued action.

src/reactions/ReactionsReader.test.tsx

@@ -260,7 +260,6 @@ test("handles a reaction", () => {
   const reactionEventId = "$my_event_id:example.org";
   const reaction = ReactionSet[1];
 
-  vitest.useFakeTimers();
   vitest.setSystemTime(0);
 
   withTestScheduler(({ schedule, time, expectObservable }) => {
@@ -269,7 +268,7 @@ test("handles a reaction", () => {
         testScope(),
         rtcSession.asMockedSession(),
       );
-      schedule(`abc`, {
+      schedule(`ab`, {
         a: () => {},
         b: () => {
           rtcSession.room.emit(
@@ -293,9 +292,6 @@ test("handles a reaction", () => {
             {} as IRoomTimelineData,
           );
         },
-        c: () => {
-          vitest.advanceTimersByTime(REACTION_ACTIVE_TIME_MS);
-        },
       });
       expectObservable(reactions$).toBe(
         `ab ${REACTION_ACTIVE_TIME_MS - 1}ms c`,
@@ -304,7 +300,8 @@ test("handles a reaction", () => {
           b: {
             [`${localRtcMember.userId}:${localRtcMember.deviceId}`]: {
               reactionOption: reaction,
-              expireAfter: new Date(REACTION_ACTIVE_TIME_MS),
+              // The reaction arrives on frame 1
+              expireAfter: new Date(1 + REACTION_ACTIVE_TIME_MS),
             },
           },
           // Expect reaction to expire.
@@ -444,7 +441,6 @@ test("that reactions cannot be spammed", () => {
   const reactionA = ReactionSet[1];
   const reactionB = ReactionSet[2];
 
-  vitest.useFakeTimers();
   vitest.setSystemTime(0);
 
   withTestScheduler(({ schedule, expectObservable }) => {
@@ -453,7 +449,7 @@ test("that reactions cannot be spammed", () => {
         testScope(),
         rtcSession.asMockedSession(),
       );
-      schedule("abcd", {
+      schedule("abc", {
         a: () => {},
         b: () => {
           rtcSession.room.emit(
@@ -499,9 +495,6 @@ test("that reactions cannot be spammed", () => {
             {} as IRoomTimelineData,
           );
         },
-        d: () => {
-          vitest.advanceTimersByTime(REACTION_ACTIVE_TIME_MS);
-        },
       });
       expectObservable(reactions$).toBe(
         `ab- ${REACTION_ACTIVE_TIME_MS - 2}ms d`,
@@ -510,7 +503,8 @@ test("that reactions cannot be spammed", () => {
           b: {
             [`${localRtcMember.userId}:${localRtcMember.deviceId}`]: {
               reactionOption: reactionA,
-              expireAfter: new Date(REACTION_ACTIVE_TIME_MS),
+              // The reaction arrives on frame 1
+              expireAfter: new Date(1 + REACTION_ACTIVE_TIME_MS),
             },
           },
           d: {},

src/utils/marbles.test.ts

@@ -0,0 +1,180 @@
+/*
+Copyright 2026 Element Creations Ltd.
+
+SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
+Please see LICENSE in the repository root for full details.
+*/
+
+import { Observable, Subject, timer } from "rxjs";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+import { withTestScheduler } from "./test";
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+describe("withTestScheduler", () => {
+  it("times cold Observables from each subscription", () => {
+    withTestScheduler(({ cold, expectObservable, schedule }) => {
+      const source$ = cold("-a-b|");
+      expectObservable(source$).toBe("-a-b|");
+      const late: string[] = [];
+      schedule("--s", {
+        s: () => {
+          source$.subscribe((v) => late.push(v));
+        },
+      });
+      schedule("------e", {
+        e: () => {
+          expect(late).toEqual(["a", "b"]);
+        },
+      });
+    });
+  });
+
+  it("times hot Observables from the start of the run", () => {
+    withTestScheduler(({ hot, expectObservable, schedule }) => {
+      const source$ = hot("a-b-c");
+      // Frame 0 of a hot Observable is delivered once the run starts, after
+      // expectations registered during the continuation have subscribed
+      expectObservable(source$).toBe("a-b-c");
+      const late: string[] = [];
+      schedule("---s", {
+        s: () => {
+          source$.subscribe((v) => late.push(v));
+        },
+      });
+      schedule("-----e", {
+        e: () => {
+          expect(late).toEqual(["c"]);
+        },
+      });
+    });
+  });
+
+  it("honours subscription marbles", () => {
+    withTestScheduler(({ cold, expectObservable }) => {
+      expectObservable(cold("-a-b-c"), "^---!").toBe("-a-b");
+    });
+  });
+
+  it("records completion and errors", () => {
+    withTestScheduler(({ cold, expectObservable }) => {
+      expectObservable(cold("a|")).toBe("a|");
+      const boom = new Error("boom");
+      expectObservable(cold("-#", {}, boom)).toBe("-#", {}, boom);
+    });
+  });
+
+  it("subscribes expectations before actions scheduled on frame 0 run", () => {
+    withTestScheduler(({ schedule, expectObservable, behavior }) => {
+      const subject = new Subject<number>();
+      schedule("a-b", { a: () => subject.next(1), b: () => subject.next(2) });
+      expectObservable(subject).toBe("a-b", { a: 1, b: 2 });
+
+      const value$ = behavior("x-y", { x: 1, y: 2 });
+      expect(value$.value).toBe(1);
+      expectObservable(value$).toBe("x-y", { x: 1, y: 2 });
+    });
+  });
+
+  it("runs rxjs timers however long they are", () => {
+    withTestScheduler(({ expectObservable }) => {
+      expectObservable(timer(90_000)).toBe("90s (a|)", { a: 0 });
+    });
+  });
+
+  it("runs plain timers and Date on the same clock as the marbles", () => {
+    withTestScheduler(({ expectObservable, schedule }) => {
+      const start = Date.now();
+      const source$ = new Observable<string>((subscriber) => {
+        setTimeout(() => {
+          subscriber.next("a");
+          subscriber.complete();
+        }, 500);
+      });
+      expectObservable(source$).toBe("500ms (a|)");
+      schedule("--d", {
+        d: () => {
+          expect(Date.now() - start).toBe(2);
+        },
+      });
+    });
+  });
+
+  it("runs plain timers due on a frame before the rxjs actions due on it", () => {
+    withTestScheduler(({ schedule }) => {
+      const order: string[] = [];
+      setTimeout(() => order.push("plain"), 5);
+      schedule("-----r-e", {
+        r: () => {
+          order.push("rxjs");
+        },
+        e: () => {
+          expect(order).toEqual(["plain", "rxjs"]);
+        },
+      });
+    });
+  });
+
+  it("schedules an rxjs timer started from a plain timer at the right frame", () => {
+    withTestScheduler(({ expectObservable, hot }) => {
+      // Something rxjs is due later, so the clock must not run straight to it
+      expectObservable(hot("300ms z")).toBe("300ms z");
+      const subject = new Subject<number>();
+      setTimeout(() => timer(50).subscribe(subject), 100);
+      expectObservable(subject).toBe("150ms (a|)", { a: 0 });
+    });
+  });
+
+  it("does not run the clock on to an action that a plain timer took back", () => {
+    withTestScheduler(({ flush }) => {
+      const start = Date.now();
+      const subscription = timer(10).subscribe();
+      setTimeout(() => subscription.unsubscribe(), 5);
+      flush();
+      expect(Date.now() - start).toBe(5);
+    });
+  });
+
+  it("flushes part-way, and keeps counting frames from the start", () => {
+    withTestScheduler(({ hot, expectObservable, flush }) => {
+      const source$ = hot("-a-b");
+      const seen: string[] = [];
+      source$.subscribe((v) => seen.push(v));
+      flush();
+      expect(seen).toEqual(["a", "b"]);
+      expectObservable(hot("-c")).toBe("----c");
+    });
+  });
+
+  it("fails on a mismatch and restores real timers", () => {
+    expect(() =>
+      withTestScheduler(({ cold, expectObservable }) => {
+        expectObservable(cold("a")).toBe("b");
+      }),
+    ).toThrow();
+    expect(vi.isFakeTimers()).toBe(false);
+  });
+
+  it("keeps a date that was set without fake timers", () => {
+    vi.setSystemTime(1000);
+    withTestScheduler(({ schedule }) => {
+      schedule("--a", {
+        a: () => {
+          expect(Date.now()).toBe(1002);
+        },
+      });
+    });
+    expect(vi.isFakeTimers()).toBe(false);
+    expect(Date.now()).toBe(1000);
+  });
+
+  it("refuses a caller's fake timers that leave Date real", () => {
+    vi.useFakeTimers({ toFake: ["setTimeout"] });
+    expect(() => withTestScheduler(() => {})).toThrow(
+      "withTestScheduler needs Date faked along with the timers",
+    );
+  });
+});

src/utils/test.ts

@@ -128,12 +128,34 @@ export function testScope(): ObservableScope {
 }
 
 /**
- * Run Observables with a scheduler that virtualizes time, for testing purposes.
+ * Runs a test against virtual time, with marble diagrams describing when
+ * things happen. One frame is one millisecond. Timers and Date are faked for
+ * the duration of the run and advance with the marbles, so the code under test
+ * may time things with rxjs or with plain timers alike.
  */
 export function withTestScheduler(
   continuation: (helpers: OurRunHelpers) => void,
 ): void {
-  const scheduler = new TestScheduler((actual, expected) => {
+  const ownsClock = !vi.isFakeTimers();
+  // A date set without fake timers is a mock of its own, to give back after
+  const mockedDate = ownsClock ? vi.getMockedSystemTime() : null;
+  if (ownsClock) vi.useFakeTimers();
+  else if ((Date as { isFake?: boolean }).isFake !== true)
+    throw new Error("withTestScheduler needs Date faked along with the timers");
+  try {
+    runWithTestScheduler(continuation);
+  } finally {
+    if (ownsClock) {
+      vi.useRealTimers();
+      if (mockedDate !== null) vi.setSystemTime(mockedDate);
+    }
+  }
+}
+
+function runWithTestScheduler(
+  continuation: (helpers: OurRunHelpers) => void,
+): void {
+  const scheduler = new FakeClockTestScheduler((actual, expected) => {
     expect(actual).toStrictEqual(expected);
   });
   const scope = new ObservableScope();
@@ -186,6 +208,95 @@ export function withTestScheduler(
   scope.end();
 }
 
+// A run that keeps rescheduling would otherwise never end
+const maxStepsPerRun = 100_000;
+
+/**
+ * A TestScheduler whose virtual time is vitest's fake clock, so that timers
+ * set outside rxjs, and Date, advance in step with the marble diagrams. Raw
+ * timers due on a frame run before the rxjs actions due on that frame.
+ */
+class FakeClockTestScheduler extends TestScheduler {
+  public constructor(
+    assertDeepEqual: (actual: unknown, expected: unknown) => boolean | void,
+  ) {
+    super(assertDeepEqual);
+    const start = Date.now();
+    // Every read of the frame has to see the fake clock, so that an rxjs
+    // timer started from a raw timer's callback is scheduled relative to when
+    // that callback ran, not to the last rxjs action
+    Object.defineProperty(this, "frame", {
+      get: () => Date.now() - start,
+      set: () => {},
+    });
+  }
+
+  public override flush(): void {
+    // The inherited flush sets up the hot Observables, runs the actions and
+    // then checks the expectations. The actions have to take turns with raw
+    // timers in time order, so they are run here, and the inherited flush is
+    // left with nothing but the expectations to check.
+    while (this.hotObservables.length > 0) this.hotObservables.shift()!.setup();
+    this.runActions();
+    super.flush();
+  }
+
+  private runActions(): void {
+    const { actions, maxFrames } = this;
+    let steps = 0;
+    const step = (): void => {
+      if (++steps > maxStepsPerRun)
+        throw new Error(
+          `Still running after ${maxStepsPerRun} timer steps; a timer that keeps rescheduling has to be cleared before the run ends`,
+        );
+      vi.advanceTimersToNextTimer();
+    };
+    try {
+      for (;;) {
+        const action = actions[0];
+        if (action === undefined) {
+          if (vi.getTimerCount() === 0) break;
+          step();
+          continue;
+        }
+        if (action.delay > maxFrames) break;
+        if (!this.reach(action, step)) continue;
+        actions.shift();
+        const error: unknown = action.execute(action.state, action.delay);
+        if (error) throw error;
+      }
+    } catch (error) {
+      for (const action of actions.splice(0)) action.unsubscribe();
+      throw error;
+    }
+  }
+
+  /**
+   * Brings the clock up to the action's frame, running the raw timers due
+   * before or on it. False if one of them queued an action that is due sooner,
+   * or took this one back.
+   */
+  private reach(
+    action: (typeof this.actions)[number],
+    step: () => void,
+  ): boolean {
+    const wait = action.delay - this.frame;
+    if (vi.getTimerCount() === 0) {
+      if (wait > 0) vi.advanceTimersByTime(wait);
+      return true;
+    }
+    // The sentinel stops the clock from running past the frame
+    const sentinel = setTimeout(() => {}, wait);
+    try {
+      do step();
+      while (this.frame < action.delay && this.actions[0] === action);
+    } finally {
+      clearTimeout(sentinel);
+    }
+    return this.actions[0] === action;
+  }
+}
+
 interface EmitterMock<T> {
   on: (...args: unknown[]) => T;
   off: (...args: unknown[]) => T;