gitbook/packages/gitbook/e2e/util.ts

import { argosScreenshot } from '@argos-ci/playwright';
import {
    type BrowserContext,
    type FrameLocator,
    type Page,
    type Response,
    expect,
    test,
} from '@playwright/test';
import deepMerge from 'deepmerge';
import rison from 'rison';
import type { DeepPartial } from 'ts-essentials';

import {
    CustomizationAIMode,
    CustomizationBackground,
    CustomizationCodeTheme,
    CustomizationCorners,
    CustomizationDefaultFont,
    CustomizationDefaultMonospaceFont,
    CustomizationDefaultThemeMode,
    CustomizationDepth,
    type CustomizationHeaderItem,
    CustomizationHeaderPreset,
    CustomizationIconsStyle,
    CustomizationLinksStyle,
    CustomizationLocale,
    CustomizationPageActionType,
    CustomizationSearchStyle,
    CustomizationSidebarBackgroundStyle,
    CustomizationSidebarListStyle,
    CustomizationTheme,
    type CustomizationThemedColor,
    type SiteCustomizationSettings,
    SiteExternalLinksTarget,
} from '@gitbook/api';

import { getContentTestURL, getTestURL } from '../tests/utils';

export interface Test {
    name: string;
    /**
     * URL to visit for testing.
     */
    url: string | (() => string | Promise<string>);
    cookies?: Parameters<BrowserContext['addCookies']>[0];
    /**
     * Test to run
     */
    run?: (page: Page, response: Response | null) => Promise<unknown>;
    /**
     * Re-applied right before every viewport screenshot (after Argos
     * stabilization), so it survives re-renders triggered by viewport resizing.
     *
     * Use this — rather than mutating the DOM once in `run` — to normalize
     * non-deterministic content (e.g. AI responses). A one-time mutation in `run`
     * is clobbered when React re-renders on resize (e.g. crossing the mobile
     * breakpoint), so only the first viewport ends up normalized.
     */
    normalizeBeforeScreenshot?: (page: Page) => Promise<void> | void;
    /**
     * Mode for the test.
     */
    mode?: 'page' | 'image';
    /**
     * Whether the test should be fullscreened during testing.
     */
    fullPage?: boolean;
    /**
     * Whether to take a screenshot of the test or set a threshold for the screenshot.
     */
    screenshot?:
        | false
        | {
              /**
               * Screenshot threshold.
               * From 0 to 1, where 0 is the most strict and 1 is the most permissive.
               * @default 0.5
               */
              threshold?: number;
              /**
               * Whether to wait for the table of contents to finish scrolling before taking the screenshot.
               */
              waitForTOCScrolling?: boolean;
          };
    /**
     * Whether to only run this test.
     */
    only?: boolean;
}

export type TestsCase = {
    name: string;
    skip?: boolean;
    tests: Test[];
    contentBaseURL?: string;
    /**
     * Whether screenshots in this test case should capture the full scrollable page by default.
     */
    fullPage?: boolean;
};

export const allLocales: CustomizationLocale[] = [
    CustomizationLocale.Fr,
    CustomizationLocale.Es,
    CustomizationLocale.Ja,
    CustomizationLocale.Zh,
];

export const allThemeModes: CustomizationDefaultThemeMode[] = [
    CustomizationDefaultThemeMode.Light,
    CustomizationDefaultThemeMode.Dark,
];

export const allTintColors: {
    label: string;
    value: CustomizationThemedColor | undefined;
}[] = [
    {
        label: 'Off',
        value: undefined,
    },
    { label: 'Primary', value: { light: '#346DDB', dark: '#346DDB' } },
    { label: 'Custom', value: { light: '#C62C68', dark: '#EF96B8' } },
];

export const allThemes: CustomizationTheme[] = [
    CustomizationTheme.Clean,
    CustomizationTheme.Muted,
    CustomizationTheme.Bold,
    CustomizationTheme.Gradient,
];

export const allDeprecatedThemePresets: CustomizationHeaderPreset[] = [
    CustomizationHeaderPreset.Default,
    CustomizationHeaderPreset.Bold,
    CustomizationHeaderPreset.Contrast,
    CustomizationHeaderPreset.Custom,
];

export const allSidebarBackgroundStyles: CustomizationSidebarBackgroundStyle[] = [
    CustomizationSidebarBackgroundStyle.Default,
    CustomizationSidebarBackgroundStyle.Filled,
];

export const allSearchStyles: CustomizationSearchStyle[] = [
    CustomizationSearchStyle.Prominent,
    CustomizationSearchStyle.Subtle,
];

// Common customization settings

export const headerLinks: CustomizationHeaderItem[] = [
    {
        title: 'Secondary button',
        to: { kind: 'url', url: 'https://www.gitbook.com' },
        style: 'button-secondary',
        links: [],
    },
    {
        title: 'Primary button',
        to: { kind: 'url', url: 'https://www.gitbook.com' },
        style: 'button-primary',
        links: [],
    },
];

