From 92696942adf3062605ff8ca8c11edb4584f75357 Mon Sep 17 00:00:00 2001
From: atanasster <atanasster@gmail.com>
Date: Tue, 19 May 2020 11:32:47 -0400
Subject: [PATCH] fix: replace onRender effect

---
 plugins/axe-plugin/src/index.tsx              | 14 ++-
 ui/blocks/src/Story/Story.tsx                 | 88 ++++++++++---------
 ui/blocks/src/context/block/ErrorBoundary.tsx |  3 +-
 .../context/components/ComponentsContext.tsx  | 12 +--
 4 files changed, 64 insertions(+), 53 deletions(-)

diff --git a/plugins/axe-plugin/src/index.tsx b/plugins/axe-plugin/src/index.tsx
index 2e360ab2d..0c111877e 100644
--- a/plugins/axe-plugin/src/index.tsx
+++ b/plugins/axe-plugin/src/index.tsx
@@ -38,6 +38,7 @@ export const AxeAllyBlock: FC<AxeAllyBlockProps> = ({
       resetTabCounter();
       runAxe(canvas)
         .then(results => {
+          console.log(results);
           setResults(results);
           isRunning.current = false;
         })
@@ -48,7 +49,9 @@ export const AxeAllyBlock: FC<AxeAllyBlockProps> = ({
     }
   };
   const onRender = () => {
-    collectResults();
+    try {
+      collectResults();
+    } catch (e) {}
   };
 
   return (
@@ -57,9 +60,12 @@ export const AxeAllyBlock: FC<AxeAllyBlockProps> = ({
         <>
           <React.Suspense fallback={<div>testing...</div>}>
             <BaseAllyBlock options={axeOptions}>
-              <React.Profiler id="axe-plugin-scan" onRender={onRender}>
-                <Story key={storyId} id={storyId} ref={storyRef} />
-              </React.Profiler>
+              <Story
+                key={storyId}
+                id={storyId}
+                ref={storyRef}
+                onRender={onRender}
+              />
             </BaseAllyBlock>
           </React.Suspense>
         </>
diff --git a/ui/blocks/src/Story/Story.tsx b/ui/blocks/src/Story/Story.tsx
index 14a34d739..388ec6bc9 100644
--- a/ui/blocks/src/Story/Story.tsx
+++ b/ui/blocks/src/Story/Story.tsx
@@ -1,5 +1,5 @@
 /** @jsx jsx */
-import { FC, createElement, forwardRef } from 'react';
+import { FC, createElement, forwardRef, useEffect } from 'react';
 import { jsx, Box } from 'theme-ui';
 import { deepMerge, StoryRenderFn } from '@component-controls/specification';
 import { getControlValues } from '@component-controls/core';
@@ -10,55 +10,59 @@ import {
 
 export type StoryProps = StoryBlockContainerProps & {
   ref?: React.Ref<HTMLDivElement>;
+  onRender?: () => void;
 };
 
 /**
  * block component to render story function with decorators
  */
 export const Story: FC<StoryProps> = forwardRef(
-  (props: StoryProps, ref: React.Ref<HTMLDivElement>) => (
-    <StoryBlockContainer {...props}>
-      {(context, rest) => {
-        const { story, options = {} } = context;
-        if (story && story.renderFn) {
-          try {
-            const values = getControlValues(story.controls);
-            const { decorators: globalDecorators = [] } = options;
-            const { decorators: storyDecorators = [] } = story;
-            const decorators = deepMerge(globalDecorators, storyDecorators);
-            const renderFn = decorators.reduce(
-              (acc: StoryRenderFn, item: StoryRenderFn) => () =>
-                item(acc, { context }),
-              //@ts-ignore
-              () => story.renderFn(values, { context }),
-            );
-            return (
-              <Box
-                id={story.id}
-                sx={{
-                  px: 4,
-                  py: 3,
-                }}
-                {...rest}
-              >
-                <div
-                  className="story-render-container"
-                  style={{ all: 'unset' }}
-                  ref={ref}
+  ({ onRender, ...props }: StoryProps, ref: React.Ref<HTMLDivElement>) => {
+    useEffect(() => onRender && onRender());
+    return (
+      <StoryBlockContainer {...props}>
+        {(context, rest) => {
+          const { story, options = {} } = context;
+          if (story && story.renderFn) {
+            try {
+              const values = getControlValues(story.controls);
+              const { decorators: globalDecorators = [] } = options;
+              const { decorators: storyDecorators = [] } = story;
+              const decorators = deepMerge(globalDecorators, storyDecorators);
+              const renderFn = decorators.reduce(
+                (acc: StoryRenderFn, item: StoryRenderFn) => () =>
+                  item(acc, { context }),
+                //@ts-ignore
+                () => story.renderFn(values, { context }),
+              );
+              return (
+                <Box
+                  id={story.id}
+                  sx={{
+                    px: 4,
+                    py: 3,
+                  }}
+                  {...rest}
                 >
-                  {createElement(renderFn)}
-                </div>
-              </Box>
-            );
-          } catch (e) {
-            console.error(e);
+                  <div
+                    className="story-render-container"
+                    style={{ all: 'unset' }}
+                    ref={ref}
+                  >
+                    {createElement(renderFn)}
+                  </div>
+                </Box>
+              );
+            } catch (e) {
+              console.error(e);
+            }
           }
-        }
-        console.error('Story function not found', props, story);
-        return null;
-      }}
-    </StoryBlockContainer>
-  ),
+          console.error('Story function not found', props, story);
+          return null;
+        }}
+      </StoryBlockContainer>
+    );
+  },
 );
 
 Story.displayName = 'Story';
diff --git a/ui/blocks/src/context/block/ErrorBoundary.tsx b/ui/blocks/src/context/block/ErrorBoundary.tsx
index 01e7318c1..9f98fc323 100644
--- a/ui/blocks/src/context/block/ErrorBoundary.tsx
+++ b/ui/blocks/src/context/block/ErrorBoundary.tsx
@@ -13,7 +13,8 @@ export class ErrorBoundary extends Component {
     const { error } = this.state;
 
     if (error) {
-      return <h1>Error caught.</h1>;
+      console.log(error);
+      return <h1>Unexpected error.</h1>;
     }
 
     return children;
diff --git a/ui/blocks/src/context/components/ComponentsContext.tsx b/ui/blocks/src/context/components/ComponentsContext.tsx
index 093abe3df..2d704d213 100644
--- a/ui/blocks/src/context/components/ComponentsContext.tsx
+++ b/ui/blocks/src/context/components/ComponentsContext.tsx
@@ -39,12 +39,6 @@ export const useComponentsContext = ({
   const [{ story, kind, component, componentPackage }, setStoryData] = useState(
     getStoryData(storyId),
   );
-
-  if (!story) {
-    return {
-      components: {},
-    };
-  }
   useEffect(() => {
     setStoryData(getStoryData(storyId));
     const onChange = () => {
@@ -56,6 +50,12 @@ export const useComponentsContext = ({
       removeObserver(onChange);
     };
   }, [storyId]);
+
+  if (!story) {
+    return {
+      components: {},
+    };
+  }
   let components: StoryComponents = {};
   if (of === CURRENT_STORY) {
     if (component) {