fix state & e2e tests

The state merging of useAddonState needed a fix because the merger function wasn't called with the correct data, due to react18 using async rendering, making this.state be the old state.
storybookjs · Oct 19, 2023 · 78bcb62 · 78bcb62
1 parent ff50f69
commit 78bcb62
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 34 deletions.
diff --git a/code/addons/interactions/src/Panel.tsx b/code/addons/interactions/src/Panel.tsx
@@ -148,32 +148,33 @@ export const Panel = memo<{ storyId: string }>(function PanelMemoized({ storyId
       },
       [STORY_RENDER_PHASE_CHANGED]: (event) => {
         if (event.newPhase === 'preparing') {
-          set((s) => ({
+          set({
             controlStates: INITIAL_CONTROL_STATES,
             isErrored: false,
             pausedAt: undefined,
             interactions: [],
             isPlaying: false,
-            isRerunAnimating: false,
-            scrollTarget,
-            collapsed: new Set() as Set<Call['id']>,
             hasException: false,
             caughtException: undefined,
             interactionsCount: 0,
-          }));
+          });
           return;
         }
-        set((s) => ({
-          ...s,
-          isPlaying: event.newPhase === 'playing',
-          pausedAt: undefined,
-          ...(event.newPhase === 'rendering'
-            ? {
-                isErrored: false,
-                caughtException: undefined,
-              }
-            : {}),
-        }));
+        set((s) => {
+          const newState: typeof s = {
+            ...s,
+            isPlaying: event.newPhase === 'playing',
+            pausedAt: undefined,
+            ...(event.newPhase === 'rendering'
+              ? {
+                  isErrored: false,
+                  caughtException: undefined,
+                }
+              : {}),
+          };
+
+          return newState;
+        });
       },
       [STORY_THREW_EXCEPTION]: () => {
         set((s) => ({ ...s, isErrored: true }));

diff --git a/code/e2e-tests/addon-actions.spec.ts b/code/e2e-tests/addon-actions.spec.ts
@@ -5,13 +5,10 @@ import { SbPage } from './util';
 const storybookUrl = process.env.STORYBOOK_URL || 'http://localhost:8001';
 
 test.describe('addon-actions', () => {
-  test.beforeEach(async ({ page }) => {
-    await page.goto(storybookUrl);
-    await new SbPage(page).waitUntilLoaded();
-  });
-
   test('should trigger an action', async ({ page }) => {
+    await page.goto(storybookUrl);
     const sbPage = new SbPage(page);
+    sbPage.waitUntilLoaded();
 
     await sbPage.navigateToStory('example/button', 'primary');
     const root = sbPage.previewRoot();

diff --git a/code/lib/manager-api/src/index.tsx b/code/lib/manager-api/src/index.tsx
@@ -167,7 +167,11 @@ class ManagerProvider extends Component<ManagerProviderProps, State> {
 
     const store = new Store({
       getState: () => this.state,
-      setState: (stateChange: Partial<State>, callback) => this.setState(stateChange, callback),
+      setState: (stateChange: Partial<State>, callback) => {
+        this.setState(stateChange, () => callback(this.state));
+
+        return this.state;
+      },
     });
 
     const routeData = { location, path, viewMode, singleStory, storyId, refId };
@@ -411,7 +415,9 @@ export function useSharedState<S>(stateId: string, defaultState?: S) {
   }, [quicksync]);
 
   const setState = async (s: S | API_StateMerger<S>, options?: Options) => {
-    const result = await api.setAddonState<S>(stateId, s, options);
+    await api.setAddonState<S>(stateId, s, options);
+    const result = api.getAddonState(stateId);
+
     STORYBOOK_ADDON_STATE[stateId] = result;
     return result;
   };
@@ -457,7 +463,8 @@ export function useSharedState<S>(stateId: string, defaultState?: S) {
   return [
     state,
     async (newStateOrMerger: S | API_StateMerger<S>, options?: Options) => {
-      const result = await setState(newStateOrMerger, options);
+      await setState(newStateOrMerger, options);
+      const result = api.getAddonState(stateId);
       emit(`${SHARED_STATE_CHANGED}-manager-${stateId}`, result);
     },
   ] as [S, (newStateOrMerger: S | API_StateMerger<S>, options?: Options) => void];

diff --git a/code/lib/manager-api/src/modules/addons.ts b/code/lib/manager-api/src/modules/addons.ts
@@ -133,16 +133,14 @@ export const init: ModuleFn<SubAPI, SubState> = ({ provider, store, fullAPI }) =
       newStateOrMerger: S | API_StateMerger<S>,
       options?: Options
     ): Promise<S> {
-      let nextState;
-      const { addons: existing } = store.getState();
-      if (typeof newStateOrMerger === 'function') {
-        const merger = newStateOrMerger as API_StateMerger<S>;
-        nextState = merger(api.getAddonState<S>(addonId));
-      } else {
-        nextState = newStateOrMerger;
-      }
+      const merger = (
+        typeof newStateOrMerger === 'function' ? newStateOrMerger : () => newStateOrMerger
+      ) as API_StateMerger<S>;
       return store
-        .setState({ addons: { ...existing, [addonId]: nextState } }, options)
+        .setState(
+          (s) => ({ ...s, addons: { ...s.addons, [addonId]: merger(s.addons[addonId]) } }),
+          options
+        )
         .then(() => api.getAddonState(addonId));
     },
     getAddonState: (addonId) => {

diff --git a/code/lib/manager-api/src/store.ts b/code/lib/manager-api/src/store.ts
@@ -103,7 +103,9 @@ export default class Store {
     }
 
     const newState: State = await new Promise((resolve) => {
-      this.upstreamSetState(patch, resolve);
+      this.upstreamSetState(patch, () => {
+        resolve(this.getState());
+      });
     });
 
     if (persistence !== 'none') {
@@ -114,6 +116,7 @@ export default class Store {
     if (callback) {
       callback(newState);
     }
+
     return newState;
   }
 }