vitest-environment-puppeteer/dist/index.js

Puppeteer environment for Vitest.

import { cosmiconfig } from 'cosmiconfig';
import chalk from 'chalk';
import 'node:os';

const DEFAULT_CONFIG = {
    browserContext: "default",
    exitOnPageError: true
};
const explorer = cosmiconfig("vitest-puppeteer");
const getDefaultConfig = ()=>{
    if (process.env.CI) {
        return {
            ...DEFAULT_CONFIG,
            launch: {
                args: [
                    "--no-sandbox",
                    "--disable-setuid-sandbox",
                    "--disable-background-timer-throttling",
                    "--disable-backgrounding-occluded-windows",
                    "--disable-renderer-backgrounding"
                ]
            }
        };
    }
    return DEFAULT_CONFIG;
};
const readConfigFile = async (configFile, searchFrom)=>{
    if (!configFile) {
        const result = await explorer.search(searchFrom);
        return result ? result.config : null;
    }
    const result = await explorer.load(configFile);
    return result ? result.config : null;
};
const readConfig = async (searchFrom)=>{
    const defaultConfig = getDefaultConfig();
    const config = await readConfigFile(process.env.VITEST_PUPPETEER_CONFIG, searchFrom);
    if (!config) return defaultConfig;
    return {
        ...defaultConfig,
        ...config
    };
};

const CTRL_C = "\u0003";
const CTRL_D = "\u0004";
const ENTER = "\r";
const blockStdin = ()=>{
    console.log(chalk.blue("\n\n🕵️‍  Code is paused, press enter to resume"));
    return new Promise((resolve)=>{
        const { stdin } = process;
        const onKeyPress = (key)=>{
            if (key === CTRL_C || key === CTRL_D || key === ENTER) {
                stdin.removeListener("data", onKeyPress);
                if (!listening) {
                    if (stdin.isTTY) {
                        stdin.setRawMode(false);
                    }
                    stdin.pause();
                }
                resolve();
            }
        };
        const listening = stdin.listenerCount("data") > 0;
        if (!listening) {
            if (stdin.isTTY) {
                stdin.setRawMode(true);
            }
            stdin.resume();
            stdin.setEncoding("utf8");
        }
        stdin.on("data", onKeyPress);
    });
};

const getPuppeteer = async ()=>{
    try {
        // @ts-expect-error Puppeteer docs say this will return PuppeteerNode
        return import('puppeteer');
    } catch  {
        // @ts-expect-error Puppeteer docs say this will return PuppeteerNode
        return import('puppeteer-core');
    }
};
const readWorkersCount = ()=>{
    if (!process.env.WORKERS_COUNT) {
        throw new Error("Invariant: WORKERS_COUNT not found");
    }
    return Number(process.env.WORKERS_COUNT);
};
const readWsEndpoints = ()=>{
    if (!process.env.PUPPETEER_WS_ENDPOINTS) {
        throw new Error("Invariant: PUPPETEER_WS_ENDPOINTS not found");
    }
    return JSON.parse(process.env.PUPPETEER_WS_ENDPOINTS);
};
const getJestWorkerId = ()=>{
    if (!process.env.VITEST_WORKER_ID) {
        throw new Error("Invariant: VITEST_WORKER_ID not found");
    }
    return Number(process.env.VITEST_WORKER_ID);
};
const getWorkerIndex = ()=>{
    // Jest worker ID starts at 1
    return getJestWorkerId() - 1;
};
const getWorkerWsEndpointIndex = ()=>{
    return Math.min(readWorkersCount() - 1, getWorkerIndex());
};
const getWorkerWsEndpoint = ()=>{
    const wsEndpoints = readWsEndpoints();
    const index = getWorkerWsEndpointIndex();
    const endpoint = wsEndpoints[index];
    if (!endpoint) {
        throw new Error("Invariant: worker WS endpoint not found");
    }
    return endpoint;
};
const connectBrowserFromWorker = async (config)=>{
    const wsEndpoint = getWorkerWsEndpoint();
    const puppeteer = await getPuppeteer();
    return puppeteer.connect({
        ...config.connect,
        ...config.launch,
        browserURL: undefined,
        browserWSEndpoint: wsEndpoint
    });
};

