Extend hooks api with prepare and getHttpHeaders properties

.storybook/main.js

@@ -14,10 +14,6 @@ let stories = [
   '../stories/pages/**/*.stories.@(js|jsx|ts|tsx)',
 ];
 
-if (process.env.STRESS_TEST) {
-  stories.push('../stories/stress-test/*.stories.@(js|jsx|ts|tsx)');
-}
-
 if (process.env.TEST_FAILURES) {
   stories = ['../stories/expected-failures/*.stories.@(js|jsx|ts|tsx)'];
 }

README.md

@@ -24,6 +24,8 @@ Storybook test runner turns all of your stories into executable tests.
   - [2 - Run tests with --coverage flag](#2---run-tests-with---coverage-flag)
   - [3 - Merging code coverage with coverage from other tools](#3---merging-code-coverage-with-coverage-from-other-tools)
 - [Experimental test hook API](#experimental-test-hook-api)
+  - [prepare](#prepare)
+  - [getHttpHeaders](#gethttpheaders)
   - [DOM snapshot recipe](#dom-snapshot-recipe)
   - [Image snapshot recipe](#image-snapshot-recipe)
   - [Render lifecycle](#render-lifecycle)
@@ -401,10 +403,42 @@ To enable use cases like visual or DOM snapshots, the test runner exports test h
 
 There are three hooks: `setup`, `preRender`, and `postRender`. `setup` executes once before all the tests run. `preRender` and `postRender` execute within a test before and after a story is rendered.
 
-The render functions are async functions that receive a [Playwright Page](https://playwright.dev/docs/pages) and a context object with the current story `id`, `title`, and `name`. They are globally settable by `@storybook/test-runner`'s `setPreRender` and `setPostRender` APIs.
+The render functions are async functions that receive a [Playwright Page](https://playwright.dev/docs/pages) and a context object with the current story's `id`, `title`, and `name`. They are globally settable by `@storybook/test-runner`'s `setPreRender` and `setPostRender` APIs.
 
 All three functions can be set up in the configuration file `.storybook/test-runner.js` which can optionally export any of these functions.
 
+Apart from these hooks, there are additional properties you can set in `.storybook/test-runner.js`:
+
+#### prepare
+
+The test-runner has a default `prepare` function which gets the browser in the right environment before testing the stories. You can override this behavior, in case you might want to hack the behavior of the browser. For example, you might want to set a cookie, or add query parameters to the visiting URL, or do some authentication before reaching the Storybook URL. You can do that by overriding the `prepare` function.
+
+The `prepare` function receives an object containing:
+
+- `browserContext`: a [Playwright Browser Context](https://playwright.dev/docs/api/class-browsercontext) instance
+- `page`: a [Playwright Page](https://playwright.dev/docs/api/class-page) instance.
+- `testRunnerConfig`: the test runner configuration object, coming from the `.storybook/test-runner.js`.
+
+For reference, please use the [default `prepare`](https://github.com/storybookjs/test-runner/blob/next/src/setup-page.ts#L12) function as a starting point.
+
+> **NOTE:** If you override the default prepare behavior, even though this is powerful, you will be responsible for properly preparing the browser. Future changes to the default prepare function will not get included in your project, so you will have to keep an eye out for changes in upcoming releases.
+
+#### getHttpHeaders
+
+The test-runner makes a few `fetch` calls to check the status of a Storybook instance, and to get the index of the Storybook's stories. Additionally, it visits a page using Playwright. In all of these scenarios, it's possible, depending on where your Storybook is hosted, that you might need to set some HTTP headers. For example, if your Storybook is hosted behind a basic authentication, you might need to set the `Authorization` header. You can do so by passing a `getHttpHeaders` function to your test-runner config. That function receives the `url` of the fetch calls and page visits, and should return an object with the headers to be set.
+
+```js
+// .storybook/test-runner.js
+module.exports = {
+  getHttpHeaders: async (url) => {
+    const token = url.includes('prod') ? 'XYZ' : 'ABC';
+    return {
+      Authorization: `Bearer ${token}`,
+    };
+  },
+};
+```
+
 > **NOTE:** These test hooks are experimental and may be subject to breaking changes. We encourage you to test as much as possible within the story's play function.
 
 ### DOM snapshot recipe

bin/test-storybook.js

@@ -11,6 +11,7 @@ const path = require('path');
 const tempy = require('tempy');
 const { getCliOptions } = require('../dist/cjs/util/getCliOptions');
 const { getStorybookMetadata } = require('../dist/cjs/util/getStorybookMetadata');
+const { getTestRunnerConfig } = require('../dist/cjs/util/getTestRunnerConfig');
 const { transformPlaywrightJson } = require('../dist/cjs/playwright/transformPlaywrightJson');
 
 // Do this as the first thing so that any code reading it knows the right env.
@@ -19,6 +20,8 @@ process.env.NODE_ENV = 'test';
 process.env.STORYBOOK_TEST_RUNNER = 'true';
 process.env.PUBLIC_URL = '';
 
+let getHttpHeaders = (_url) => Promise.resolve({});
+
 // Makes the script crash on unhandled rejections instead of silently
 // ignoring them. In the future, promise rejections that are not handled will
 // terminate the Node.js process with a non-zero exit code.
@@ -125,7 +128,8 @@ async function executeJestPlaywright(args) {
 
 async function checkStorybook(url) {
   try {
-    const res = await fetch(url, { method: 'HEAD' });
+    const headers = await getHttpHeaders(url);
+    const res = await fetch(url, { method: 'HEAD', headers });
     if (res.status !== 200) throw new Error(`Unxpected status: ${res.status}`);
   } catch (e) {
     console.error(
@@ -144,8 +148,13 @@ async function checkStorybook(url) {
 async function getIndexJson(url) {
   const indexJsonUrl = new URL('index.json', url).toString();
   const storiesJsonUrl = new URL('stories.json', url).toString();
+  const headers = await getHttpHeaders(url);
+  const fetchOptions = { headers };
 
-  const [indexRes, storiesRes] = await Promise.all([fetch(indexJsonUrl), fetch(storiesJsonUrl)]);
+  const [indexRes, storiesRes] = await Promise.all([
+    fetch(indexJsonUrl, fetchOptions),
+    fetch(storiesJsonUrl, fetchOptions),
+  ]);
 
   if (indexRes.ok) {
     try {
@@ -219,6 +228,13 @@ const main = async () => {
     process.exit(0);
   }
 
+  process.env.STORYBOOK_CONFIG_DIR = runnerOptions.configDir;
+
+  const testRunnerConfig = getTestRunnerConfig(runnerOptions.configDir);
+  if (testRunnerConfig.getHttpHeaders) {
+    getHttpHeaders = testRunnerConfig.getHttpHeaders;
+  }
+
   // set this flag to skip reporting coverage in watch mode
   isWatchMode = jestOptions.watch || jestOptions.watchAll;
 
@@ -261,8 +277,6 @@ const main = async () => {
     process.env.TEST_MATCH = '**/*.test.js';
   }
 
-  process.env.STORYBOOK_CONFIG_DIR = runnerOptions.configDir;
-
   const { storiesPaths, lazyCompilation } = getStorybookMetadata();
   process.env.STORYBOOK_STORIES_PATTERN = storiesPaths;
 

playwright/custom-environment.js

@@ -6,7 +6,7 @@ const PlaywrightEnvironment = require('jest-playwright-preset/lib/PlaywrightEnvi
 class CustomEnvironment extends PlaywrightEnvironment {
   async setup() {
     await super.setup();
-    await setupPage(this.global.page);
+    await setupPage(this.global.page, this.global.context);
   }
 
   async teardown() {

playwright/jest-setup.js

@@ -14,6 +14,6 @@ if (testRunnerConfig) {
 }
 
 // If the transformed tests need a dependency, it has to be globally available
-// in order to work both in defaul (file transformation) and stories/index.json mode.
+// in order to work both in default (file transformation) and stories/index.json mode.
 globalThis.__sbSetupPage = setupPage;
 globalThis.__sbCollectCoverage = process.env.STORYBOOK_COLLECT_COVERAGE === 'true';

src/playwright/hooks.ts

@@ -1,4 +1,4 @@
-import type { Page } from 'playwright';
+import type { BrowserContext, Page } from 'playwright';
 import type { StoryContext } from '@storybook/csf';
 
 export type TestContext = {
@@ -7,12 +7,30 @@ export type TestContext = {
   name: string;
 };
 
+export type PrepareContext = {
+  page: Page;
+  browserContext: BrowserContext;
+  testRunnerConfig: TestRunnerConfig;
+};
+
 export type TestHook = (page: Page, context: TestContext) => Promise<void>;
+export type HttpHeaderSetter = (url: string) => Promise<Record<string, any>>;
+export type PrepareSetter = (context: PrepareContext) => Promise<void>;
 
 export interface TestRunnerConfig {
   setup?: () => void;
   preRender?: TestHook;
   postRender?: TestHook;
+  /**
+   * Adds http headers to the test-runner's requests. This is useful if you need to set headers such as `Authorization` for your Storybook instance.
+   */
+  getHttpHeaders?: HttpHeaderSetter;
+  /**
+   * Overrides the default prepare behavior of the test-runner. Good for customizing the environment before testing, such as authentication etc.
+   *
+   * If you override the default prepare behavior, even though this is powerful, you will be responsible for properly preparing the browser. Future changes to the default prepare function will not get included in your project, so you will have to keep an eye out for changes in upcoming releases.
+   */
+  prepare?: PrepareSetter;
 }
 
 export const setPreRender = (preRender: TestHook) => {

src/playwright/transformPlaywright.test.ts

@@ -95,7 +95,7 @@ describe('Playwright', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: "\${"Example/foo/bar"}/\${"A"}". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -165,7 +165,7 @@ describe('Playwright', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: "\${"Example/foo/bar"}/\${"A"}". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -236,7 +236,7 @@ describe('Playwright', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: "\${"Example/Header"}/\${"A"}". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;

src/playwright/transformPlaywright.ts

@@ -54,7 +54,7 @@ export const testPrefixer = template(
         if(err.toString().includes('Execution context was destroyed')) {
           console.log(\`An error occurred in the following story, most likely because of a navigation: "\${%%title%%}/\${%%name%%}". Retrying...\`);
           await jestPlaywright.resetPage();
-          await globalThis.__sbSetupPage(globalThis.page);
+          await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
           await testFn();
         } else {
           throw err;

src/playwright/transformPlaywrightJson.test.ts

@@ -75,7 +75,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Header\\"}/\${\\"Logged In\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -129,7 +129,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Header\\"}/\${\\"Logged Out\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -185,7 +185,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Page\\"}/\${\\"Logged In\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -266,7 +266,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Page\\"}/\${\\"Logged In\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -375,7 +375,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Header\\"}/\${\\"Logged In\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -429,7 +429,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Header\\"}/\${\\"Logged Out\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -485,7 +485,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Page\\"}/\${\\"Logged In\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;
@@ -579,7 +579,7 @@ describe('Playwright Json', () => {
                 if (err.toString().includes('Execution context was destroyed')) {
                   console.log(\`An error occurred in the following story, most likely because of a navigation: \\"\${\\"Example/Page\\"}/\${\\"Logged In\\"}\\". Retrying...\`);
                   await jestPlaywright.resetPage();
-                  await globalThis.__sbSetupPage(globalThis.page);
+                  await globalThis.__sbSetupPage(globalThis.page, globalThis.context);
                   await testFn();
                 } else {
                   throw err;

src/setup-page.ts

@@ -1,5 +1,32 @@
-import type { Page } from 'playwright';
+import type { Page, BrowserContext } from 'playwright';
 import readPackageUp from 'read-pkg-up';
+import { PrepareContext } from './playwright/hooks';
+import { getTestRunnerConfig } from './util';
+
+/**
+ * This is a default prepare function which can be overridden by the user.
+ *
+ * In order to test stories, we need to prepare the environment first.
+ * This is done by accessing the /iframe.html page of the Storybook instance. The test-runner then injects a script which prepares the environment for testing, then visits the stories.
+ */
+const defaultPrepare = async ({ page, browserContext, testRunnerConfig }: PrepareContext) => {
+  const targetURL = process.env.TARGET_URL;
+  const iframeURL = new URL('iframe.html', targetURL).toString();
+
+  if (testRunnerConfig?.getHttpHeaders) {
+    const headers = await testRunnerConfig.getHttpHeaders(iframeURL);
+    await browserContext.setExtraHTTPHeaders(headers);
+  }
+
+  await page.goto(iframeURL, { waitUntil: 'load' }).catch((err) => {
+    if (err.message?.includes('ERR_CONNECTION_REFUSED')) {
+      const errorMessage = `Could not access the Storybook instance at ${targetURL}. Are you sure it's running?\n\n${err.message}`;
+      throw new Error(errorMessage);
+    }
+
+    throw err;
+  });
+};
 
 const sanitizeURL = (url: string) => {
   let finalURL = url;
@@ -22,7 +49,7 @@ const sanitizeURL = (url: string) => {
   return finalURL;
 };
 
-export const setupPage = async (page: Page) => {
+export const setupPage = async (page: Page, browserContext: BrowserContext) => {
   const targetURL = process.env.TARGET_URL;
 
   const viewMode = process.env.VIEW_MODE || 'story';
@@ -36,20 +63,17 @@ export const setupPage = async (page: Page) => {
     : 1000;
 
   if ('TARGET_URL' in process.env && !process.env.TARGET_URL) {
-    console.log(
-      `Received TARGET_URL but with a falsy value: ${process.env.TARGET_URL}, will fallback to ${targetURL} instead.`
+    console.warn(
+      `Received TARGET_URL but with a falsy value: ${process.env.TARGET_URL}. Please fix it.`
     );
   }
 
-  const iframeURL = new URL('iframe.html', process.env.TARGET_URL).toString();
-  await page.goto(iframeURL, { waitUntil: 'load' }).catch((err) => {
-    if (err.message?.includes('ERR_CONNECTION_REFUSED')) {
-      const errorMessage = `Could not access the Storybook instance at ${targetURL}. Are you sure it's running?\n\n${err.message}`;
-      throw new Error(errorMessage);
-    }
-
-    throw err;
-  });
+  const testRunnerConfig = getTestRunnerConfig();
+  if (testRunnerConfig?.prepare) {
+    await testRunnerConfig.prepare({ page, browserContext, testRunnerConfig });
+  } else {
+    await defaultPrepare({ page, browserContext, testRunnerConfig });
+  }
 
   // if we ever want to log something from the browser to node
   await page.exposeBinding('logToPage', (_, message) => console.log(message));

src/util/getTestRunnerConfig.ts

@@ -5,7 +5,9 @@ import { TestRunnerConfig } from '../playwright/hooks';
 let testRunnerConfig: TestRunnerConfig;
 let loaded = false;
 
-export const getTestRunnerConfig = (configDir: string): TestRunnerConfig | undefined => {
+export const getTestRunnerConfig = (
+  configDir: string = process.env.STORYBOOK_CONFIG_DIR
+): TestRunnerConfig | undefined => {
   // testRunnerConfig can be undefined
   if (loaded) {
     return testRunnerConfig;