export async function waitForCookiesDialog(page: Page) {
    const dialog = page.getByTestId('cookies-dialog');
    await expect(dialog).toBeVisible({
        // Cookies dialog may take some times to appear
        timeout: 10_000,
    });
}

export async function waitForHydration(page: Page) {
    await page.locator('html.hydrated').waitFor();
}

/**
 * Wait for the GitBook admin toolbar to be present.
 *
 * The toolbar only renders when signed in to GitBook. It is hidden from
 * screenshots (see `argosCSS`) because it auto-expands with an animation, so
 * use this to assert it is rendered without capturing its flaky visual state.
 */
export async function waitForAdminToolbar(page: Page) {
    await expect(page.getByTestId('admin-toolbar')).toBeVisible({
        timeout: 10_000,
    });
}

export async function waitForNotFound(_page: Page, response: Response | null) {
    expect(response).not.toBeNull();
    expect(response?.status()).toBe(404);
}

/**
 * Wait for an AI chat response to be fully settled before asserting or
 * screenshotting it.
 *
 * The chat exposes `aria-busy` on its container (`[data-testid="ai-chat"]`),
 * which stays true from the moment a message is sent until the stream — including
 * the follow-up suggestion phase — completes. Gating on it avoids the two main
 * sources of flakiness: capturing a "thinking" placeholder or a half-streamed
 * answer, and running the content normalization while React is still re-rendering
 * (which would clobber the replacements).
 *
 * Argos also waits for `aria-busy` to clear during its own stabilization
 * (`waitForAriaBusy`), so this is both an explicit gate and a backstop.
 *
 * Accepts a `Page` or a `FrameLocator` (for the embedded assistant in an iframe).
 */
export async function waitForAIChatResponse(scope: Page | FrameLocator) {
    await expect(scope.getByTestId('ai-chat')).toHaveAttribute('aria-busy', 'false', {
        timeout: 60_000,
    });
}

export async function setTimeToMorning(page: Page) {
    const now = new Date();
    now.setHours(8, 0, 0, 0); // 8:00:00.000 AM (local time)

    await page.clock.install({ time: now });
}

export async function waitForCoverImages(page: Page, options?: { darkMode?: boolean }) {
    const selector = options?.darkMode
        ? 'img[alt="Page cover"].dark\\:inline'
        : 'img[alt="Page cover"]:not(.dark\\:inline)';
    // Wait for cover images to exist (not the shimmer placeholder)
    await expect(page.locator(selector)).toBeVisible({
        timeout: 10_000,
    });
}

/**
 * Transform test cases into Playwright tests and run it.
 */
