From 36868f5251e0c8f8e8da184dd25deec2c0fb9f83 Mon Sep 17 00:00:00 2001
From: Aruna Herath <arunabherath@gmail.com>
Date: Sun, 28 Aug 2016 20:15:32 +0530
Subject: [PATCH] Use addDecorator api to wrap stories

This requires some changes to the wrap function because not its called every time a story is mounted. Now it uses a map of knob stores instead of localKnobStores.
---
 src/index.js          | 15 ++++++++++++---
 src/stories/Button.js |  2 +-
 src/stories/index.js  |  9 +++++----
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/index.js b/src/index.js
index 81d8184a74a4..2c559a63a5a5 100644
--- a/src/index.js
+++ b/src/index.js
@@ -18,6 +18,7 @@ function register() {
 }
 
 let knobStore = {};
+const stories = {};
 
 function createKnob(name, value, type) {
   if (knobStore[name]) {
@@ -30,12 +31,20 @@ function createKnob(name, value, type) {
 
 function wrap(storyFn) {
   const channel = addons.getChannel();
-  const localKnobStore = {};
 
   return context => {
+    if (!stories[context.kind]) {
+      stories[context.kind] = {};
+    }
+
+    if (!stories[context.kind][context.story]) {
+      stories[context.kind][context.story] = {};
+    }
+
     // Change the global knobStore to the one local to this story
-    knobStore = localKnobStore;
-    return <Wrap {...{ context, storyFn, channel, store: localKnobStore }} />;
+    knobStore = stories[context.kind][context.story];
+
+    return <Wrap {...{ context, storyFn, channel, store: knobStore }} />;
   };
 }
 
diff --git a/src/stories/Button.js b/src/stories/Button.js
index 0c9af5a67424..25645c7aadfc 100644
--- a/src/stories/Button.js
+++ b/src/stories/Button.js
@@ -11,7 +11,7 @@ const buttonStyles = {
 
 const Button = ({ onClick, style, color, width, children, disabled }) => (
   <button
-    style={{ ...buttonStyles, ...style, ...{ backgroundColor: color, width: `${width}px` } }}
+    style={{ ...buttonStyles, ...{ backgroundColor: color, width: `${width}px` }, ...style }}
     onClick={onClick}
     disabled={disabled}
   >
diff --git a/src/stories/index.js b/src/stories/index.js
index ac8159e03ccb..0f198d2f79e3 100644
--- a/src/stories/index.js
+++ b/src/stories/index.js
@@ -5,15 +5,16 @@ import { createKnob, wrap } from '../index';
 import Button from './Button';
 
 storiesOf('Button', module)
-  .add('default view', wrap(() => (
+  .addDecorator((story, context) => (wrap(story)(context)))
+  .add('default view', () => (
     <Button
       onClick={ action('button clicked') }
       style={createKnob('style', { width: '70px' }, 'object')}
     >
       {createKnob('children', 'Hello')}
     </Button>
-  )))
-  .add('default view with different knobs', wrap(() => (
+  ))
+  .add('default view with different knobs', () => (
     <Button
       onClick={ action('button clicked') }
       color={createKnob('color', '#abc')}
@@ -22,7 +23,7 @@ storiesOf('Button', module)
     >
       {createKnob('text', 'Hello')}
     </Button>
-  )))
+  ))
   .add('Story without any knobs', () => (
     <Button>Hello</Button>
   ));