sse-session/tests/utils/async-push-iterator.test.ts

import { expect, test, vi } from 'vitest';

import { AsyncPushIterator } from '../../src/utils/async-push-iterator.js';

/**
 * Collects every value from the iterator into an array.
 *
 * @param iterator - Iterator under test.
 */
const collectAll = async <T>(iterator: AsyncPushIterator<T>): Promise<T[]> => {
    const results: T[] = [];

    for await (const value of iterator) {
        results.push(value);
    }

    return results;
};

/**
 * Tests that values pushed while a consumer is already waiting are delivered in order.
 */
const testPushComposedPushAndConsume = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    const result = new Promise<number[]>((resolve) => {
        void (async (): Promise<void> => {
            resolve(await collectAll(iterator));
        })();
    });

    iterator.push(1);
    iterator.push(2);
    iterator.push(3);
    iterator.close();

    await expect(result).resolves.toEqual([ 1, 2, 3 ]);
};

/**
 * Tests that values pushed before `for await...of` starts are buffered and yielded
 * once the consumer begins reading.
 */
const testPushComposedBuffersValuesPushedBeforeLoopStarts = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    iterator.push(1);
    iterator.push(2);
    iterator.push(3);

    const result = collectAll(iterator);

    iterator.close();

    await expect(result).resolves.toEqual([ 1, 2, 3 ]);
};

/**
 * Tests that the iterator completes with no values when nothing was pushed.
 */
const testPushComposedResolvesWithNoValues = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    const result = new Promise<number[]>((resolve) => {
        void (async (): Promise<void> => {
            resolve(await collectAll(iterator));
        })();
    });

    iterator.close();

    await expect(result).resolves.toEqual([]);
};

/**
 * Tests that values pushed after {@link AsyncPushIterator.close} are ignored.
 */
const testPushComposedIgnoresValuesAfterClose = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    const result = new Promise<number[]>((resolve) => {
        void (async (): Promise<void> => {
            resolve(await collectAll(iterator));
        })();
    });

    iterator.push(1);
    iterator.push(2);
    iterator.push(3);
    iterator.close();
    iterator.push(4);

    await expect(result).resolves.toEqual([ 1, 2, 3 ]);
};

/**
 * Tests that only one async consumer can read from the composed ReadableStream at a time.
 *
 * Unlike the hand-rolled async-push-iterator, the second consumer fails with a
 * stream lock error rather than TooManyAsyncIteratorsError.
 */
const testPushComposedRejectsMultipleConsumers = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    const failureFlag = vi.fn();

    const successfulIterator = (): Promise<number[]> =>
        new Promise((resolve) => {
            void (async (): Promise<void> => {
                resolve(await collectAll(iterator));
            })();
        });

    const failedIterator = (): Promise<void> =>
        new Promise((resolve, reject) => {
            void (async (): Promise<void> => {
                try {
                    /* eslint-disable-next-line */
                    for await (const _value of iterator) {
                    }
                } catch (error) {
                    failureFlag();
                    reject(error);
                }

                resolve();
            })();
        });

    const promises = [ successfulIterator(), failedIterator().catch(() => {}) ];

    iterator.close();

    await Promise.all(promises);

    expect(failureFlag).toHaveBeenCalledOnce();
};

/**
 * Tests that closing before iteration starts lets the loop finish immediately.
 */
const testPushComposedResolvesWhenClosedBeforeLoop = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    iterator.close();

    await expect(collectAll(iterator)).resolves.toEqual([]);
};

/**
 * Tests that breaking out of `for await...of` early does not cancel the stream.
 *
 * {@link AsyncPushIterator} uses `preventCancel: true` so producers can keep pushing
 * and a later consumer can read the remaining values.
 */
const testPushComposedAllowsPushingAfterEarlyBreak = async (): Promise<void> => {
    const iterator = new AsyncPushIterator<number>();

    iterator.push(1);

    const firstPass: number[] = [];

    for await (const value of iterator) {
        firstPass.push(value);
        break;
    }

    iterator.push(2);
    iterator.push(3);
    iterator.close();

    const secondPass = await collectAll(iterator);

    expect(firstPass).toEqual([ 1 ]);
    expect(secondPass).toEqual([ 2, 3 ]);
};

const runTests = async (): Promise<void> => {
    test('AsyncPushIterator (composed): pushes and consumes values', testPushComposedPushAndConsume);
    test('AsyncPushIterator (composed): buffers values pushed before the for-await loop starts', testPushComposedBuffersValuesPushedBeforeLoopStarts);
    test('AsyncPushIterator (composed): resolves with no values when nothing was pushed', testPushComposedResolvesWithNoValues);
    test('AsyncPushIterator (composed): ignores values pushed after close', testPushComposedIgnoresValuesAfterClose);
    test('AsyncPushIterator (composed): rejects multiple consumers', testPushComposedRejectsMultipleConsumers);
    test('AsyncPushIterator (composed): resolves immediately when closed before the loop starts', testPushComposedResolvesWhenClosedBeforeLoop);
    test('AsyncPushIterator (composed): keeps the stream open after an early break', testPushComposedAllowsPushingAfterEarlyBreak);
};

await runTests();