export function runTestCases(testCases: TestsCase[]) {
    for (const testCase of testCases) {
        if (testCase.skip) {
            continue;
        }

        test.describe(testCase.name, () => {
            for (const testEntry of testCase.tests) {
                const { mode = 'page' } = testEntry;
                const testFn = testEntry.only ? test.only : test;
                testFn(testEntry.name, async ({ page, context }) => {
                    const testEntryPathname =
                        typeof testEntry.url === 'function' ? await testEntry.url() : testEntry.url;
                    const url = testCase.contentBaseURL
                        ? getContentTestURL(
                              new URL(testEntryPathname, testCase.contentBaseURL).toString()
                          )
                        : getTestURL(testEntryPathname);

                    if (testEntry.cookies) {
                        await context.addCookies(
                            testEntry.cookies.map((cookie) => ({
                                ...cookie,
                                domain: new URL(url).host,
                                path: '/',
                            }))
                        );
                    }

                    // Reset the cross-space navigation state on every document load so the
                    // "Back to <space>" shortcut never leaks between navigations/tests. It is
                    // detected client-side from this sessionStorage, and a stale value (e.g.
                    // after a retry or a cross-space redirect) makes it appear or not
                    // non-deterministically, causing flaky screenshots.
                    await page.addInitScript(() => {
                        try {
                            sessionStorage.removeItem('gitbook-space-navigation:last');
                            sessionStorage.removeItem('gitbook-space-navigation:back');
                            sessionStorage.removeItem('gitbook-space-navigation:from-picker');
                        } catch {}
                    });

                    // Set the header to disable the Vercel toolbar
                    // But only on the main document as it'd cause CORS issues on other resources
                    await page.route('**/*', async (route, request) => {
                        if (request.resourceType() === 'document') {
                            await route.continue({
                                headers: {
                                    ...request.headers(),
                                    'x-vercel-skip-toolbar': '1',
                                },
                            });
                        } else {
                            await route.continue();
                        }
                    });

                    // Wait only for `domcontentloaded` rather than the default `load`: these
                    // are real customer sites whose third-party subresources can hang and
                    // never fire `load`, aborting the navigation. Argos stabilization (run in
                    // `beforeScreenshot`) still waits for images/fonts before capturing.
                    const response = await page.goto(url, { waitUntil: 'domcontentloaded' });
                    if (testEntry.run) {
                        await testEntry.run(page, response);
                    }
                    const screenshotOptions = testEntry.screenshot;
                    if (screenshotOptions !== false) {
                        const screenshotName = `${testCase.name} - ${testEntry.name}`;
                        if (mode === 'image') {
                            await argosScreenshot(page, screenshotName, {
                                viewports: ['macbook-13'],
                                threshold: screenshotOptions?.threshold ?? undefined,
                                fullPage: true,
                            });
                        } else {
                            await argosScreenshot(page, screenshotName, {
                                viewports: ['macbook-16', 'macbook-13', 'ipad-2', 'iphone-x'],
                                argosCSS: `
                            /* Hide Intercom */
                            .intercom-lightweight-app {
                                display: none !important;
                            }
                            /* Hide the GitBook admin toolbar: it auto-expands with an
                               animation, so its state at capture time is non-deterministic.
                               Its presence is asserted separately via waitForAdminToolbar. */
                            [data-testid="admin-toolbar"] {
                                display: none !important;
                            }
                                `,
                                threshold: screenshotOptions?.threshold ?? undefined,
                                fullPage: testEntry.fullPage ?? testCase.fullPage ?? false,
                                beforeScreenshot: async ({ runStabilization }) => {
                                    await runStabilization();
                                    if (screenshotOptions?.waitForTOCScrolling !== false) {
                                        await waitForTOCScrolling(page);
                                    }
                                    await waitForIcons(page);
                                    // Re-apply per viewport, last — after any resize-driven
                                    // re-render — so normalized content survives to capture.
                                    await testEntry.normalizeBeforeScreenshot?.(page);
                                },
                            });
                        }
                    }
                });
            }
        });
    }
}

/**
 * Create a URL with customization settings.
 */
export function getCustomizationURL(partial: DeepPartial<SiteCustomizationSettings>): string {
    // We replicate the theme migration logic from the API to the tests, because the don't get these settings from the API.
    // We can remove this once the migration to the new themes have been completed and the new theme styles are verified
    // Map the theme preset (+ tint) to one of the new themes
    const newTheme = (() => {
        if (partial.styling?.theme) {
            return partial.styling.theme;
        }

        switch (partial.header?.preset) {
            case CustomizationHeaderPreset.Bold:
            case CustomizationHeaderPreset.Contrast:
            case CustomizationHeaderPreset.Custom:
                return CustomizationTheme.Bold;

            case CustomizationHeaderPreset.None:
            case CustomizationHeaderPreset.Default:
                if (partial.styling?.tint) {
                    return CustomizationTheme.Muted;
                }

                return CustomizationTheme.Clean;
            default:
                return CustomizationTheme.Clean;
        }
    })();

    /**
     * Default customization settings.
     *
     * The customization object passed to the URL should be a valid API settings object. Hence we extend the test with necessary defaults.
     */
    const DEFAULT_CUSTOMIZATION: SiteCustomizationSettings = {
        styling: {
            theme: newTheme,
            primaryColor: { light: '#346DDB', dark: '#346DDB' },
            infoColor: { light: '#787878', dark: '#787878' },
            warningColor: { light: '#FE9A00', dark: '#FE9A00' },
            dangerColor: { light: '#FB2C36', dark: '#FB2C36' },
            successColor: { light: '#00C950', dark: '#00C950' },
            corners: CustomizationCorners.Rounded,
            depth: CustomizationDepth.Subtle,
            font: CustomizationDefaultFont.Inter,
            monospaceFont: CustomizationDefaultMonospaceFont.IBMPlexMono,
            background: CustomizationBackground.Plain,
            icons: CustomizationIconsStyle.Regular,
            links: CustomizationLinksStyle.Default,
            codeTheme: {
                default: {
                    light: CustomizationCodeTheme.DefaultLight,
                    dark: CustomizationCodeTheme.DefaultDark,
                },
                openapi: {
                    light: CustomizationCodeTheme.DefaultLight,
                    dark: CustomizationCodeTheme.DefaultDark,
                },
            },
            sidebar: {
                background: CustomizationSidebarBackgroundStyle.Default,
                list: CustomizationSidebarListStyle.Default,
            },
            search: CustomizationSearchStyle.Subtle,
        },
        internationalization: {
            locale: CustomizationLocale.En,
        },
        insights: {
            trackingCookie: true,
        },
        favicon: {},
        header: {
            preset: CustomizationHeaderPreset.Default,
            links: [],
        },
        footer: {
            groups: [],
        },
        themes: {
            default: CustomizationDefaultThemeMode.System,
            toggeable: true,
        },
        feedback: {
            enabled: false,
        },
        ai: {
            mode: CustomizationAIMode.None,
        },
        externalLinks: {
            target: SiteExternalLinksTarget.Self,
        },
        advancedCustomization: {
            enabled: true,
        },
        pagination: {
            enabled: true,
        },
        pageActions: {
            items: [
                CustomizationPageActionType.Assistant,
                CustomizationPageActionType.Markdown,
                CustomizationPageActionType.ExternalAi,
                CustomizationPageActionType.Mcp,
                CustomizationPageActionType.Pdf,
            ],
        },
        trademark: {
            enabled: true,
        },
        privacyPolicy: {
            url: 'https://www.gitbook.com/privacy',
        },
        socialPreview: {},
        socialAccounts: [],
    };

    const encoded = rison.encode_object(
        deepMerge(DEFAULT_CUSTOMIZATION, partial, {
            arrayMerge: (_target, source) => source,
        })
    );

    const searchParams = new URLSearchParams();
    searchParams.set('customization', encoded);

    return `?${searchParams.toString()}`;
}

