From 840731a0fcda6ef5fb356373503a9650efe75e5c Mon Sep 17 00:00:00 2001
From: Muhammed Thanish <mnmtanish@gmail.com>
Date: Sat, 16 Apr 2016 19:59:23 +0530
Subject: [PATCH] Support story decorators

---
 dist/client/client_api.js          | 24 +++++++++++++++++++++---
 src/client/__tests__/client_api.js | 23 +++++++++++++++++++++++
 src/client/client_api.js           | 22 +++++++++++++++++++---
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/dist/client/client_api.js b/dist/client/client_api.js
index 0eb2d67568e5..85d0103980aa 100644
--- a/dist/client/client_api.js
+++ b/dist/client/client_api.js
@@ -39,12 +39,30 @@ var ClientApi = function () {
         });
       }
 
-      var add = function add(storyName, fn) {
+      var decorators = [];
+      var api = {};
+
+      api.add = function (storyName, getStory) {
+        // Wrap the getStory function with each decorator. The first
+        // decorator will wrap the story function. The second will
+        // wrap the first decorator and so on.
+        var fn = decorators.reduce(function (decorated, decorator) {
+          return function () {
+            return decorator(decorated);
+          };
+        }, getStory);
+
+        // Add the fully decorated getStory function.
         _this._storyStore.addStory(kind, storyName, fn);
-        return { add: add };
+        return api;
       };
 
-      return { add: add };
+      api.addDecorator = function (decorator) {
+        decorators.push(decorator);
+        return api;
+      };
+
+      return api;
     }
   }, {
     key: 'action',
diff --git a/src/client/__tests__/client_api.js b/src/client/__tests__/client_api.js
index 38713d9407b6..cb67a7a1414a 100644
--- a/src/client/__tests__/client_api.js
+++ b/src/client/__tests__/client_api.js
@@ -74,6 +74,29 @@ describe('client.ClientApi', () => {
     });
   });
 
+  describe('storiesOf.addDecorator', () => {
+    it('should wrap a given story', () => {
+      const api = getClientApi();
+      const handle = () => ['h'];
+      const decorate1 = fn => fn().concat('d1');
+      const decorate2 = fn => fn().concat('d2');
+
+      api._storyStore.addStory = sinon.stub();
+      api.storiesOf('kind')
+        .addDecorator(decorate1)
+        .addDecorator(decorate2)
+        .add('name', handle);
+
+      const args = api._storyStore.addStory.args[0];
+      expect(args[0]).to.be.equal('kind');
+      expect(args[1]).to.be.equal('name');
+      expect(args[2]).to.be.a('function');
+
+      const decoratedFn = args[2];
+      expect(decoratedFn()).to.deep.equal(['h', 'd1', 'd2']);
+    });
+  });
+
   describe('action', () => {
     it('should send action info to the syncedStore', () => {
       const api = getClientApi();
diff --git a/src/client/client_api.js b/src/client/client_api.js
index 2cf80fbc2ea6..2f623a4260f0 100644
--- a/src/client/client_api.js
+++ b/src/client/client_api.js
@@ -11,12 +11,28 @@ export default class ClientApi {
       });
     }
 
-    const add = (storyName, fn) => {
+    const decorators = [];
+    const api = {};
+
+    api.add = (storyName, getStory) => {
+      // Wrap the getStory function with each decorator. The first
+      // decorator will wrap the story function. The second will
+      // wrap the first decorator and so on.
+      const fn = decorators.reduce((decorated, decorator) => {
+        return () => decorator(decorated);
+      }, getStory);
+
+      // Add the fully decorated getStory function.
       this._storyStore.addStory(kind, storyName, fn);
-      return { add };
+      return api;
+    };
+
+    api.addDecorator = decorator => {
+      decorators.push(decorator);
+      return api;
     };
 
-    return { add };
+    return api;
   }
 
   action(name) {