From cf38bf7d43aee40a18be3c66dda2d7af4ac9f990 Mon Sep 17 00:00:00 2001
From: Arunoda Susiripala <arunoda.susiripala@gmail.com>
Date: Fri, 9 Sep 2016 09:12:40 +0530
Subject: [PATCH] Make the wrapStory decorator simple to use.

---
 example/stories/index.js |  2 +-
 src/KnobManager.js       | 20 +++++++++-----------
 src/index.js             |  4 ++--
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/example/stories/index.js b/example/stories/index.js
index 5b981f918f4..5dda111f416 100644
--- a/example/stories/index.js
+++ b/example/stories/index.js
@@ -12,7 +12,7 @@ import {
 import Button from './Button';
 
 storiesOf('Button', module)
-  .addDecorator((story, context) => (withKnobs(story)(context)))
+  .addDecorator(withKnobs)
   .add('default view', () => (
     <Button
       onClick={ action('button clicked') }
diff --git a/src/KnobManager.js b/src/KnobManager.js
index 8ba091f3c8a..e02669d84a8 100644
--- a/src/KnobManager.js
+++ b/src/KnobManager.js
@@ -30,18 +30,16 @@ export default class KnobManager {
     return knobStore.get(name).value;
   }
 
-  wrapStory(storyFn, channel) {
-    return context => {
-      const key = `${context.kind}:::${context.story}`;
-      let knobStore = this.knobStoreMap[key];
+  wrapStory(channel, storyFn, context) {
+    const key = `${context.kind}:::${context.story}`;
+    let knobStore = this.knobStoreMap[key];
 
-      if (!knobStore) {
-        knobStore = this.knobStoreMap[key] = new KnobStore();
-      }
-      this.knobStore = knobStore;
+    if (!knobStore) {
+      knobStore = this.knobStoreMap[key] = new KnobStore();
+    }
+    this.knobStore = knobStore;
 
-      const props = { context, storyFn, channel, knobStore };
-      return (<WrapStory {...props} />);
-    };
+    const props = { context, storyFn, channel, knobStore };
+    return (<WrapStory {...props} />);
   }
 }
diff --git a/src/index.js b/src/index.js
index 96b76b207b4..42730340575 100644
--- a/src/index.js
+++ b/src/index.js
@@ -23,7 +23,7 @@ export function object(name, value) {
   return manager.knob(name, { type: 'object', value });
 }
 
-export function withKnobs(storyFn) {
+export function withKnobs(storyFn, context) {
   const channel = addons.getChannel();
-  return manager.wrapStory(storyFn, channel);
+  return manager.wrapStory(channel, storyFn, context);
 }