Testing Async Observables with TestScheduler

Hello and welcome! In our last lesson, we successfully fortified our EventBus, making it robust and memory-safe by ensuring listeners are automatically cleaned up on unsubscription. With our core observable and event bus implementations now stable, we face a new and exciting challenge: testing complex, time-dependent asynchronous behavior.

So far, we've touched upon using Vitest's fakeTimers to test simple delays. However, as we prepare to implement more sophisticated time-based operators like debounceTime and throttleTime, orchestrating and asserting complex sequences of events becomes difficult and unwieldy.

This brings us to today's learning outcome: Implement a TestScheduler to deterministically test asynchronous observable chains using virtual time.

We will build the foundational engine that allows us to test asynchronous code in a perfectly synchronous, predictable, and instantaneous way. This powerful concept, pioneered by libraries like RxJS, is the gold standard for testing reactive systems.

1. The Problem with Time in Tests

As an experienced developer, you know that tests involving real-time delays (setTimeout) are slow, non-deterministic, and fragile. While tools like vi.useFakeTimers() improve the situation by allowing you to manipulate the clock (vi.advanceTimersByTime()), they can still be cumbersome for complex reactive streams.

Imagine testing a debounceTime(100) operator. You'd have to:

1. Emit a value.
2. Advance the clock by 99ms.
3. Assert that nothing has been emitted.
4. Advance the clock by 1ms.
5. Assert that the value has now been emitted.
6. Repeat this for multiple, overlapping emissions.

This imperative style is verbose and hard to read. The solution is to create a scheduler that owns a "virtual" clock and runs tests in a more declarative way.

To understand this concept better, let's start with the official RxJS documentation on the subject.

As you've read, the TestScheduler provides a run method that creates a sandboxed environment. Inside this environment, time-based operators automatically use the scheduler's virtual clock instead of setTimeout. This allows us to describe timed event streams and assert against them without any actual waiting.

2. Designing Our TestScheduler

Inspired by the RxJS implementation, we will build the core components of our own TestScheduler. For today, we will focus on the engine that manages virtual time and scheduled tasks. In the next lesson, we will build the parser for the "marble diagram" string syntax.

Our TestScheduler needs three main components:

1. A Virtual Clock: A simple counter, which we'll call frame, to track the passage of virtual time. By convention, one frame equals one virtual millisecond.
2. An Action Queue: An array of tasks to be performed. Each task will have a frame at which it should execute and an action function to run.
3. A flush() Method: The engine that processes the action queue. It will sort tasks by their scheduled frame, advance the virtual clock to the correct time, and execute the corresponding action.

Let's start by defining the structure in a new file: src/testing/TestScheduler.ts.

// src/testing/TestScheduler.ts

// Represents a scheduled task.
interface ScheduledAction {
  frame: number;
  action: () => void;
}

export class TestScheduler {
  private frame = 0;
  private actions: ScheduledAction[] = [];
  
  // A helper to push a task onto the queue.
  public schedule(action: () => void, delay = 0) {
    this.actions.push({ frame: this.frame + delay, action });
  }

  // The core execution engine.
  public flush() {
    // Process actions in chronological order.
    this.actions.sort((a, b) => a.frame - b.frame);

    while (this.actions.length > 0) {
      // Dequeue the next action.
      const action = this.actions.shift()!;
      
      // Advance virtual time to the action's scheduled frame.
      this.frame = action.frame;
      
      // Execute the action.
      action.action();
    }
  }
}

This is the heart of our virtual time machine. The schedule method allows us to queue work, and flush executes that work in perfect chronological order, instantly.

3. Creating Test-Friendly Operators

For our operators to be testable with the TestScheduler, they need a way to use its schedule method instead of the global setTimeout. A common pattern is to allow an optional scheduler to be passed to time-based operators.

Let's imagine we're building the delay operator (which we will formally do in a future lesson). A test-friendly implementation would look like this:

// src/operators/delay.ts (Hypothetical for now)
import { Observable } from '../Observable';
import { TestScheduler } from '../testing/TestScheduler'; // For type-hinting

// A generic Scheduler interface that our operators will use.
// TestScheduler will implement this.
export interface Scheduler {
  schedule(action: () => void, delay?: number): void;
}

export function delay<T>(duration: number, scheduler?: Scheduler) {
  return (source: Observable<T>): Observable<T> => {
    return new Observable(observer => {
      const subscription = source.subscribe({
        next: value => {
          // Use the scheduler if provided, otherwise fallback to setTimeout.
          if (scheduler) {
            scheduler.schedule(() => observer.next(value), duration);
          } else {
            setTimeout(() => observer.next(value), duration);
          }
        },
        error: err => observer.error(err),
        complete: () => {
          if (scheduler) {
            scheduler.schedule(() => observer.complete(), duration);
          } else {
            setTimeout(() => observer.complete(), duration);
          }
        },
      });
      return () => subscription.unsubscribe();
    });
  };
}

Notice how the operator now uses the provided scheduler to queue its emissions. This hooks it into our virtual time system during tests.

4. The run Helper and Scheduled Assertions

The final piece is to create an ergonomic API for writing tests, similar to the testScheduler.run() method from RxJS. This method should handle the setup, execution, and assertion scheduling automatically.

Assertions in a scheduled environment are tricky. A standard expect() call would execute immediately, before our flush() method has run any delayed actions. The solution is to schedule the assertions themselves.

Our run method will provide a helper, expectObservable, which records the emissions of an observable and schedules a final check to be run after all actions are flushed.

Let's implement this more advanced structure.

// src/testing/TestScheduler.ts (continued)
import { Observable } from '../Observable';

// A standardized way to represent observable notifications for testing.
export interface TestNotification<T> {
  frame: number;
  kind: 'N' | 'E' | 'C'; // Next, Error, Complete
  value?: T;
  error?: any;
}

// Update the Scheduler interface to be part of the same file.
export interface Scheduler {
  schedule(action: () => void, delay?: number): void;
}

export class TestScheduler implements Scheduler {
  private frame = 0;
  private actions: ScheduledAction[] = [];
  // A queue for assertions to be run at the very end.
  private assertions: (() => void)[] = [];

  constructor(private assertDeepEqual: (actual: any, expected: any) => void) {}

  public schedule(action: () => void, delay = 0) {
    this.actions.push({ frame: this.frame + delay, action });
  }

  private flush() {
    this.actions.sort((a, b) => a.frame - b.frame);
    while (this.actions.length > 0) {
      const action = this.actions.shift()!;
      this.frame = action.frame;
      action.action();
    }
    // After all actions are done, run the scheduled assertions.
    this.assertions.forEach(assert => assert());
  }

  // The main entry point for a test.
  public run(callback: (helpers: { expectObservable: Function }) => void) {
    // 1. Reset state for a clean test run.
    this.frame = 0;
    this.actions = [];
    this.assertions = [];
    
    // 2. Define and provide the helpers.
    const helpers = {
      expectObservable: <T>(source: Observable<T>) => {
        const actual: TestNotification<T>[] = [];
        const subscription = { closed: false };

        // Schedule the subscription to happen at the beginning of the test (frame 0).
        this.schedule(() => {
          source.subscribe({
            next: (value: T) => actual.push({ frame: this.frame, kind: 'N', value }),
            error: (err: any) => actual.push({ frame: this.frame, kind: 'E', error: err }),
            complete: () => actual.push({ frame: this.frame, kind: 'C' }),
          });
        });

        return {
          // This method schedules the final assertion.
          toBe: (expected: TestNotification<T>[]) => {
            this.assertions.push(() => {
              this.assertDeepEqual(actual, expected);
            });
          },
        };
      },
    };

    // 3. Execute the user's test logic, which will schedule actions and assertions.
    callback(helpers);

    // 4. Flush the queues to run everything.
    this.flush();
  }
}