const testTimeoutSymbol = Symbol.for("TEST_TIMEOUT_SYMBOL");
const handlePageError = (error)=>{
    process.emit("uncaughtException", error);
};
const getBrowser = (global)=>{
    if (!global.browser) {
        throw new Error("Cannot access browser before launching browser.");
    }
    return global.browser;
};
const getContext = (global)=>{
    if (!global.context) {
        throw new Error("Cannot access context before launching context.");
    }
    return global.context;
};
const connectBrowser = async (global)=>{
    if (global.browser) {
        throw new Error("Cannot connect browser before closing previous browser.");
    }
    global.browser = await connectBrowserFromWorker(global.puppeteerConfig);
};
const disconnectBrowser = async (global)=>{
    if (!global.browser) return;
    await global.browser.disconnect();
    global.browser = undefined;
};
const getPage = (global)=>{
    if (!global.page) {
        throw new Error("Cannot access page before launching browser.");
    }
    return global.page;
};
const openPage = async (global)=>{
    if (global.page) {
        throw new Error("Cannot open page before closing previous page.");
    }
    const page = await getContext(global).newPage();
    if (global.puppeteerConfig.exitOnPageError) {
        page.on("pageerror", handlePageError);
    }
    global.page = page;
};
const closePage = async (global)=>{
    if (!global.page) return;
    if (global.puppeteerConfig.exitOnPageError) {
        global.page.off("pageerror", handlePageError);
    }
    await global.page.close({
        runBeforeUnload: Boolean(global.puppeteerConfig.runBeforeUnloadOnClose)
    });
    global.page = undefined;
};
const createContext = async (global)=>{
    if (global.context) {
        throw new Error("Cannot create context before closing previous context.");
    }
    const configBrowserContext = global.puppeteerConfig.browserContext ?? "default";
    const browser = getBrowser(global);
    switch(configBrowserContext){
        case "default":
            global.context = browser.defaultBrowserContext();
            break;
        case "incognito":
            global.context = "createBrowserContext" in browser ? await browser.createBrowserContext() : await browser.createIncognitoBrowserContext();
            break;
        default:
            throw new Error(`browserContext should be either 'incognito' or 'default'. Received '${configBrowserContext}'`);
    }
};
const closeContext = async (global)=>{
    if (!global.context) return;
    const browser = getBrowser(global);
    // If a custom context was created, close it
    if (global.context !== browser.defaultBrowserContext()) {
        await global.context.close();
    }
    global.context = undefined;
};
const initAll = async (global)=>{
    await connectBrowser(global);
    await createContext(global);
    await openPage(global);
};
const closeAll = async (global)=>{
    await closePage(global);
    await closeContext(global);
    await disconnectBrowser(global);
};
const PuppeteerEnvironment = {
    name: "puppeteer",
    transformMode: "ssr",
    async setup (global) {
        global.puppeteerConfig = await readConfig();
        global.vitestPuppeteer = {
            debug: async ()=>{
                // Set timeout to 4 days
                global[testTimeoutSymbol] = 345600000;
                // Run a debugger (in case Puppeteer has been launched with `{ devtools: true }`)
                await getPage(global).evaluate(()=>{
                    debugger;
                });
                return blockStdin();
            },
            resetPage: async ()=>{
                await closePage(global);
                await openPage(global);
            },
            resetBrowser: async ()=>{
                await closeAll(global);
                await initAll(global);
            }
        };
        await initAll(global);
        return {
            async teardown () {
                await closeAll(global);
            }
        };
    }
};

export { PuppeteerEnvironment as TestEnvironment, PuppeteerEnvironment as default };