extend-expect.ts

import type { Assertion, ExpectStatic } from 'vitest';

import { TO_FINISH_WITH_OPTIONS } from './consts.js';
import type { RunTestResult } from './run-test-result.js';
import type { Interval, ToFinishWithOptions } from './types.js';

export const extendExpect = (expect: ExpectStatic): ExpectStatic => {
    expect.extend({
        toBeArray: (actual) => {
            const pass = Array.isArray(actual);
            return {
                pass,
                message: () => `Expected "${actual}" to be a array`,
            };
        },
        toBeBoolean: (received) => {
            const actual = typeof received;
            const expected = 'boolean';
            return {
                pass: actual === expected,
                message: () => `Expected "${actual}" to be a boolean`,
                expected: 'boolean',
            };
        },
        toBeEmptyArray: (actual) => {
            const pass = Array.isArray(actual) && actual.length === 0;
            return {
                pass,
                message: () => `Expected ${actual} to be an empty array`,
                actual,
                expected: [],
            };
        },
        toBeNonEmptyArray: (actual) => {
            const pass = Array.isArray(actual) && actual.length > 0;
            return {
                pass,
                message: () => `Expected ${actual} to be a non-empty array`,
            };
        },
        toBeNonEmptyString: (actual) => {
            const pass = typeof actual === 'string' && actual.length > 0;
            return {
                pass,
                message: () => `Expected "${actual}" to be a non-empty string`,
            };
        },
        toBeNumber: (actual) => {
            const pass = typeof actual === 'number';
            return {
                pass,
                message: () => `Expected "${actual}" to be a number`,
            };
        },
        toBeNonEmptyObject: (actual) => {
            const isObject = typeof actual === 'object' && !Array.isArray(actual);
            const isNonEmpty = Object.entries(actual).length > 0;
            const pass = isObject && isNonEmpty;
            return {
                pass,
                message: () => `Expected "${actual}" to be a non-emtpy object`,
            };
        },
        toBeObject: (actual) => {
            const pass = typeof actual === 'object' && !Array.isArray(actual);
            return {
                pass,
                message: () => `Expected "${actual}" to be an object`,
            };
        },
        toBeString: (received) => {
            const actual = typeof received;
            const expected = 'string';
            return {
                pass: actual === expected,
                message: () => `Expected "${actual}" to be a string`,
                actual,
                expected,
            };
        },
        toBeTrue: (actual) => {
            const pass = actual === true;
            return {
                pass,
                message: () => `Expected "${actual}" to be true`,
                expected: true,
                actual,
            };
        },
        toBeFalse: (actual) => {
            const pass = actual === false;
            return {
                pass,
                message: () => `Expected "${actual}" to be false`,
            };
        },
        toBeWholeNumber: (received) => {
            const pass = Number.isInteger(received);
            return {
                pass,
                message: () => `Expected "${received}" to be a whole number`,
            };
        },
        toBeWithinRange: (received, lower, upper) => {
            const pass = received >= lower && received <= upper;
            return {
                pass,
                message: () => `Expected ${received} to be within range <${lower}, ${upper}>`,
            };
        },
        toFinishWith: async <PpeEvent extends string>(
            received: RunTestResult,
            userOptions: ToFinishWithOptions<PpeEvent>,
        ) => {
            const options = {
                ...TO_FINISH_WITH_OPTIONS,
                ...userOptions,
            };

            const failedAssertions: string[] = [];

            const diffs: Diffs = {
                pass: true,
                actual: ['Run:'],
                expected: ['Run:'],
            };
            {
                const expected = options?.status;
                const actual = received.status;
                const statusTest = expected === actual;
                if (statusTest === false) {
                    diffs.pass = false;
                    diffs.actual.push(`status=${actual}`);
                    diffs.expected.push(`status=${expected}`);
                    failedAssertions.push(`Failed status check, expected "${expected}", got "${actual}".`);
                }
            }

            const {
                chargedEventCounts,
                stats: { durationMillis },
            } = await received.getRunInfo();
            const datasetItemCount = (await received.getDataset()).items.length;
            const stats = await received.getStatistics();

            const checkInterval = (
                value: number | undefined,
                key: ['datasetItemCount', 'duration', 'failedRequests', 'requestsRetries'][number],
                label: string,
            ) => {
                const result = isWithinInterval(diffs, value, options, key);
                if (result === false) {
                    failedAssertions.push(
                        `Failed ${label} check, expected ${JSON.stringify(options[key])}, got ${value}.`,
                    );
                }
            };

            checkInterval(datasetItemCount, 'datasetItemCount', 'dataset item count');
            checkInterval(durationMillis, 'duration', 'duration');
            checkInterval(stats?.requestsFailed, 'failedRequests', 'failed requests');
            checkInterval(stats?.requestsRetries, 'requestsRetries', 'requests retries');

            if (options.maxRetriesPerRequest !== null) {
                const maxRetriesPerRequestObserved = (stats?.requestRetryHistogram ?? [0]).length - 1;
                if (maxRetriesPerRequestObserved > options.maxRetriesPerRequest) {
                    diffs.pass = false;
                    diffs.actual.push(`maxRetriesPerRequest=${maxRetriesPerRequestObserved}`);
                    diffs.expected.push(`maxRetriesPerRequest<=${options.maxRetriesPerRequest}`);
                    failedAssertions.push(
                        `Failed max retries observed check, expected <=${options.maxRetriesPerRequest}, got ${maxRetriesPerRequestObserved}.`,
                    );
                }
            }
            const ppeDiffs: Diffs = {
                pass: true,
                actual: ['PPE Events:'],
                expected: ['PPE Events:'],
            };

            if (options.chargedEventCounts) {
                const expected = options.chargedEventCounts;
                const actual = chargedEventCounts ?? {};
                const uniquePpeEvents = new Set([
                    ...Array.from(Object.keys(expected)),
                    ...Array.from(Object.keys(actual)),
                ]);

                for (const ppeEvent of uniquePpeEvents) {
                    if (!(ppeEvent in expected)) {
                        isWithinInterval(ppeDiffs, actual[ppeEvent], { [ppeEvent]: 0 }, ppeEvent);
                        continue;
                    }
                    if (!(ppeEvent in actual)) {
                        ppeDiffs.pass = false;
                        ppeDiffs.actual.push(`${ppeEvent}=X`);
                        ppeDiffs.expected.push(`${ppeEvent}=${expected[ppeEvent as PpeEvent]}`);
                        continue;
                    }
                    isWithinInterval(ppeDiffs, actual[ppeEvent], expected, ppeEvent as PpeEvent);
                }

                if (ppeDiffs.pass === false) {
                    diffs.pass = ppeDiffs.pass;
                    diffs.actual.push(ppeDiffs.actual.join('\n    '));
                    diffs.expected.push(ppeDiffs.expected.join('\n    '));
                    failedAssertions.push(
                        `Failed PPE event counts check, expected ${JSON.stringify(
                            options.chargedEventCounts,
                        )}, got ${JSON.stringify(chargedEventCounts)}.`,
                    );
                }
            }

            {
                const { forbiddenLogs } = options;
                const log = await received.getLog();
                const occuredLogs = [];
                for (const forbiddenLog of forbiddenLogs) {
                    if (log.includes(forbiddenLog)) {
                        occuredLogs.push(forbiddenLog);
                    }
                }
                if (occuredLogs.length > 0) {
                    diffs.pass = false;
                    diffs.actual.push(` logs=[${occuredLogs.join(', ')}]`);
                    diffs.expected.push(` logs=[]`);
                    failedAssertions.push(
                        `Failed forbidden logs check, expected [] but got [${occuredLogs.join(', ')}].`,
                    );
                }
            }

            return {
                pass: diffs.pass,
                message: () => `Run did not finish as expected. Failed assertions: ${failedAssertions.join(' ')}`,
                actual: diffs.actual.join('\n  '),
                expected: diffs.expected.join('\n  '),
            };
        },
        toStartWith: (received, prefix) => {
            const pass = received?.startsWith(prefix);
            return {
                pass,
                message: () => `Expected "${received}" to start with "${prefix}"`,
            };
        },
    });

    // here we are switching logic of expect and expect.soft
    // expect => expect.soft
    // expect.hard => expect
    const { soft } = expect;
    const oldExpect = { ...expect };
    const softExpect = <T>(actual: T, message?: string): Assertion => {
        return soft(actual, message);
    };

    for (const [key, value] of Object.entries(oldExpect)) {
        // @ts-expect-error: No idea how to type this properly
        softExpect[key] = value;
    }

    softExpect.hard = <T>(actual: T, message?: string): Assertion<T> => {
        return expect(actual, message);
    };

    return softExpect as unknown as ExpectStatic;
};

type Diffs = {
    pass: boolean;
    actual: string[];
    expected: string[];
};

const isWithinInterval = <T extends string>(
    diffs: Diffs,
    actual: number | undefined,
    options: Record<T, Interval | null>,
    intervalOption: T,
) => {
    const expected = options[intervalOption];
    if (expected === null) {
        // check is disabled if expected value is null
        return undefined;
    }
    if (typeof expected === 'number') {
        if (actual !== expected) {
            // eslint-disable-next-line no-param-reassign --- Don't want to touch this code :)
            diffs.pass = false;
            diffs.actual.push(`${intervalOption}=${actual}`);
            diffs.expected.push(`${intervalOption}=${expected}`);
            return false;
        }
    } else if (typeof expected === 'object') {
        const { min, max } = expected;
        if (actual === undefined || (min !== undefined && actual < min) || (max !== undefined && actual > max)) {
            // eslint-disable-next-line no-param-reassign
            diffs.pass = false;
            diffs.actual.push(`${intervalOption}=${actual}`);
            diffs.expected.push(`${intervalOption}=<${min ?? ''},${max ?? ''}>`);
            return false;
        }
    }
    return undefined;
};