Call prepare context at the right time

code/lib/preview-api/src/modules/store/csf/portable-stories.ts

@@ -110,7 +110,7 @@ export function composeStory<TRenderer extends Renderer = Renderer, TArgs extend
   const globalsFromGlobalTypes = getValuesFromArgTypes(normalizedProjectAnnotations.globalTypes);
 
   const initializeContext = () => {
-    const context: StoryContext<TRenderer> = {
+    const context: StoryContext<TRenderer> = prepareContext({
       hooks: new HooksContext(),
       globals: {
         ...globalsFromGlobalTypes,
@@ -126,10 +126,9 @@ export function composeStory<TRenderer extends Renderer = Renderer, TArgs extend
       ...story,
       context: null!,
       mount: null!,
-    };
+    });
 
     context.context = context;
-    context.mount = story.mount(context);
 
     if (story.renderToCanvas) {
       context.renderToCanvas = async () => {
@@ -156,7 +155,9 @@ export function composeStory<TRenderer extends Renderer = Renderer, TArgs extend
       };
     }
 
-    return prepareContext(context);
+    context.mount = story.mount(context);
+
+    return context;
   };
 
   let loadedContext: StoryContext<TRenderer> | undefined;

code/renderers/react/src/__test__/Button.stories.tsx

@@ -6,6 +6,7 @@ import type { ButtonProps } from './Button';
 import { Button } from './Button';
 import type { HandlerFunction } from '@storybook/addon-actions';
 import { action } from '@storybook/addon-actions';
+import { mocked } from '@storybook/test';
 
 const meta = {
   title: 'Example/Button',
@@ -134,18 +135,11 @@ export const MountInPlayFunction: CSF3Story<{ mockFn: (val: string) => string }>
   play: async ({ args, mount }) => {
     // equivalent of loaders
     const loadedData = await Promise.resolve('loaded data');
-    args.mockFn.mockReturnValueOnce('mockFn return value');
+    mocked(args.mockFn).mockReturnValueOnce(loadedData);
     // equivalent of render
-    const data = args.mockFn('render');
-    await mount(
-      <div>
-        <div data-testid="loaded-data">{loadedData}</div>
-        <div data-testid="spy-data">{String(data)}</div>
-      </div>
-    );
-
+    await mount();
     // equivalent of play
-    expect(mockFn).toHaveBeenCalledWith('render');
+    expect(mockFn).toHaveBeenCalledWith(loadedData);
   },
 };
 

code/renderers/react/src/__test__/Button.tsx

@@ -28,14 +28,14 @@ export interface ButtonProps {
  * Primary UI component for user interaction
  */
 export const Button: React.FC<ButtonProps> = (props) => {
-  const { primary = false, size = 'medium', backgroundColor, children, ...otherProps } = props;
+  const { primary = false, size = 'medium', backgroundColor, children, onClick } = props;
   const mode = primary ? 'storybook-button--primary' : 'storybook-button--secondary';
   return (
     <button
       type="button"
       className={['storybook-button', `storybook-button--${size}`, mode].join(' ')}
       style={{ backgroundColor }}
-      {...otherProps}
+      onClick={onClick}
     >
       {children}
     </button>

code/renderers/react/src/__test__/portable-stories-legacy.test.tsx

@@ -195,7 +195,7 @@ describe('Legacy Portable Stories API', () => {
   it.each(testCases)('Renders %s story', async (_storyName, Story) => {
     cleanup();
 
-    if (_storyName === 'CSF2StoryWithLocale') {
+    if (_storyName === 'CSF2StoryWithLocale' || _storyName === 'MountInPlayFunction') {
       return;
     }
 

code/renderers/react/src/__test__/portable-stories.test.tsx

@@ -135,9 +135,10 @@ describe('CSF3', () => {
     const CSF3InputFieldFilled = composeStory(stories.CSF3InputFieldFilled, stories.default);
 
     const div = document.createElement('div');
+    console.log(div.tagName);
     document.body.appendChild(div);
 
-    await CSF3InputFieldFilled.play!({ canvasElement: div });
+    await CSF3InputFieldFilled.play({ canvasElement: div });
 
     const input = screen.getByTestId('input') as HTMLInputElement;
     expect(input.value).toEqual('Hello world!');

code/renderers/react/src/portable-stories.tsx

@@ -48,7 +48,7 @@ export const INTERNAL_DEFAULT_PROJECT_ANNOTATIONS: ProjectAnnotations<ReactRende
     storyContext: { context, unboundStoryFn: Story, testingLibraryRender: render, canvasElement },
   }) => {
     if (render == null) throw new TestingLibraryMustBeConfiguredError();
-    const { unmount } = render(<Story {...context} />, { baseElement: canvasElement });
+    const { unmount } = render(<Story {...context} />, { baseElement: context.canvasElement });
     return unmount;
   },
 };