From bb0801069b2c6e3b97676ae282ed11c4678057a1 Mon Sep 17 00:00:00 2001
From: Matthew Brookes <github@nospam.33m.co>
Date: Thu, 17 Oct 2019 22:17:01 +0100
Subject: [PATCH] Only strip outer div, fix and extend tests

---
 docs/src/modules/components/Demo.js          |  2 +-
 docs/src/modules/utils/getJsxPreview.js      | 13 ++--
 docs/src/modules/utils/getJsxPreview.test.js | 65 +++++++++++++++++++-
 3 files changed, 69 insertions(+), 11 deletions(-)

diff --git a/docs/src/modules/components/Demo.js b/docs/src/modules/components/Demo.js
index 7e1fe455354666..1df8ca92ad480d 100644
--- a/docs/src/modules/components/Demo.js
+++ b/docs/src/modules/components/Demo.js
@@ -266,7 +266,7 @@ function Demo(props) {
 
   const match = useMediaQuery(theme => theme.breakpoints.up('sm'));
 
-  const jsx = getJsxPreview(demoData.raw || '', demoOptions);
+  const jsx = getJsxPreview(demoData.raw || '', demoOptions.defaultCodeOpen);
   const showPreview =
     !demoOptions.hideHeader && jsx !== demoData.raw && jsx.split(/\n/).length <= 20;
 
diff --git a/docs/src/modules/utils/getJsxPreview.js b/docs/src/modules/utils/getJsxPreview.js
index 2bad66a602dcba..dcc46b9176e457 100644
--- a/docs/src/modules/utils/getJsxPreview.js
+++ b/docs/src/modules/utils/getJsxPreview.js
@@ -1,23 +1,20 @@
-// docs/pages/components/material-icons.js doesn't provide the source
-export default function getJsxPreview(code, demoOptions) {
+export default function getJsxPreview(code, defaultCodeOpen) {
   /* The regex matches the content of the return statement in the default export,
    * stripping any wrapper divs:
    *
    * `export default.*`
    * `\n  return (\n` or `\n  return `
    * `    <div.*>\n` (optional)
-   * `      <div.*>\n` (optional)
    *  everything until:
-   * `\n    </div>` (optional)
    * `\n  </div>` (optional)
    * `  );\n}` or `;\n}`
    */
   let jsx = code.match(
-    /export default .*(?:\n {2}return \(\n|\n {2}return )(?: {4}<div.*?>\n)?(?: {6}<div.*?>\n)?(.*?)(\n {6}<\/div>)?(\n {4}<\/div>)?(\n {2}\);\n}|;\n})/s,
+    /export default .*(?:\n {2}return \(\n|\n {2}return )(?: {4}<div.*?>\n)?(.*?)(\n {4}<\/div>)?(\n {2}\);\n}|;\n})/s,
   );
-  /* Just the match, otherwise the full source if no match or disabled,
-so as not to break the Collapse transition. */
-  jsx = jsx && demoOptions.defaultCodeOpen !== false ? jsx[1] : code;
+  // Just the match, otherwise the full source if either no match or preview disabled,
+  // so as not to break the Collapse transition.
+  jsx = jsx && defaultCodeOpen !== false ? jsx[1] : code;
 
   // Remove leading spaces from each line
   return jsx.split(/\n/).reduce(
diff --git a/docs/src/modules/utils/getJsxPreview.test.js b/docs/src/modules/utils/getJsxPreview.test.js
index 0a74ac9317e59f..cc72b965c6872c 100644
--- a/docs/src/modules/utils/getJsxPreview.test.js
+++ b/docs/src/modules/utils/getJsxPreview.test.js
@@ -2,7 +2,7 @@ import { expect } from 'chai';
 import getJsxPreview from './getJsxPreview';
 
 describe('getJsxPreview', () => {
-  it('should extract a preview', () => {
+  it('should extract one-line a preview', () => {
     expect(
       getJsxPreview(`
 import React from 'react';
@@ -11,8 +11,69 @@ import Rating from '@material-ui/lab/Rating';
 export default function HalfRating() {
   return <Rating name="half-rating" value={2.5} precision={0.5} />;
 }
-    `),
+`, true),
     ).to.equal(`<Rating name="half-rating" value={2.5} precision={0.5} />
+`);
+  });
+
+  it('should extract a multi-line preview', () => {
+    expect(
+      getJsxPreview(`
+export default function UseWidth() {
+  return (
+    <ThemeProvider theme={theme}>
+      <MyComponent />
+    </ThemeProvider>
+  );
+}
+`, true),
+    ).to.equal(`<ThemeProvider theme={theme}>
+  <MyComponent />
+</ThemeProvider>
+`);
+  });
+
+  it('should exclude an outer div', () => {
+    expect(
+      getJsxPreview(`
+export default function UseWidth() {
+  return (
+    <div className={classes.root}>
+      <MyComponent />
+    </div>
+  );
+}
+`, true),
+    ).to.equal(`<MyComponent />
+`);
+  });
+
+  it('should return all if no match', () => {
+    expect(
+      getJsxPreview(`
+export function UseWidth() {
+  return ( <MyComponent />;
+}
+`, true),
+    ).to.equal(`
+export function UseWidth() {
+  return ( <MyComponent />;
+}
+
+`);
+  });
+  it('should return all if defaultCodeOpen is false', () => {
+    expect(
+      getJsxPreview(`
+export default function UseWidth() {
+  return <MyComponent />;
+}
+`, false),
+    ).to.equal(`
+export default function UseWidth() {
+  return <MyComponent />;
+}
+
 `);
   });
 });