Addon-Vitest: Support vite/vitest config with deferred export

code/addons/vitest/src/updateVitestFile.test.ts

@@ -291,6 +291,91 @@ describe('updateConfigFile', () => {
     `);
   });
 
+  it('updates config which is not exported immediately', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.config.3.2.template.ts', {
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      import { defineConfig } from 'vite'
+      import viteReact from '@vitejs/plugin-react'
+      import { fileURLToPath, URL } from 'url'
+
+      const config = defineConfig({
+        resolve: {
+          preserveSymlinks: true,
+          alias: {
+            '@': fileURLToPath(new URL('./src', import.meta.url)),
+          },
+        },
+        plugins: [
+          viteReact(),
+        ],
+      })
+
+      export default config
+    `);
+
+    const before = babel.generate(target).code;
+    const updated = updateConfigFile(source, target);
+    expect(updated).toBe(true);
+
+    const after = babel.generate(target).code;
+
+    expect(getDiff(before, after)).toMatchInlineSnapshot(`
+  "  import { defineConfig } from 'vite';
+    import viteReact from '@vitejs/plugin-react';
+    import { fileURLToPath, URL } from 'url';
+
+  + import path from 'node:path';
+  + import { storybookTest } from '@storybook/addon-vitest/vitest-plugin';
+  + const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+  + 
+  + // More info at: https://storybook.js.org/docs/next/writing-tests/integrations/vitest-addon
+  + 
+    const config = defineConfig({
+      resolve: {
+        preserveSymlinks: true,
+        alias: {
+          '@': fileURLToPath(new URL('./src', import.meta.url))
+        }
+      },
+
+  -   plugins: [viteReact()]
+  - 
+  +   plugins: [viteReact()],
+  +   test: {
+  +     projects: [{
+  +       extends: true,
+  +       plugins: [
+  +       // The plugin will run tests for the stories defined in your Storybook config
+  +       // See options at: https://storybook.js.org/docs/next/writing-tests/integrations/vitest-addon#storybooktest
+  +       storybookTest({
+  +         configDir: path.join(dirname, '.storybook')
+  +       })],
+  +       test: {
+  +         name: 'storybook',
+  +         browser: {
+  +           enabled: true,
+  +           headless: true,
+  +           provider: 'playwright',
+  +           instances: [{
+  +             browser: 'chromium'
+  +           }]
+  +         },
+  +         setupFiles: ['../.storybook/vitest.setup.ts']
+  +       }
+  +     }]
+  +   }
+  + 
+    });
+    export default config;"
+`);
+  });
+
   it('edits projects property of test config', async () => {
     const source = babel.babelParse(
       await loadTemplate('vitest.config.3.2.template.ts', {

code/addons/vitest/src/updateVitestFile.ts

@@ -54,6 +54,60 @@ const mergeProperties = (
   }
 };
 
+/**
+ * Resolves the target's default export to the actual config object expression we can merge into.
+ * Handles: export default defineConfig({}), export default {}, and export default config (where
+ * config is a variable holding defineConfig({}) or {}).
+ */
+const getTargetConfigObject = (
+  target: BabelFile['ast'],
+  exportDefault: t.ExportDefaultDeclaration
+): t.ObjectExpression | null => {
+  const decl = exportDefault.declaration;
+  if (decl.type === 'ObjectExpression') {
+    return decl;
+  }
+  if (
+    decl.type === 'CallExpression' &&
+    decl.callee.type === 'Identifier' &&
+    decl.callee.name === 'defineConfig' &&
+    decl.arguments[0]?.type === 'ObjectExpression'
+  ) {
+    return decl.arguments[0] as t.ObjectExpression;
+  }
+  if (decl.type === 'Identifier') {
+    const varName = decl.name;
+    const varDecl = target.program.body.find(
+      (n): n is t.VariableDeclaration =>
+        n.type === 'VariableDeclaration' &&
+        n.declarations.some((d) => d.id.type === 'Identifier' && d.id.name === varName)
+    );
+    if (!varDecl) {
+      return null;
+    }
+    const declarator = varDecl.declarations.find(
+      (d) => d.id.type === 'Identifier' && d.id.name === varName
+    );
+    if (!declarator?.init) {
+      return null;
+    }
+    const init = declarator.init;
+    if (
+      init.type === 'CallExpression' &&
+      init.callee.type === 'Identifier' &&
+      init.callee.name === 'defineConfig' &&
+      init.arguments[0]?.type === 'ObjectExpression'
+    ) {
+      return init.arguments[0] as t.ObjectExpression;
+    }
+    if (init.type === 'ObjectExpression') {
+      return init;
+    }
+    return null;
+  }
+  return null;
+};
+
 /**
  * Merges a source Vitest configuration AST into a target configuration AST.
  *
@@ -98,27 +152,42 @@ export const updateConfigFile = (source: BabelFile['ast'], target: BabelFile['as
   }
 
   // Check if this is a function notation that we don't support
+  const rejectFunctionNotation = (decl: t.ExportDefaultDeclaration['declaration']) => {
+    if (
+      decl.type === 'CallExpression' &&
+      decl.callee.type === 'Identifier' &&
+      decl.callee.name === 'defineConfig' &&
+      decl.arguments.length > 0 &&
+      decl.arguments[0].type === 'ArrowFunctionExpression'
+    ) {
+      return true;
+    }
+    return false;
+  };
   if (
     targetExportDefault.declaration.type === 'CallExpression' &&
-    targetExportDefault.declaration.callee.type === 'Identifier' &&
-    targetExportDefault.declaration.callee.name === 'defineConfig' &&
-    targetExportDefault.declaration.arguments.length > 0 &&
-    targetExportDefault.declaration.arguments[0].type === 'ArrowFunctionExpression'
+    rejectFunctionNotation(targetExportDefault.declaration)
   ) {
-    // This is function notation that we don't support
     return false;
   }
+  if (targetExportDefault.declaration.type === 'Identifier') {
+    const varName = targetExportDefault.declaration.name;
+    const varDecl = target.program.body.find(
+      (n): n is t.VariableDeclaration =>
+        n.type === 'VariableDeclaration' &&
+        n.declarations.some((d) => d.id.type === 'Identifier' && d.id.name === varName)
+    );
+    const declarator = varDecl?.declarations.find(
+      (d) => d.id.type === 'Identifier' && d.id.name === varName
+    );
+    if (declarator?.init?.type === 'CallExpression' && rejectFunctionNotation(declarator.init)) {
+      return false;
+    }
+  }
 
-  // Check if we can handle mergeConfig patterns
+  // Check if we can handle mergeConfig patterns (including export default config where config = defineConfig({}))
   let canHandleConfig = false;
-  if (targetExportDefault.declaration.type === 'ObjectExpression') {
-    canHandleConfig = true;
-  } else if (
-    targetExportDefault.declaration.type === 'CallExpression' &&
-    targetExportDefault.declaration.callee.type === 'Identifier' &&
-    targetExportDefault.declaration.callee.name === 'defineConfig' &&
-    targetExportDefault.declaration.arguments[0]?.type === 'ObjectExpression'
-  ) {
+  if (getTargetConfigObject(target, targetExportDefault) !== null) {
     canHandleConfig = true;
   } else if (
     targetExportDefault.declaration.type === 'CallExpression' &&
@@ -173,16 +242,9 @@ export const updateConfigFile = (source: BabelFile['ast'], target: BabelFile['as
         sourceNode.declaration.arguments[0].type === 'ObjectExpression'
       ) {
         const { properties } = sourceNode.declaration.arguments[0];
-        if (exportDefault.declaration.type === 'ObjectExpression') {
-          mergeProperties(properties, exportDefault.declaration.properties);
-          updated = true;
-        } else if (
-          exportDefault.declaration.type === 'CallExpression' &&
-          exportDefault.declaration.callee.type === 'Identifier' &&
-          exportDefault.declaration.callee.name === 'defineConfig' &&
-          exportDefault.declaration.arguments[0]?.type === 'ObjectExpression'
-        ) {
-          mergeProperties(properties, exportDefault.declaration.arguments[0].properties);
+        const targetConfigObject = getTargetConfigObject(target, exportDefault);
+        if (targetConfigObject !== null) {
+          mergeProperties(properties, targetConfigObject.properties);
           updated = true;
         } else if (
           exportDefault.declaration.type === 'CallExpression' &&