From a784aa1b9b1335880db027791cd58a882c8108d5 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Tue, 11 Feb 2020 20:56:13 +1100
Subject: [PATCH] Listen to `CHANGE_STORY_STATE` and change it in the store

---
 lib/client-api/src/story_store.test.ts | 14 ++++++++++++--
 lib/client-api/src/story_store.ts      |  6 +++++-
 lib/core-events/src/index.ts           |  2 ++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/lib/client-api/src/story_store.test.ts b/lib/client-api/src/story_store.test.ts
index 613e47b42fa..303880b2050 100644
--- a/lib/client-api/src/story_store.test.ts
+++ b/lib/client-api/src/story_store.test.ts
@@ -1,6 +1,6 @@
 import createChannel from '@storybook/channel-postmessage';
 import { toId } from '@storybook/csf';
-import addons from '@storybook/addons';
+import addons, { mockChannel } from '@storybook/addons';
 import Events from '@storybook/core-events';
 
 import StoryStore from './story_store';
@@ -96,7 +96,7 @@ describe('preview.story_store', () => {
 
     it('synchronously emits STORY_STATE_CHANGED if different', () => {
       const onStateChanged = jest.fn();
-      const testChannel = createChannel({ page: 'preview' });
+      const testChannel = mockChannel();
       testChannel.on(Events.STORY_STATE_CHANGED, onStateChanged);
 
       const store = new StoryStore({ channel: testChannel });
@@ -108,6 +108,16 @@ describe('preview.story_store', () => {
       store.setStoryState('a--1', { baz: 'bing' });
       expect(onStateChanged).toHaveBeenCalledWith('a--1', { foo: 'bar', baz: 'bing' });
     });
+
+    it('should update if the CHANGE_STORY_STATE event is received', () => {
+      const testChannel = mockChannel();
+      const store = new StoryStore({ channel: testChannel });
+      addStoryToStore(store, 'a', '1', () => 0);
+
+      testChannel.emit(Events.CHANGE_STORY_STATE, 'a--1', { foo: 'bar' });
+
+      expect(store.getRawStory('a', '1').state).toEqual({ foo: 'bar' });
+    });
   });
 
   describe('storySort', () => {
diff --git a/lib/client-api/src/story_store.ts b/lib/client-api/src/story_store.ts
index 263bc6e1d2b..806da07e915 100644
--- a/lib/client-api/src/story_store.ts
+++ b/lib/client-api/src/story_store.ts
@@ -84,12 +84,16 @@ export default class StoryStore extends EventEmitter {
     this._stories = {};
     this._revision = 0;
     this._selection = {} as any;
-    this._channel = params.channel;
     this._error = undefined;
+
+    if (params.channel) this.setChannel(params.channel);
   }
 
   setChannel = (channel: Channel) => {
     this._channel = channel;
+    channel.on(Events.CHANGE_STORY_STATE, (id: string, newState: StoryState) =>
+      this.setStoryState(id, newState)
+    );
   };
 
   addGlobalMetadata({ parameters, decorators }: StoryMetadata) {
diff --git a/lib/core-events/src/index.ts b/lib/core-events/src/index.ts
index a7a2ace516d..51ebca850ac 100644
--- a/lib/core-events/src/index.ts
+++ b/lib/core-events/src/index.ts
@@ -10,6 +10,7 @@ enum events {
   STORY_ADDED = 'storyAdded',
   STORY_CHANGED = 'storyChanged',
   STORY_UNCHANGED = 'storyUnchanged',
+  CHANGE_STORY_STATE = 'changeStoryState',
   STORY_STATE_CHANGED = 'storyStateChanged',
   FORCE_RE_RENDER = 'forceReRender',
   REGISTER_SUBSCRIPTION = 'registerSubscription',
@@ -45,6 +46,7 @@ export const {
   REGISTER_SUBSCRIPTION,
   STORY_INIT,
   STORY_ADDED,
+  CHANGE_STORY_STATE,
   STORY_STATE_CHANGED,
   STORY_RENDER,
   STORY_RENDERED,