/**
 * Wait for all icons present on the page to be loaded.
 */
export async function waitForIcons(page: Page) {
    await page.waitForFunction(() => {
        type IconURLStates = Record<
            string,
            { state: 'pending'; uri: null } | { state: 'loaded'; uri: string }
        >;
        const iconStatesWindow = window as Window & { __ICONS_STATES__?: IconURLStates };
        const urlStates: IconURLStates = iconStatesWindow.__ICONS_STATES__ || {};
        iconStatesWindow.__ICONS_STATES__ = urlStates;

        const fetchSvgAsDataUri = async (url: string): Promise<string> => {
            const response = await fetch(url);
            if (!response.ok) {
                throw new Error(`Failed to fetch SVG: ${response.status}`);
            }

            const svgText = await response.text();
            const encoded = encodeURIComponent(svgText).replace(/'/g, '%27').replace(/"/g, '%22');

            return `data:image/svg+xml;charset=utf-8,${encoded}`;
        };

        const loadUrl = (url: string) => {
            // Mark the URL as pending.
            urlStates[url] = { state: 'pending', uri: null };
            fetchSvgAsDataUri(url).then((uri) => {
                urlStates[url] = { state: 'loaded', uri };
            });
        };

        const icons = Array.from(document.querySelectorAll('svg.gb-icon'));
        const results = icons.map((icon) => {
            if (!(icon instanceof SVGElement)) {
                throw new Error('Icon is not an SVGElement');
            }

            // Ignore icons that are not visible.
            if (!icon.checkVisibility()) {
                return true;
            }

            const state = icon.getAttribute('data-argos-state');

            if (state === 'pending') {
                return false;
            }

            if (state === 'loaded') {
                return true;
            }

            const maskImage = icon.querySelector('[data-testid="mask-image"]');
            if (!maskImage) {
                const inlineContent = icon.querySelector(
                    'path, circle, ellipse, line, polygon, polyline, rect, g, use'
                );
                if (inlineContent) {
                    icon.setAttribute('data-argos-state', 'loaded');
                    return true;
                }

                throw new Error('Icon has no inline SVG content or mask-image element');
            }

            const url = maskImage.getAttribute('href');
            // If URL is invalid we throw an error.
            if (!url) {
                throw new Error('No mask-image url');
            }

            // If the URL is already queued for loading, we return the state.
            if (urlStates[url]) {
                if (urlStates[url].state === 'loaded') {
                    icon.setAttribute('data-argos-state', 'pending');
                    icon.style.maskImage = `url("${urlStates[url].uri}")`;
                    requestAnimationFrame(() => {
                        icon.setAttribute('data-argos-state', 'loaded');
                    });
                    return false;
                }

                return false;
            }

            loadUrl(url);
            return false;
        });

        return results.every((x) => x);
    });
}

/**
 * Wait for TOC to be correctly scrolled into view.
 */
async function waitForTOCScrolling(page: Page) {
    const viewport = await page.viewportSize();
    if (viewport && viewport.width >= 1024 && !page.url().includes('~gitbook/embed/demo')) {
        // The embed demo is an iframe, which means the viewport is only a fraction of the main document. So there is no open TOC to scroll to.
        const toc = page.getByTestId('table-of-contents');
        await expect(toc).toBeVisible();
        await page.evaluate(() => {
            const tocScrollContainer = document.querySelector(
                '[data-testid="table-of-contents"] [data-testid="toc-scroll-container"]'
            );
            if (!tocScrollContainer) {
                throw new Error('TOC scroll container not found');
            }
            tocScrollContainer.scrollTo(0, 0);
        });
    }
}