CSF: Enhance csf-factories codemods

code/lib/cli-storybook/src/codemod/helpers/config-to-csf-factory.test.ts

@@ -73,15 +73,13 @@ describe('main/preview codemod: general parsing functionality', () => {
     ).resolves.toMatchInlineSnapshot(`
       import { defineMain } from '@storybook/react-vite/node';
 
-      const config = {
-        framework: '@storybook/react-vite',
+      export default defineMain({
         tags: [],
         viteFinal: () => {
           return config;
         },
-      };
-
-      export default config;
+        framework: '@storybook/react-vite',
+      });
     `);
   });
   it('should wrap defineMain call from named exports format', async () => {
@@ -244,4 +242,21 @@ describe('preview specific functionality', () => {
       });
     `);
   });
+  it('should work', async () => {
+    await expect(
+      transform(dedent`
+        export const decorators = [1]
+        export default {
+          parameters: {},
+        }
+      `)
+    ).resolves.toMatchInlineSnapshot(`
+      import { definePreview } from '@storybook/react-vite';
+
+      export default definePreview({
+        decorators: [1],
+        parameters: {},
+      });
+    `);
+  });
 });

code/lib/cli-storybook/src/codemod/helpers/config-to-csf-factory.ts

@@ -45,8 +45,54 @@ export async function configToCsfFactory(
    * Transform into: `export default defineMain({ tags: [], parameters: {} })`
    */
   if (config._exportsObject && hasNamedExports) {
-    config._exportsObject.properties.push(...defineConfigProps);
+    // when merging named exports with default exports, add the named exports first in the list
+    config._exportsObject.properties = [...defineConfigProps, ...config._exportsObject.properties];
     programNode.body = removeExportDeclarations(programNode, exportDecls);
+
+    // After merging, ensure the default export is wrapped with defineMain/definePreview
+    const defineConfigCall = t.callExpression(t.identifier(methodName), [config._exportsObject]);
+
+    let exportDefaultNode = null as unknown as t.ExportDefaultDeclaration;
+    let declarationNodeIndex = -1;
+
+    programNode.body.forEach((node) => {
+      // Detect Syntax 1: export default <identifier>
+      if (t.isExportDefaultDeclaration(node) && t.isIdentifier(node.declaration)) {
+        const declarationName = node.declaration.name;
+
+        declarationNodeIndex = programNode.body.findIndex(
+          (n) =>
+            t.isVariableDeclaration(n) &&
+            n.declarations.some(
+              (d) =>
+                t.isIdentifier(d.id) &&
+                d.id.name === declarationName &&
+                t.isObjectExpression(d.init)
+            )
+        );
+
+        if (declarationNodeIndex !== -1) {
+          exportDefaultNode = node;
+          // remove the original declaration as it will become a default export
+          const declarationNode = programNode.body[declarationNodeIndex];
+          if (t.isVariableDeclaration(declarationNode)) {
+            const id = declarationNode.declarations[0].id;
+            const variableName = t.isIdentifier(id) && id.name;
+
+            if (variableName) {
+              programNode.body.splice(declarationNodeIndex, 1);
+            }
+          }
+        }
+      } else if (t.isExportDefaultDeclaration(node) && t.isObjectExpression(node.declaration)) {
+        // Detect Syntax 2: export default { ... }
+        exportDefaultNode = node;
+      }
+    });
+
+    if (exportDefaultNode !== null) {
+      exportDefaultNode.declaration = defineConfigCall;
+    }
   } else if (config._exportsObject) {
     /**
      * Scenario 2: Default exports

code/lib/cli-storybook/src/codemod/helpers/story-to-csf-factory.test.ts

@@ -303,16 +303,39 @@ describe('stories codemod', () => {
         };
         const data = {};
         export const A = () => {};
-        // not supported yet (story as function)
         export function B() { };
         // not supported yet (story redeclared)
         const C = { ...A, args: data, };
-        export { C };
+        const D = { args: data };
+        export { C, D as E };
         `);
 
+      expect(transformed).toMatchInlineSnapshot(`
+        import preview from '#.storybook/preview';
+
+        import { A as Component } from './Button';
+        import * as Stories from './Other.stories';
+        import someData from './fixtures';
+
+        const meta = preview.meta({
+          component: Component,
+
+          // not supported yet (story coming from another file)
+          args: Stories.A.args,
+        });
+
+        const data = {};
+        export const A = meta.story(() => {});
+        export const B = meta.story(() => {});
+        // not supported yet (story redeclared)
+        const C = { ...A.input, args: data };
+        const D = { args: data };
+        export { C, D as E };
+      `);
+
       expect(transformed).toContain('A = meta.story');
-      // @TODO: when we support these, uncomment these lines
-      // expect(transformed).toContain('B = meta.story');
+      expect(transformed).toContain('B = meta.story');
+      // @TODO: when we support these, uncomment this line
       // expect(transformed).toContain('C = meta.story');
     });
 
