Merge pull request #12460 from StephanBijzitter/storyshots-puppeteer-…

…element Storyshots: Allow taking a screenshot of just a specific element
storybookjs · Sep 12, 2020 · b15b04e · b15b04e
2 parents 6c4b5d4 + e48ad84
commit b15b04e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/addons/storyshots/storyshots-puppeteer/README.md b/addons/storyshots/storyshots-puppeteer/README.md
@@ -305,3 +305,16 @@ initStoryshots({
 ```
 
 `getScreenshotOptions` receives an object `{ context: {kind, story}, url}`. _kind_ is the kind of the story and the _story_ its name. _url_ is the URL the browser will use to screenshot.
+
+To create a screenshot of just a single element (with its children), rather than the page or current viewport, an ElementHandle can be returned from `beforeScreenshot`:
+```js
+import initStoryshots from '@storybook/addon-storyshots';
+import { imageSnapshot } from '@storybook/addon-storyshots-puppeteer';
+
+const beforeScreenshot = (page) => page.$('#root > *');
+
+initStoryshots({
+  suite: 'Image storyshots',
+  test: imageSnapshot({ storybookUrl: 'http://localhost:6006', beforeScreenshot }),
+});
+```
diff --git a/addons/storyshots/storyshots-puppeteer/src/config.ts b/addons/storyshots/storyshots-puppeteer/src/config.ts
@@ -1,5 +1,5 @@
 import { MatchImageSnapshotOptions } from 'jest-image-snapshot';
-import { Base64ScreenShotOptions, Browser, DirectNavigationOptions, Page } from 'puppeteer';
+import { Base64ScreenShotOptions, Browser, DirectNavigationOptions, Page, ElementHandle } from 'puppeteer';
 
 export interface Context {
   kind: string;
@@ -33,7 +33,7 @@ export interface PuppeteerTestConfig extends CommonConfig {
 export interface ImageSnapshotConfig extends CommonConfig {
   getMatchOptions: (options: Options) => MatchImageSnapshotOptions;
   getScreenshotOptions: (options: Options) => Base64ScreenShotOptions;
-  beforeScreenshot: (page: Page, options: Options) => void;
+  beforeScreenshot: (page: Page, options: Options) => void | ElementHandle;
   afterScreenshot: (options: { image: string; context: Context }) => void;
 }
 

diff --git a/addons/storyshots/storyshots-puppeteer/src/imageSnapshot.ts b/addons/storyshots/storyshots-puppeteer/src/imageSnapshot.ts
@@ -12,8 +12,8 @@ export const imageSnapshot = (customConfig: Partial<ImageSnapshotConfig> = {}) =
     ...config,
     async testBody(page, options) {
       expect.assertions(1);
-      await beforeScreenshot(page, options);
-      const image = await page.screenshot(getScreenshotOptions(options));
+      const element = await beforeScreenshot(page, options);
+      const image = await (element || page).screenshot(getScreenshotOptions(options));
       await afterScreenshot({ image, context: options.context });
       expect(image).toMatchImageSnapshot(getMatchOptions(options));
     },