Merge pull request #19145 from storybookjs/fix-broken-angular-stories

Fix issue in instrumenter with `waitFor`
storybookjs · Sep 12, 2022 · bea5e2c · bea5e2c
2 parents 7b49cae + 274101c
commit bea5e2c
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 17 deletions.
diff --git a/code/examples/angular-cli/src/stories/addons/interactions/addon-interactions.stories.ts b/code/examples/angular-cli/src/stories/addons/interactions/addon-interactions.stories.ts
@@ -66,16 +66,19 @@ Submitted.play = async (context) => {
   const canvas = within(context.canvasElement);
   await userEvent.click(canvas.getByText('Submit'));
 
-  await waitFor(async () => {
-    await expect(
-      canvas.getByRole('heading', {
-        name: /you submitted the following:/i,
-      })
-    ).toBeInTheDocument();
-    await expect(canvas.getByTestId('hero-name').textContent).toEqual('Storm');
-    await expect(canvas.getByTestId('hero-alterego').textContent).toEqual('Ororo Munroe');
-    await expect(canvas.getByTestId('hero-power').textContent).toEqual('Weather Changer');
-  });
+  await waitFor(
+    async () => {
+      await expect(
+        canvas.getByRole('heading', {
+          name: /you submitted the following:/i,
+        })
+      ).toBeInTheDocument();
+      await expect(canvas.getByTestId('hero-name').textContent).toEqual('Storm');
+      await expect(canvas.getByTestId('hero-alterego').textContent).toEqual('Ororo Munroe');
+      await expect(canvas.getByTestId('hero-power').textContent).toEqual('Weather Changer');
+    },
+    { timeout: 2000 }
+  );
 };
 
 export const SubmittedAndEditedAfter: StoryFn<HeroForm> = Template.bind({});

diff --git a/code/lib/instrumenter/src/instrumenter.test.ts b/code/lib/instrumenter/src/instrumenter.test.ts
@@ -1,3 +1,4 @@
+/// <reference types="@types/jest" />;
 /* eslint-disable no-underscore-dangle */
 
 import { addons, mockChannel } from '@storybook/addons';
@@ -241,6 +242,43 @@ describe('Instrumenter', () => {
     );
   });
 
+  it('handles exceptions when making calls inside callbacks', () => {
+    const fn = (callback?: Function) => callback && callback();
+    const { fn1, fn2, fn3 } = instrument({
+      fn1: fn,
+      fn2: fn,
+      fn3: fn,
+    });
+    const error = new Error('foo');
+    let thrownError;
+    fn1(() => {
+      try {
+        fn2(() => {
+          throw error;
+        });
+      } catch (err) {
+        thrownError = err;
+      }
+      fn3();
+    });
+    expect(callSpy).toHaveBeenCalledWith(
+      expect.objectContaining({ id: 'kind--story [0] fn1', ancestors: [] })
+    );
+    expect(callSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: 'kind--story [0] fn1 [0] fn2',
+        ancestors: ['kind--story [0] fn1'],
+      })
+    );
+    expect(thrownError).toBe(error);
+    expect(callSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: 'kind--story [0] fn1 [1] fn3',
+        ancestors: ['kind--story [0] fn1'],
+      })
+    );
+  });
+
   it('tracks the parent call id for async callbacks', async () => {
     const fn = (callback?: Function) => Promise.resolve(callback && callback());
     const { fn1, fn2, fn3 } = instrument({ fn1: fn, fn2: fn, fn3: fn });
@@ -275,6 +313,43 @@ describe('Instrumenter', () => {
     );
   });
 
+  it('handles exceptions when making calls inside async callbacks', async () => {
+    const fn = (callback?: Function) => Promise.resolve(callback && callback());
+    const { fn1, fn2, fn3 } = instrument({
+      fn1: fn,
+      fn2: fn,
+      fn3: fn,
+    });
+    const error = new Error('foo');
+    let thrownError;
+    await fn1(async () => {
+      try {
+        await fn2(async () => {
+          throw error;
+        });
+      } catch (err) {
+        thrownError = err;
+      }
+      await fn3();
+    });
+    expect(callSpy).toHaveBeenCalledWith(
+      expect.objectContaining({ id: 'kind--story [0] fn1', ancestors: [] })
+    );
+    expect(callSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: 'kind--story [0] fn1 [0] fn2',
+        ancestors: ['kind--story [0] fn1'],
+      })
+    );
+    expect(thrownError).toBe(error);
+    expect(callSpy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        id: 'kind--story [0] fn1 [1] fn3',
+        ancestors: ['kind--story [0] fn1'],
+      })
+    );
+  });
+
   it('emits a "sync" event with debounce after a patched function is invoked', () => {
     const { fn } = instrument({ fn: (...args: any) => {} }, { intercept: true });
     jest.useFakeTimers();

diff --git a/code/lib/instrumenter/src/instrumenter.ts b/code/lib/instrumenter/src/instrumenter.ts
@@ -497,13 +497,19 @@ export class Instrumenter {
           this.setState(call.storyId, { cursor: 0, ancestors: [...ancestors, call.id] });
           const restore = () => this.setState(call.storyId, { cursor, ancestors });
 
-          // Invoke the actual callback function.
-          const res = arg(...args);
-
-          // Reset cursor and ancestors to their original values before we entered the callback.
-          if (res instanceof Promise) return res.then(restore, restore);
-          restore();
-          return res;
+          // Invoke the actual callback function, taking care to reset the cursor and ancestors
+          // to their original values before we entered the callback, once the callback completes.
+          let willRestore = false;
+          try {
+            const res = arg(...args);
+            if (res instanceof Promise) {
+              willRestore = true; // We need to wait for the promise to finish before restoring
+              return res.finally(restore);
+            }
+            return res;
+          } finally {
+            if (!willRestore) restore();
+          }
         };
       });