Core: Move prepareContext into store.getStoryContext()

code/frameworks/angular/src/client/docs/sourceDecorator.ts

@@ -38,8 +38,8 @@ export const sourceDecorator = (
 
   useEffect(() => {
     if (toEmit) {
-      const { id, args } = context;
-      channel.emit(SNIPPET_RENDERED, { id, args, source: toEmit, format: 'angular' });
+      const { id, unmappedArgs } = context;
+      channel.emit(SNIPPET_RENDERED, { id, args: unmappedArgs, source: toEmit, format: 'angular' });
     }
   });
 

code/lib/preview-api/src/modules/preview-web/PreviewWithSelection.tsx

@@ -380,7 +380,7 @@ export class PreviewWithSelection<TFramework extends Renderer> extends Preview<T
 
     if (isStoryRender(render)) {
       if (!render.story) throw new Error('Render has not been prepared!');
-      const { parameters, initialArgs, argTypes, args } = this.storyStore.getStoryContext(
+      const { parameters, initialArgs, argTypes, unmappedArgs } = this.storyStore.getStoryContext(
         render.story
       );
 
@@ -390,15 +390,15 @@ export class PreviewWithSelection<TFramework extends Renderer> extends Preview<T
           parameters,
           initialArgs,
           argTypes,
-          args,
+          args: unmappedArgs,
         });
       }
 
       // For v6 mode / compatibility
       // If the implementation changed, or args were persisted, the args may have changed,
       // and the STORY_PREPARED event above may not be respected.
       if (implementationChanged || persistedArgs) {
-        this.channel.emit(STORY_ARGS_UPDATED, { storyId, args });
+        this.channel.emit(STORY_ARGS_UPDATED, { storyId, args: unmappedArgs });
       }
     } else if (global.FEATURES?.storyStoreV7) {
       if (!this.storyStore.projectAnnotations) throw new Error('Store not initialized');

code/lib/preview-api/src/modules/preview-web/render/StoryRender.test.ts

@@ -103,48 +103,4 @@ describe('StoryRender', () => {
     await render.renderToElement({} as any);
     expect(story.playFunction).not.toHaveBeenCalled();
   });
-
-  it('passes the initialArgs to loaders and render function if forceInitialArgs is true', async () => {
-    const story = {
-      id: 'id',
-      title: 'title',
-      name: 'name',
-      tags: [],
-      initialArgs: { a: 'b' },
-      applyLoaders: jest.fn(),
-      unboundStoryFn: jest.fn(),
-      playFunction: jest.fn(),
-      prepareContext: jest.fn((ctx) => ctx),
-    };
-
-    const renderToScreen = jest.fn();
-
-    const render = new StoryRender(
-      new Channel(),
-      { getStoryContext: () => ({ args: { a: 'c ' } }) } as any,
-      renderToScreen as any,
-      {} as any,
-      entry.id,
-      'story',
-      { forceInitialArgs: true },
-      story as any
-    );
-
-    await render.renderToElement({} as any);
-
-    expect(story.applyLoaders).toHaveBeenCalledWith(
-      expect.objectContaining({
-        args: { a: 'b' },
-      })
-    );
-
-    expect(renderToScreen).toHaveBeenCalledWith(
-      expect.objectContaining({
-        storyContext: expect.objectContaining({
-          args: { a: 'b' },
-        }),
-      }),
-      expect.any(Object)
-    );
-  });
 });

code/lib/preview-api/src/modules/preview-web/render/StoryRender.ts

