From e59fab57e2142ec89e22b476cfb251f9a74e31b0 Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Fri, 5 Jul 2019 15:51:56 +0800
Subject: [PATCH] Core: Update story_store to keep _data and _legacydata in
 sync

---
 lib/client-api/src/story_store.js      | 13 ++++++++++
 lib/client-api/src/story_store.test.js | 35 ++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/lib/client-api/src/story_store.js b/lib/client-api/src/story_store.js
index b0e8abb9906..df3b813cdcf 100644
--- a/lib/client-api/src/story_store.js
+++ b/lib/client-api/src/story_store.js
@@ -114,7 +114,13 @@ export default class StoryStore extends EventEmitter {
 
   remove = id => {
     const { _data } = this;
+    const { kind, name } = _data[id];
     delete _data[id];
+
+    const kindData = this._legacydata[toKey(kind)];
+    if (kindData) {
+      delete kindData.stories[toKey(name)];
+    }
   };
 
   addStory(
@@ -271,6 +277,13 @@ export default class StoryStore extends EventEmitter {
   removeStoryKind(kind) {
     if (this.hasStoryKind(kind)) {
       this._legacydata[toKey(kind)].stories = {};
+
+      this._data = Object.entries(this._data).reduce((acc, [id, story]) => {
+        if (story.kind !== kind) {
+          Object.assign(acc, { [id]: story });
+        }
+        return acc;
+      }, {});
     }
   }
 
diff --git a/lib/client-api/src/story_store.test.js b/lib/client-api/src/story_store.test.js
index f27cf64cc24..1ac523c4bbe 100644
--- a/lib/client-api/src/story_store.test.js
+++ b/lib/client-api/src/story_store.test.js
@@ -95,6 +95,41 @@ describe('preview.story_store', () => {
       const store = new StoryStore({ channel });
       store.removeStoryKind('kind');
     });
+    it('should remove the kind in both modern and legacy APIs', () => {
+      const store = new StoryStore({ channel });
+      store.addStory(...make('kind-1', 'story-1.1', () => 0));
+      store.addStory(...make('kind-1', 'story-1.2', () => 0));
+      store.addStory(...make('kind-2', 'story-2.1', () => 0));
+      store.addStory(...make('kind-2', 'story-2.2', () => 0));
+
+      store.removeStoryKind('kind-1');
+
+      // _legacydata
+      expect(store.hasStory('kind-1', 'story-1.1')).toBeFalsy();
+      expect(store.hasStory('kind-2', 'story-2.1')).toBeTruthy();
+
+      // _data
+      expect(store.fromId(toId('kind-1', 'story-1.1'))).toBeFalsy();
+      expect(store.fromId(toId('kind-2', 'story-2.1'))).toBeTruthy();
+    });
+  });
+
+  describe('remove', () => {
+    it('should remove the kind in both modern and legacy APIs', () => {
+      const store = new StoryStore({ channel });
+      store.addStory(...make('kind-1', 'story-1.1', () => 0));
+      store.addStory(...make('kind-1', 'story-1.2', () => 0));
+
+      store.remove(toId('kind-1', 'story-1.1'));
+
+      // _legacydata
+      expect(store.hasStory('kind-1', 'story-1.1')).toBeFalsy();
+      expect(store.hasStory('kind-1', 'story-1.2')).toBeTruthy();
+
+      // _data
+      expect(store.fromId(toId('kind-1', 'story-1.1'))).toBeFalsy();
+      expect(store.fromId(toId('kind-1', 'story-1.2'))).toBeTruthy();
+    });
   });
 
   describe('getStoryAndParameters', () => {