From e48ad84e1b1f9872a5e0c31e35583341aa37d684 Mon Sep 17 00:00:00 2001
From: StephanBijzitter <stephanbijzitter@gmail.com>
Date: Fri, 11 Sep 2020 20:42:49 +0200
Subject: [PATCH] Allow taking a screenshot of just a specific element

---
 addons/storyshots/storyshots-puppeteer/README.md    | 13 +++++++++++++
 .../storyshots/storyshots-puppeteer/src/config.ts   |  4 ++--
 .../storyshots-puppeteer/src/imageSnapshot.ts       |  4 ++--
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/addons/storyshots/storyshots-puppeteer/README.md b/addons/storyshots/storyshots-puppeteer/README.md
index acb9eb79e3cb..4776f0d23b00 100644
--- 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 }),
+});
+```
\ No newline at end of file
diff --git a/addons/storyshots/storyshots-puppeteer/src/config.ts b/addons/storyshots/storyshots-puppeteer/src/config.ts
index a421ea029992..c8f6792593c5 100644
--- 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
index e8cceac61523..d2e950401f7a 100644
--- 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));
     },