Merge pull request #5439 from storybooks/fix/background-tool

FIX the addon-background tool not appearing, or appearing only after a few navigations

addons/backgrounds/src/Tool.js

@@ -4,16 +4,42 @@ import PropTypes from 'prop-types';
 import memoize from 'memoizerific';
 
 import { logger } from '@storybook/client-logger';
-import { STORY_CHANGED } from '@storybook/core-events';
+import { SET_STORIES } from '@storybook/core-events';
 
 import { Popout, Item, Icons, Icon, IconButton, Title, Detail, List } from '@storybook/components';
 import * as S from './components';
 
 import { PARAM_KEY } from './constants';
 
-const toList = memoize(50)(viewports => Object.entries(viewports));
+const toList = memoize(50)((items = {}) => Object.entries(items));
 const getIframe = memoize(1)(() => document.getElementById('storybook-preview-background'));
 
+const getState = (props, state) => {
+  const data = props.api.getCurrentStoryData();
+  const backgrounds = data && data.parameters && data.parameters[PARAM_KEY];
+  const list = backgrounds ? toList(backgrounds) : [];
+
+  return list && list.length
+    ? list.reduce(
+        (acc, { name, value, default: isSelected }) => {
+          acc.backgrounds.push(value);
+
+          if (isSelected) {
+            acc.selected = name;
+          }
+          return acc;
+        },
+        {
+          backgrounds: [],
+          selected: state.selected,
+        }
+      )
+    : {
+        backgrounds: [],
+        selected: undefined,
+      };
+};
+
 export default class BackgroundTool extends Component {
   constructor(props) {
     super(props);
@@ -27,38 +53,16 @@ export default class BackgroundTool extends Component {
   componentDidMount() {
     const { api } = this.props;
 
-    api.on(STORY_CHANGED, this.onStoryChange);
+    api.on(SET_STORIES, () => {
+      const { state, props } = this;
+      this.setState(getState(props, state));
+    });
   }
 
-  componentWillUnmount() {
-    const { api } = this.props;
-    api.off(STORY_CHANGED, this.onStoryChange);
+  static getDerivedStateFromProps(props, state) {
+    return getState(props, state);
   }
 
-  onStoryChange = id => {
-    const { api } = this.props;
-    const params = api.getParameters(id, PARAM_KEY);
-
-    if (params && !params.disable) {
-      const { backgrounds, selected } = params.reduce(
-        (acc, { name, value, default: isSelected }) => {
-          Object.assign(acc.backgrounds, { [name]: value });
-          if (isSelected) {
-            acc.selected = name;
-          }
-          return acc;
-        },
-        {
-          backgrounds: {},
-          selected: undefined,
-        }
-      );
-      this.setState({ backgrounds, selected }, this.apply);
-    } else {
-      this.setState({ backgrounds: {}, selected: undefined }, this.apply);
-    }
-  };
-
   apply = () => {
     const iframe = getIframe();
     if (iframe) {
@@ -85,13 +89,8 @@ export default class BackgroundTool extends Component {
 
   render() {
     const { backgrounds, selected } = this.state;
-    const list = toList(backgrounds);
 
-    if (!list.length) {
-      return null;
-    }
-
-    return (
+    return backgrounds.length ? (
       <Popout key="backgrounds">
         <IconButton key="background" title="Backgrounds">
           <Icons icon="photo" />
@@ -114,7 +113,7 @@ export default class BackgroundTool extends Component {
               </Fragment>
             ) : null}
 
-            {list.map(([key, value]) => (
+            {backgrounds.map(([key, value]) => (
               <Item
                 key={key}
                 onClick={() => {
@@ -130,7 +129,7 @@ export default class BackgroundTool extends Component {
           </List>
         )}
       </Popout>
-    );
+    ) : null;
   }
 }
 BackgroundTool.propTypes = {

lib/ui/src/core/stories.js

@@ -28,6 +28,25 @@ const initStoriesApi = ({
   storyId: initialStoryId,
   viewMode: initialViewMode,
 }) => {
+  const getData = storyId => {
+    const { storiesHash } = store.getState();
+
+    return storiesHash[storyId];
+  };
+  const getCurrentStoryData = () => {
+    const { storyId } = store.getState();
+
+    return getData(storyId);
+  };
+  const getParameters = (storyId, addon) => {
+    const data = getData(storyId);
+    if (!data) {
+      return null;
+    }
+    const { parameters } = data;
+    return addon ? parameters[addon] : parameters;
+  };
+
   const jumpToStory = direction => {
     const { storiesHash, viewMode, storyId } = store.getState();
 
@@ -59,21 +78,6 @@ const initStoriesApi = ({
     }
   };
 
-  const getData = storyId => {
-    const { storiesHash } = store.getState();
-
-    return storiesHash[storyId];
-  };
-
-  const getParameters = (storyId, addon) => {
-    const data = getData(storyId);
-    if (!data) {
-      return null;
-    }
-    const { parameters } = data;
-    return addon ? parameters[addon] : parameters;
-  };
-
   const jumpToComponent = direction => {
     const state = store.getState();
     const { storiesHash, viewMode, storyId } = state;
@@ -209,6 +213,7 @@ const initStoriesApi = ({
     api: {
       storyId: toId,
       selectStory,
+      getCurrentStoryData,
       setStories,
       jumpToComponent,
       jumpToStory,