This implementation is powerful. The run method orchestrates the entire test:

1. It resets the scheduler's state.
2. It calls the test function, passing in the expectObservable helper.
3. expectObservable subscribes to the stream and records all notifications with their virtual timestamps. Its toBe method queues the final assertion.
4. After the test function returns, run calls flush, which executes all scheduled operator logic and, finally, the assertions.

5. Putting it all Together: A Unit Test

Let's write a test for our hypothetical delay operator using the TestScheduler. We will create a test file src/operators/delay.test.ts.

// src/operators/delay.test.ts
import { describe, it, expect } from 'vitest';
import { of } from '../creators/of';
import { TestScheduler, TestNotification } from '../testing/TestScheduler';
import { delay } from './delay'; // Our hypothetical operator

describe('delay', () => {
  it('should delay emissions by the specified duration', () => {
    // 1. Instantiate the scheduler with an assertion function from our test framework.
    const scheduler = new TestScheduler((actual, expected) => {
      expect(actual).toEqual(expected);
    });

    // 2. Run the test within the scheduler's context.
    scheduler.run(({ expectObservable }) => {
      const source = of(1);
      const delayTime = 200;
      
      // 3. Apply the operator, passing the scheduler to it.
      const delayedStream = source.pipe(delay(delayTime, scheduler));

      // 4. Define the expected output as an array of notifications.
      // The value '1' should be emitted at frame 200.
      // Completion should also happen at frame 200.
      const expected: TestNotification<number>[] = [
        { frame: 200, kind: 'N', value: 1 },
        { frame: 200, kind: 'C' },
      ];

      // 5. Schedule the assertion.
      expectObservable(delayedStream).toBe(expected);
    });
  });
});

This test is declarative, synchronous, and runs instantly. It clearly describes the source stream (of(1)), the operation (delay(200)), and the expected outcome at precise moments in virtual time. This is a massive improvement over vi.advanceTimersByTime().

The blog post "RxJS Unit Testing" provides a great walkthrough of a similar process.

Test your understanding!

In our TestScheduler.run implementation, we scheduled the subscription to the source observable like this: this.schedule(() => { source.subscribe(...) });.

Why is it important to wrap the subscribe call in this.schedule() instead of just calling source.subscribe(...) directly inside the expectObservable helper?

Show answer

Scheduling the subscription is crucial for maintaining the correct order of operations within virtual time.

If we called source.subscribe(...) directly, it would execute immediately when expectObservable is called. If the source observable emits values synchronously (e.g., from of(1, 2, 3)), those emissions would be recorded at frame: 0 before any other actions might have been scheduled by the test setup.

By wrapping it in this.schedule(() => ..., 0), we place the act of subscribing into the action queue at the current frame (frame: 0). This ensures it is executed as part of the flush process, in the correct order relative to other actions also scheduled for frame 0. It makes the behavior consistent for both synchronous and asynchronous source observables.

Conclusion

Excellent work today! You have built the most critical piece of infrastructure for advanced reactive testing. By creating a TestScheduler, you've unlocked the ability to test complex, time-based asynchronous logic in a completely synchronous, deterministic, and readable way.

Key Takeaways:

• Virtual Time: A TestScheduler replaces the real clock with a virtual frame counter, allowing tests to run instantly.
• Action Queue: The scheduler works by processing a queue of ScheduledAction objects in chronological order.
• Scheduled Assertions: Helpers like expectObservable are necessary to defer assertions until after the flush cycle is complete, ensuring all asynchronous actions have been executed.
• Testable Operators: Asynchronous operators must be designed to accept an optional Scheduler to be hooked into the test environment.

In our next lesson, we will build upon this foundation. We will design and implement a parser for marble diagram strings, allowing us to replace the verbose TestNotification[] array with the concise and expressive marble syntax (e.g., '--a--b--|') that is the hallmark of RxJS testing.