@@ -589,5 +612,58 @@ describe('stories codemod', () => {
         export const A = meta.story();
       `);
     });
+
+    it('should support non-conventional formats', async () => {
+      const transformed = await transform(dedent`
+        import { Meta, StoryObj as CSF3 } from '@storybook/react';
+        import { ComponentProps } from './Component';
+        import { A as Component } from './Button';
+        import * as Stories from './Other.stories';
+        import someData from './fixtures'
+        export default {
+          title: 'Component',
+          component: Component, 
+          // not supported yet (story coming from another file)
+          args: Stories.A.args
+        };
+        const data = {};
+        export const A: StoryObj = () => {};
+        export function B() { };
+        // not supported yet (story redeclared)
+        const C = { ...A, args: data, } satisfies CSF3<ComponentProps>;
+        const D = { args: data };
+        export { C, D as E };
+        `);
+
+      expect(transformed).toMatchInlineSnapshot(`
+        import preview from '#.storybook/preview';
+
+        import { A as Component } from './Button';
+        import { ComponentProps } from './Component';
+        import * as Stories from './Other.stories';
+        import someData from './fixtures';
+
+        const meta = preview.meta({
+          title: 'Component',
+          component: Component,
+
+          // not supported yet (story coming from another file)
+          args: Stories.A.args,
+        });
+
+        const data = {};
+        export const A = meta.story(() => {});
+        export const B = meta.story(() => {});
+        // not supported yet (story redeclared)
+        const C = { ...A.input, args: data } satisfies CSF3<ComponentProps>;
+        const D = { args: data };
+        export { C, D as E };
+      `);
+
+      expect(transformed).toContain('A = meta.story');
+      expect(transformed).toContain('B = meta.story');
+      // @TODO: when we support these, uncomment this line
+      // expect(transformed).toContain('C = meta.story');
+    });
   });
 });

code/lib/cli-storybook/src/codemod/helpers/story-to-csf-factory.ts

@@ -95,7 +95,7 @@ export async function storyToCsfFactory(
   // @TODO: Support unconventional formats:
   // `export function Story() { };` and `export { Story };
   // These are not part of csf._storyExports but rather csf._storyStatements and are tricky to support.
-  Object.entries(csf._storyExports).forEach(([_key, decl]) => {
+  Object.entries(csf._storyExports).forEach(([, decl]) => {
     const id = decl.id;
     const declarator = decl as t.VariableDeclarator;
     let init = t.isVariableDeclarator(declarator) ? declarator.init : undefined;
@@ -128,6 +128,34 @@ export async function storyToCsfFactory(
     }
   });
 
+  // Support function-declared stories
+  Object.entries(csf._storyExports).forEach(([exportName, decl]) => {
+    if (t.isFunctionDeclaration(decl) && decl.id) {
+      const arrowFn = t.arrowFunctionExpression(decl.params, decl.body);
+      arrowFn.async = !!decl.async;
+
+      const wrappedCall = t.callExpression(
+        t.memberExpression(t.identifier(metaVariableName), t.identifier('story')),
+        [arrowFn]
+      );
+
+      const replacement = t.exportNamedDeclaration(
+        t.variableDeclaration('const', [
+          t.variableDeclarator(t.identifier(exportName), wrappedCall),
+        ])
+      );
+
+      const pathForExport = (
+        csf as unknown as {
+          _storyPaths?: Record<string, { replaceWith?: (node: t.Node) => void }>;
+        }
+      )._storyPaths?.[exportName];
+      if (pathForExport && pathForExport.replaceWith) {
+        pathForExport.replaceWith(replacement);
+      }
+    }
+  });
+
   const storyExportDecls = new Map(
     Object.entries(csf._storyExports).filter(
       (