@@ -136,7 +136,8 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
 
   private storyContext() {
     if (!this.story) throw new Error(`Cannot call storyContext before preparing`);
-    return this.store.getStoryContext(this.story);
+    const { forceInitialArgs } = this.renderOptions;
+    return this.store.getStoryContext(this.story, { forceInitialArgs });
   }
 
   async render({
@@ -150,18 +151,8 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     if (!this.story) throw new Error('cannot render when not prepared');
     if (!canvasElement) throw new Error('cannot render when canvasElement is unset');
 
-    const {
-      id,
-      componentId,
-      title,
-      name,
-      tags,
-      applyLoaders,
-      unboundStoryFn,
-      playFunction,
-      prepareContext,
-      initialArgs,
-    } = this.story;
+    const { id, componentId, title, name, tags, applyLoaders, unboundStoryFn, playFunction } =
+      this.story;
 
     if (forceRemount && !initial) {
       // NOTE: we don't check the cancel actually worked here, so the previous
@@ -176,16 +167,10 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     const abortSignal = (this.abortController as AbortController).signal;
 
     try {
-      const getCurrentContext = () =>
-        prepareContext({
-          ...this.storyContext(),
-          ...(this.renderOptions.forceInitialArgs && { args: initialArgs }),
-        } as StoryContext);
-
       let loadedContext: Awaited<ReturnType<typeof applyLoaders>>;
       await this.runPhase(abortSignal, 'loading', async () => {
         loadedContext = await applyLoaders({
-          ...getCurrentContext(),
+          ...this.storyContext(),
           viewMode: this.viewMode,
         } as StoryContextForLoaders<TRenderer>);
       });
@@ -197,7 +182,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
         ...loadedContext!,
         // By this stage, it is possible that new args/globals have been received for this story
         // and we need to ensure we render it with the new values
-        ...getCurrentContext(),
+        ...this.storyContext(),
         abortSignal,
         // We should consider parameterizing the story types with TRenderer['canvasElement'] in the future
         canvasElement: canvasElement as any,

code/lib/preview-api/src/modules/store/StoryStore.test.ts

@@ -8,6 +8,7 @@ import type { HooksContext } from './hooks';
 
 // Spy on prepareStory/processCSFFile
 jest.mock('./csf/prepareStory', () => ({
+  ...jest.requireActual('./csf/prepareStory'),
   prepareStory: jest.fn(jest.requireActual('./csf/prepareStory').prepareStory),
 }));
 jest.mock('./csf/processCSFFile', () => ({
@@ -425,6 +426,20 @@ describe('StoryStore', () => {
       });
     });
 
+    it('can force initial args', async () => {
+      const store = new StoryStore();
+      store.setProjectAnnotations(projectAnnotations);
+      store.initialize({ storyIndex, importFn, cache: false });
+
+      const story = await store.loadStory({ storyId: 'component-one--a' });
+
+      store.args.update(story.id, { foo: 'bar' });
+
+      expect(store.getStoryContext(story, { forceInitialArgs: true })).toMatchObject({
+        args: { foo: 'a' },
+      });
+    });
+
     it('returns the same hooks each time', async () => {
       const store = new StoryStore();
       store.setProjectAnnotations(projectAnnotations);
@@ -735,7 +750,6 @@ describe('StoryStore', () => {
               "fileName": "./src/ComponentOne.stories.js",
             },
             "playFunction": undefined,
-            "prepareContext": [Function],
             "story": "A",
             "storyFn": [Function],
             "subcomponents": undefined,
@@ -781,7 +795,6 @@ describe('StoryStore', () => {
               "fileName": "./src/ComponentOne.stories.js",
             },
             "playFunction": undefined,
-            "prepareContext": [Function],
             "story": "B",
             "storyFn": [Function],
             "subcomponents": undefined,
@@ -827,7 +840,6 @@ describe('StoryStore', () => {
               "fileName": "./src/ComponentTwo.stories.js",
             },
             "playFunction": undefined,
-            "prepareContext": [Function],
             "story": "C",
             "storyFn": [Function],
             "subcomponents": undefined,

code/lib/preview-api/src/modules/store/StoryStore.ts

@@ -30,7 +30,13 @@ import { HooksContext } from '../addons';
 import { StoryIndexStore } from './StoryIndexStore';
 import { ArgsStore } from './ArgsStore';
 import { GlobalsStore } from './GlobalsStore';
-import { processCSFFile, prepareStory, prepareMeta, normalizeProjectAnnotations } from './csf';
+import {
+  processCSFFile,
+  prepareStory,
+  prepareMeta,
+  normalizeProjectAnnotations,
+  prepareContext,
+} from './csf';
 
 const CSF_CACHE_SIZE = 1000;
 const STORY_CACHE_SIZE = 10000;
@@ -276,16 +282,17 @@ export class StoryStore<TRenderer extends Renderer> {
   // A prepared story does not include args, globals or hooks. These are stored in the story store
   // and updated separtely to the (immutable) story.
   getStoryContext(
-    story: PreparedStory<TRenderer>
-  ): Omit<StoryContextForLoaders<TRenderer>, 'viewMode'> {
+    story: PreparedStory<TRenderer>,
+    { forceInitialArgs = false } = {}
+  ): Omit<StoryContextForLoaders, 'viewMode'> {
     if (!this.globals) throw new Error(`getStoryContext called before initialization`);
 
-    return {
+    return prepareContext({
       ...story,
-      args: this.args.get(story.id),
+      args: forceInitialArgs ? story.initialArgs : this.args.get(story.id),
       globals: this.globals.get(),
       hooks: this.hooks[story.id] as unknown,
-    };
+    });
   }
 
   cleanupStory(story: PreparedStory<TRenderer>): void {

code/lib/preview-api/src/modules/store/args.test.ts

@@ -276,8 +276,12 @@ describe('groupArgsByTarget', () => {
   it('groups targeted args', () => {
     const groups = groupArgsByTarget({
       args: { a: 1, b: 2, c: 3 },
-      argTypes: { a: { target: 'group1' }, b: { target: 'group2' }, c: { target: 'group2' } },
-    } as any);
+      argTypes: {
+        a: { name: 'a', target: 'group1' },
+        b: { name: 'b', target: 'group2' },
+        c: { name: 'c', target: 'group2' },
+      },
+    });
     expect(groups).toEqual({
       group1: {
         a: 1,
@@ -292,8 +296,8 @@ describe('groupArgsByTarget', () => {
   it('groups non-targetted args into a group with no name', () => {
     const groups = groupArgsByTarget({
       args: { a: 1, b: 2, c: 3 },
-      argTypes: { b: { name: 'b', target: 'group2' }, c: {} },
-    } as any);
+      argTypes: { a: { name: 'a' }, b: { name: 'b', target: 'group2' }, c: { name: 'c' } },
+    });
     expect(groups).toEqual({
       [UNTARGETED]: {
         a: 1,

code/lib/preview-api/src/modules/store/args.ts

@@ -148,7 +148,7 @@ export const UNTARGETED = 'UNTARGETED';
 export function groupArgsByTarget<TArgs extends Args = Args>({
   args,
   argTypes,
-}: StoryContext<Renderer, TArgs>) {
+}: Pick<StoryContext<Renderer, TArgs>, 'args' | 'argTypes'>) {
   const groupedArgs: Record<string, Partial<TArgs>> = {};
   (Object.entries(args) as [keyof TArgs, any][]).forEach(([name, value]) => {
     const { target = UNTARGETED } = (argTypes[name] || {}) as { target?: string };
@@ -159,6 +159,8 @@ export function groupArgsByTarget<TArgs extends Args = Args>({
   return groupedArgs;
 }
 
-export function noTargetArgs<TArgs extends Args = Args>(context: StoryContext<Renderer, TArgs>) {
+export function noTargetArgs<TArgs extends Args = Args>(
+  context: Pick<StoryContext<Renderer, TArgs>, 'args' | 'argTypes'>
+) {
   return groupArgsByTarget(context)[UNTARGETED];
 }