Update existing config file and only fall back to creating a workspace file if we cannot update the config file

code/addons/test/src/postinstall.ts

@@ -27,8 +27,8 @@ import { dedent } from 'ts-dedent';
 
 import { type PostinstallOptions } from '../../../lib/cli-storybook/src/add';
 import { SUPPORTED_FRAMEWORKS, SUPPORTED_RENDERERS } from './constants';
-import { printError, printInfo, printSuccess, step } from './postinstall-logger';
-import { updateWorkspaceFile } from './updateWorkspaceFile';
+import { printError, printInfo, printSuccess, printWarning, step } from './postinstall-logger';
+import { loadTemplate, updateConfigFile, updateWorkspaceFile } from './updateVitestFile';
 import { getAddonNames } from './utils';
 
 const ADDON_NAME = '@storybook/experimental-addon-test' as const;
@@ -40,19 +40,6 @@ const addonA11yName = '@storybook/addon-a11y';
 const findFile = async (basename: string, extensions = EXTENSIONS) =>
   findUp(extensions.map((ext) => basename + ext));
 
-const loadTemplate = async (name: string, replacements: Record<string, string>) => {
-  let template = await fs.readFile(
-    join(
-      dirname(require.resolve('@storybook/experimental-addon-test/package.json')),
-      'templates',
-      name
-    ),
-    'utf8'
-  );
-  Object.entries(replacements).forEach(([key, value]) => (template = template.replace(key, value)));
-  return template;
-};
-
 export default async function postInstall(options: PostinstallOptions) {
   printSuccess(
     '👋 Howdy!',
@@ -444,21 +431,27 @@ export default async function postInstall(options: PostinstallOptions) {
   if (vitestWorkspaceFile) {
     // If there's an existing workspace file, we update that file to include the Storybook test plugin.
     // We assume the existing workspaces include the Vite(st) config, so we won't add it.
-    const vitestSetupFilePath = relative(dirname(vitestWorkspaceFile), vitestSetupFile);
     const workspaceTemplate = await loadTemplate('vitest.workspace.template.ts', {
       EXTENDS_WORKSPACE: viteConfigFile
         ? relative(dirname(vitestWorkspaceFile), viteConfigFile)
         : '',
       CONFIG_DIR: options.configDir,
       BROWSER_CONFIG: browserConfig,
-      SETUP_FILE: vitestSetupFilePath,
+      SETUP_FILE: relative(dirname(vitestWorkspaceFile), vitestSetupFile),
     }).then((t) => t.replace(`\n  'ROOT_CONFIG',`, '').replace(/\s+extends: '',/, ''));
     const workspaceFile = await fs.readFile(vitestWorkspaceFile, 'utf8');
     const source = babelParse(workspaceTemplate);
     const target = babelParse(workspaceFile);
 
     const updated = updateWorkspaceFile(source, target);
-    if (!updated) {
+    if (updated) {
+      logger.line(1);
+      logger.plain(`${step} Updating your Vitest workspace file:`);
+      logger.plain(colors.gray(`  ${vitestWorkspaceFile}`));
+
+      const formattedContent = await formatFileContent(vitestWorkspaceFile, generate(target).code);
+      await writeFile(vitestWorkspaceFile, formattedContent);
+    } else {
       printError(
         '🚨 Oh no!',
         dedent`
@@ -475,41 +468,65 @@ export default async function postInstall(options: PostinstallOptions) {
       logger.line(1);
       return;
     }
-
-    logger.line(1);
-    logger.plain(`${step} Updating your Vitest workspace file:`);
-    logger.plain(colors.gray(`  ${vitestWorkspaceFile}`));
-
-    const formattedContent = await formatFileContent(vitestWorkspaceFile, generate(target).code);
-    await writeFile(vitestWorkspaceFile, formattedContent);
   } else if (rootConfig) {
-    // If there's an existing Vite/Vitest config, we create a workspace file so we can run Storybook tests alongside.
-    const extension = extname(rootConfig).includes('ts') ? '.ts' : '.js';
-    const newWorkspaceFile = resolve(dirname(rootConfig), `vitest.workspace${extension}`);
-    const vitestSetupFilePath = relative(dirname(newWorkspaceFile), vitestSetupFile);
-    const workspaceTemplate = await loadTemplate('vitest.workspace.template.ts', {
-      ROOT_CONFIG: relative(dirname(newWorkspaceFile), rootConfig),
+    // If there's an existing Vite/Vitest config, we update it to include the Storybook test plugin.
+    const configTemplate = await loadTemplate('vitest.config.template.ts', {
       // We only extend from Vite config (without test property), not Vitest config.
-      EXTENDS_WORKSPACE: viteConfigFile ? relative(dirname(newWorkspaceFile), viteConfigFile) : '',
       CONFIG_DIR: options.configDir,
       BROWSER_CONFIG: browserConfig,
-      SETUP_FILE: vitestSetupFilePath,
-    }).then((t) => t.replace(/\s+extends: '',/, ''));
+      SETUP_FILE: relative(dirname(rootConfig), vitestSetupFile),
+    });
+    const configFile = await fs.readFile(rootConfig, 'utf8');
+    const source = babelParse(configTemplate);
+    const target = babelParse(configFile);
 
-    logger.line(1);
-    logger.plain(`${step} Creating a Vitest workspace file:`);
-    logger.plain(colors.gray(`  ${newWorkspaceFile}`));
+    const updated = updateConfigFile(source, target);
+    if (updated) {
+      logger.line(1);
+      logger.plain(`${step} Updating your Vitest config file:`);
+      logger.plain(colors.gray(`  ${rootConfig}`));
 
-    const formattedContent = await formatFileContent(newWorkspaceFile, workspaceTemplate);
-    await writeFile(newWorkspaceFile, formattedContent);
+      const formattedContent = await formatFileContent(rootConfig, generate(target).code);
+      await writeFile(rootConfig, formattedContent);
+    } else {
+      // Fall back to creating a workspace file if we can't update the config file.
+      printWarning(
+        '⚠️ Cannot update config file',
+        dedent`
+          Could not update your existing Vitest config file:
+          ${colors.gray(rootConfig)}
+
+          Your existing config file cannot be safely updated, so instead a new Vitest
+          workspace file will be created, extending from your config file.
+
+          Please refer to the Vitest documentation to learn about the workspace file:
+          ${picocolors.cyan(`https://vitest.dev/guide/workspace.html`)}
+        `
+      );
+
+      const extension = extname(rootConfig).includes('ts') ? '.ts' : '.js';
+      const newWorkspaceFile = resolve(dirname(rootConfig), `vitest.workspace${extension}`);
+      const workspaceTemplate = await loadTemplate('vitest.workspace.template.ts', {
+        ROOT_CONFIG: relative(dirname(newWorkspaceFile), rootConfig),
+        CONFIG_DIR: options.configDir,
+        BROWSER_CONFIG: browserConfig,
+        SETUP_FILE: relative(dirname(newWorkspaceFile), vitestSetupFile),
+      });
+
+      logger.line(1);
+      logger.plain(`${step} Creating a Vitest workspace file:`);
+      logger.plain(colors.gray(`  ${newWorkspaceFile}`));
+
+      const formattedContent = await formatFileContent(newWorkspaceFile, workspaceTemplate);
+      await writeFile(newWorkspaceFile, formattedContent);
+    }
   } else {
     // If there's no existing Vitest/Vite config, we create a new Vitest config file.
     const newConfigFile = resolve(`vitest.config.${fileExtension}`);
-    const vitestSetupFilePath = relative(dirname(newConfigFile), vitestSetupFile);
     const configTemplate = await loadTemplate('vitest.config.template.ts', {
       CONFIG_DIR: options.configDir,
       BROWSER_CONFIG: browserConfig,
-      SETUP_FILE: vitestSetupFilePath,
+      SETUP_FILE: relative(dirname(newConfigFile), vitestSetupFile),
     });
 
     logger.line(1);

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

@@ -0,0 +1,370 @@
+import { describe, expect, it } from 'vitest';
+
+import * as babel from 'storybook/internal/babel';
+
+import { loadTemplate, updateConfigFile, updateWorkspaceFile } from './updateVitestFile';
+
+describe('updateConfigFile', () => {
+  it('updates vite config file', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.config.template.ts', {
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      /// <reference types="vitest/config" />
+      import { defineConfig } from 'vite'
+      import react from '@vitejs/plugin-react'
+
+      // https://vite.dev/config/
+      export default defineConfig({
+        plugins: [react()],
+        test: {
+          globals: true,
+          workspace: ['packages/*']
+        },
+      })
+    `);
+
+    const updated = updateConfigFile(source, target);
+    expect(updated).toBe(true);
+
+    const { code } = babel.generate(target);
+    expect(code).toMatchInlineSnapshot(`
+      "/// <reference types="vitest/config" />
+      import { defineConfig } from 'vite';
+      import react from '@vitejs/plugin-react';
+
+      // https://vite.dev/config/
+      import path from 'node:path';
+      import { fileURLToPath } from 'node:url';
+      import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
+      const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+      // More info at: https://storybook.js.org/docs/writing-tests/test-addon
+      export default defineConfig({
+        plugins: [react()],
+        test: {
+          globals: true,
+          workspace: ['packages/*', {
+            plugins: [
+            // The plugin will run tests for the stories defined in your Storybook config
+            // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+            storybookTest({
+              configDir: path.join(dirname, '.storybook')
+            })],
+            test: {
+              name: 'storybook',
+              browser: {
+                provider: 'playwright'
+              },
+              setupFiles: ['../.storybook/vitest.setup.ts']
+            }
+          }]
+        }
+      });"
+    `);
+  });
+
+  it('supports object notation without defineConfig', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.config.template.ts', {
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      /// <reference types="vitest/config" />
+      import react from '@vitejs/plugin-react'
+
+      // https://vite.dev/config/
+      export default {
+        plugins: [react()],
+        test: {
+          globals: true,
+          workspace: ['packages/*']
+        },
+      }
+    `);
+
+    const updated = updateConfigFile(source, target);
+    expect(updated).toBe(true);
+
+    const { code } = babel.generate(target);
+    expect(code).toMatchInlineSnapshot(`
+      "/// <reference types="vitest/config" />
+      import react from '@vitejs/plugin-react';
+
+      // https://vite.dev/config/
+      import path from 'node:path';
+      import { fileURLToPath } from 'node:url';
+      import { defineConfig } from 'vitest/config';
+      import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
+      const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+      // More info at: https://storybook.js.org/docs/writing-tests/test-addon
+      export default {
+        plugins: [react()],
+        test: {
+          globals: true,
+          workspace: ['packages/*', {
+            plugins: [
+            // The plugin will run tests for the stories defined in your Storybook config
+            // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+            storybookTest({
+              configDir: path.join(dirname, '.storybook')
+            })],
+            test: {
+              name: 'storybook',
+              browser: {
+                provider: 'playwright'
+              },
+              setupFiles: ['../.storybook/vitest.setup.ts']
+            }
+          }]
+        }
+      };"
+    `);
+  });
+
+  it('does not support function notation', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.config.template.ts', {
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      /// <reference types="vitest/config" />
+      import react from '@vitejs/plugin-react'
+
+      // https://vite.dev/config/
+      export default defineConfig(() => ({
+        plugins: [react()],
+        test: {
+          globals: true,
+          workspace: ['packages/*']
+        },
+      }))
+    `);
+
+    const updated = updateConfigFile(source, target);
+    expect(updated).toBe(false);
+  });
+
+  it('adds workspace property to test config', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.config.template.ts', {
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      /// <reference types="vitest/config" />
+      import { defineConfig } from 'vite'
+      import react from '@vitejs/plugin-react'
+
+      // https://vite.dev/config/
+      export default defineConfig({
+        plugins: [react()],
+        test: {
+          globals: true,
+        },
+      })
+    `);
+
+    const updated = updateConfigFile(source, target);
+    expect(updated).toBe(true);
+
+    const { code } = babel.generate(target);
+    expect(code).toMatchInlineSnapshot(`
+      "/// <reference types="vitest/config" />
+      import { defineConfig } from 'vite';
+      import react from '@vitejs/plugin-react';
+
+      // https://vite.dev/config/
+      import path from 'node:path';
+      import { fileURLToPath } from 'node:url';
+      import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
+      const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+      // More info at: https://storybook.js.org/docs/writing-tests/test-addon
+      export default defineConfig({
+        plugins: [react()],
+        test: {
+          globals: true,
+          workspace: [{
+            plugins: [
+            // The plugin will run tests for the stories defined in your Storybook config
+            // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+            storybookTest({
+              configDir: path.join(dirname, '.storybook')
+            })],
+            test: {
+              name: 'storybook',
+              browser: {
+                provider: 'playwright'
+              },
+              setupFiles: ['../.storybook/vitest.setup.ts']
+            }
+          }]
+        }
+      });"
+    `);
+  });
+
+  it('adds test property to vite config', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.config.template.ts', {
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      /// <reference types="vitest/config" />
+      import { defineConfig } from 'vite'
+      import react from '@vitejs/plugin-react'
+
+      // https://vite.dev/config/
+      export default defineConfig({
+        plugins: [react()],
+      })
+    `);
+
+    const updated = updateConfigFile(source, target);
+    expect(updated).toBe(true);
+
+    const { code } = babel.generate(target);
+    expect(code).toMatchInlineSnapshot(`
+      "/// <reference types="vitest/config" />
+      import { defineConfig } from 'vite';
+      import react from '@vitejs/plugin-react';
+
+      // https://vite.dev/config/
+      import path from 'node:path';
+      import { fileURLToPath } from 'node:url';
+      import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
+      const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+      // More info at: https://storybook.js.org/docs/writing-tests/test-addon
+      export default defineConfig({
+        plugins: [react()],
+        test: {
+          workspace: [{
+            plugins: [
+            // The plugin will run tests for the stories defined in your Storybook config
+            // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+            storybookTest({
+              configDir: path.join(dirname, '.storybook')
+            })],
+            test: {
+              name: 'storybook',
+              browser: {
+                provider: 'playwright'
+              },
+              setupFiles: ['../.storybook/vitest.setup.ts']
+            }
+          }]
+        }
+      });"
+    `);
+  });
+});
+
+describe('updateWorkspaceFile', () => {
+  it('updates vitest workspace file using array syntax', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.workspace.template.ts', {
+        EXTENDS_WORKSPACE: '',
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      export default ['packages/*']
+    `);
+
+    const updated = updateWorkspaceFile(source, target);
+    expect(updated).toBe(true);
+
+    const { code } = babel.generate(target);
+    expect(code).toMatchInlineSnapshot(`
+      "import path from 'node:path';
+      import { fileURLToPath } from 'node:url';
+      import { defineWorkspace } from 'vitest/config';
+      import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
+      const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+      // More info at: https://storybook.js.org/docs/writing-tests/test-addon
+      export default ['packages/*', 'ROOT_CONFIG', {
+        extends: '',
+        plugins: [
+        // The plugin will run tests for the stories defined in your Storybook config
+        // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+        storybookTest({
+          configDir: path.join(dirname, '.storybook')
+        })],
+        test: {
+          name: 'storybook',
+          browser: {
+            provider: 'playwright'
+          },
+          setupFiles: ['../.storybook/vitest.setup.ts']
+        }
+      }];"
+    `);
+  });
+
+  it('updates vitest workspace file using defineWorkspace syntax', async () => {
+    const source = babel.babelParse(
+      await loadTemplate('vitest.workspace.template.ts', {
+        EXTENDS_WORKSPACE: '',
+        CONFIG_DIR: '.storybook',
+        BROWSER_CONFIG: "{ provider: 'playwright' }",
+        SETUP_FILE: '../.storybook/vitest.setup.ts',
+      })
+    );
+    const target = babel.babelParse(`
+      import { defineWorkspace } from 'vitest/config'
+
+      export default defineWorkspace(['packages/*'])
+    `);
+
+    const updated = updateWorkspaceFile(source, target);
+    expect(updated).toBe(true);
+
+    const { code } = babel.generate(target);
+    expect(code).toMatchInlineSnapshot(`
+      "import { defineWorkspace } from 'vitest/config';
+      import path from 'node:path';
+      import { fileURLToPath } from 'node:url';
+      import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
+      const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
+
+      // More info at: https://storybook.js.org/docs/writing-tests/test-addon
+      export default defineWorkspace(['packages/*', 'ROOT_CONFIG', {
+        extends: '',
+        plugins: [
+        // The plugin will run tests for the stories defined in your Storybook config
+        // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+        storybookTest({
+          configDir: path.join(dirname, '.storybook')
+        })],
+        test: {
+          name: 'storybook',
+          browser: {
+            provider: 'playwright'
+          },
+          setupFiles: ['../.storybook/vitest.setup.ts']
+        }
+      }]);"
+    `);
+  });
+});

code/addons/test/src/updateVitestFile.ts

@@ -0,0 +1,177 @@
+import * as fs from 'node:fs/promises';
+
+import type { BabelFile } from 'storybook/internal/babel';
+
+import type { ObjectExpression } from '@babel/types';
+import { dirname, join } from 'pathe';
+
+export const loadTemplate = async (name: string, replacements: Record<string, string>) => {
+  let template = await fs.readFile(
+    join(
+      dirname(require.resolve('@storybook/experimental-addon-test/package.json')),
+      '../templates',
+      name
+    ),
+    'utf8'
+  );
+  Object.entries(replacements).forEach(([key, value]) => (template = template.replace(key, value)));
+  return template;
+};
+
+// Recursively merge object properties from source into target
+// Handles nested objects and shallowly merging of arrays
+const mergeProperties = (
+  source: ObjectExpression['properties'],
+  target: ObjectExpression['properties']
+) => {
+  for (const sourceProp of source) {
+    if (sourceProp.type === 'ObjectProperty') {
+      const targetProp = target.find(
+        (p) =>
+          sourceProp.key.type === 'Identifier' &&
+          p.type === 'ObjectProperty' &&
+          p.key.type === 'Identifier' &&
+          p.key.name === sourceProp.key.name
+      );
+      if (targetProp && targetProp.type === 'ObjectProperty') {
+        if (
+          sourceProp.value.type === 'ObjectExpression' &&
+          targetProp.value.type === 'ObjectExpression'
+        ) {
+          mergeProperties(sourceProp.value.properties, targetProp.value.properties);
+        } else if (
+          sourceProp.value.type === 'ArrayExpression' &&
+          targetProp.value.type === 'ArrayExpression'
+        ) {
+          targetProp.value.elements.push(...sourceProp.value.elements);
+        } else {
+          targetProp.value = sourceProp.value;
+        }
+      } else {
+        target.push(sourceProp);
+      }
+    }
+  }
+};
+
+export const updateConfigFile = (source: BabelFile['ast'], target: BabelFile['ast']) => {
+  let updated = false;
+  for (const sourceNode of source.program.body) {
+    if (sourceNode.type === 'ImportDeclaration') {
+      // Insert imports that don't already exist (according to their local specifier name)
+      if (
+        !target.program.body.some(
+          (targetNode) =>
+            targetNode.type === sourceNode.type &&
+            targetNode.specifiers.some((s) => s.local.name === sourceNode.specifiers[0].local.name)
+        )
+      ) {
+        const lastImport = target.program.body.findLastIndex((n) => n.type === 'ImportDeclaration');
+        target.program.body.splice(lastImport + 1, 0, sourceNode);
+      }
+    } else if (sourceNode.type === 'VariableDeclaration') {
+      // Copy over variable declarations, making sure they're inserted after any imports
+      if (
+        !target.program.body.some(
+          (targetNode) =>
+            targetNode.type === sourceNode.type &&
+            targetNode.declarations.some(
+              (d) =>
+                'name' in d.id &&
+                'name' in sourceNode.declarations[0].id &&
+                d.id.name === sourceNode.declarations[0].id.name
+            )
+        )
+      ) {
+        const lastImport = target.program.body.findLastIndex((n) => n.type === 'ImportDeclaration');
+        target.program.body.splice(lastImport + 1, 0, sourceNode);
+      }
+    } else if (sourceNode.type === 'ExportDefaultDeclaration') {
+      const exportDefault = target.program.body.find((n) => n.type === 'ExportDefaultDeclaration');
+      if (
+        exportDefault &&
+        sourceNode.declaration.type === 'CallExpression' &&
+        sourceNode.declaration.arguments.length > 0 &&
+        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);
+          updated = true;
+        }
+      }
+    }
+  }
+  return updated;
+};
+
+export const updateWorkspaceFile = (source: BabelFile['ast'], target: BabelFile['ast']) => {
+  let updated = false;
+  for (const sourceNode of source.program.body) {
+    if (sourceNode.type === 'ImportDeclaration') {
+      // Insert imports that don't already exist
+      if (
+        !target.program.body.some(
+          (targetNode) =>
+            targetNode.type === sourceNode.type &&
+            targetNode.source.value === sourceNode.source.value &&
+            targetNode.specifiers.some((s) => s.local.name === sourceNode.specifiers[0].local.name)
+        )
+      ) {
+        const lastImport = target.program.body.findLastIndex((n) => n.type === 'ImportDeclaration');
+        target.program.body.splice(lastImport + 1, 0, sourceNode);
+      }
+    } else if (sourceNode.type === 'VariableDeclaration') {
+      // Copy over variable declarations, making sure they're inserted after any imports
+      if (
+        !target.program.body.some(
+          (targetNode) =>
+            targetNode.type === sourceNode.type &&
+            targetNode.declarations.some(
+              (d) =>
+                'name' in d.id &&
+                'name' in sourceNode.declarations[0].id &&
+                d.id.name === sourceNode.declarations[0].id.name
+            )
+        )
+      ) {
+        const lastImport = target.program.body.findLastIndex((n) => n.type === 'ImportDeclaration');
+        target.program.body.splice(lastImport + 1, 0, sourceNode);
+      }
+    } else if (sourceNode.type === 'ExportDefaultDeclaration') {
+      // Merge workspace array, which is the default export on both sides but may or may not be
+      // wrapped in a defineWorkspace call
+      const exportDefault = target.program.body.find((n) => n.type === 'ExportDefaultDeclaration');
+      if (
+        exportDefault &&
+        sourceNode.declaration.type === 'CallExpression' &&
+        sourceNode.declaration.arguments.length > 0 &&
+        sourceNode.declaration.arguments[0].type === 'ArrayExpression' &&
+        sourceNode.declaration.arguments[0].elements.length > 0
+      ) {
+        const { elements } = sourceNode.declaration.arguments[0];
+        if (exportDefault.declaration.type === 'ArrayExpression') {
+          exportDefault.declaration.elements.push(...elements);
+          updated = true;
+        } else if (
+          exportDefault.declaration.type === 'CallExpression' &&
+          exportDefault.declaration.callee.type === 'Identifier' &&
+          exportDefault.declaration.callee.name === 'defineWorkspace' &&
+          exportDefault.declaration.arguments[0]?.type === 'ArrayExpression'
+        ) {
+          exportDefault.declaration.arguments[0].elements.push(...elements);
+          updated = true;
+        }
+      }
+    }
+  }
+  return updated;
+};

code/addons/test/src/updateWorkspaceFile.test.ts

@@ -1,93 +0,0 @@
-import { expect, it } from 'vitest';
-
-import * as babel from 'storybook/internal/babel';
-
-import { updateWorkspaceFile } from './updateWorkspaceFile';
-
-const source = babel.babelParse(`
-  import path from 'node:path';
-  import { fileURLToPath } from 'node:url';
-  import { defineWorkspace } from 'vitest/config';
-  import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
-
-  const dirname = typeof __dirname !== 'undefined'
-    ? __dirname
-    : path.dirname(fileURLToPath(import.meta.url));
-
-  // More info at: https://storybook.js.org/docs/writing-tests/test-addon
-  export default defineWorkspace([{
-    extends: '${'./vitest.config.ts'}',
-    plugins: [
-      // The plugin will run tests for the stories defined in your Storybook config
-      // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
-      storybookTest({ configDir: path.join(dirname, '${'.storybook'}') })
-    ],
-    test: {
-      name: 'storybook',
-    },
-  }]);
-`);
-
-it('updates vitest workspace file using array syntax', async () => {
-  const target = babel.babelParse(`
-    export default ['packages/*']
-  `);
-
-  updateWorkspaceFile(source, target);
-  const { code } = babel.generate(target);
-
-  expect(code).toMatchInlineSnapshot(`
-    "import path from 'node:path';
-    import { fileURLToPath } from 'node:url';
-    import { defineWorkspace } from 'vitest/config';
-    import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
-    const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
-
-    // More info at: https://storybook.js.org/docs/writing-tests/test-addon
-    export default ['packages/*', {
-      extends: './vitest.config.ts',
-      plugins: [
-      // The plugin will run tests for the stories defined in your Storybook config
-      // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
-      storybookTest({
-        configDir: path.join(dirname, '.storybook')
-      })],
-      test: {
-        name: 'storybook'
-      }
-    }];"
-  `);
-});
-
-it('updates vitest workspace file using defineWorkspace syntax', async () => {
-  const target = babel.babelParse(`
-    import { defineWorkspace } from 'vitest/config'
-
-    export default defineWorkspace(['packages/*'])
-  `);
-
-  updateWorkspaceFile(source, target);
-  const { code } = babel.generate(target);
-
-  expect(code).toMatchInlineSnapshot(`
-    "import { defineWorkspace } from 'vitest/config';
-    import path from 'node:path';
-    import { fileURLToPath } from 'node:url';
-    import { storybookTest } from '@storybook/experimental-addon-test/vitest-plugin';
-    const dirname = typeof __dirname !== 'undefined' ? __dirname : path.dirname(fileURLToPath(import.meta.url));
-
-    // More info at: https://storybook.js.org/docs/writing-tests/test-addon
-    export default defineWorkspace(['packages/*', {
-      extends: './vitest.config.ts',
-      plugins: [
-      // The plugin will run tests for the stories defined in your Storybook config
-      // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
-      storybookTest({
-        configDir: path.join(dirname, '.storybook')
-      })],
-      test: {
-        name: 'storybook'
-      }
-    }]);"
-  `);
-});

code/addons/test/src/updateWorkspaceFile.ts

@@ -1,64 +0,0 @@
-import type { BabelFile } from 'storybook/internal/babel';
-
-export const updateWorkspaceFile = (source: BabelFile['ast'], target: BabelFile['ast']) => {
-  let updated = false;
-  for (const sourceNode of source.program.body) {
-    if (sourceNode.type === 'ImportDeclaration') {
-      // Insert imports that don't already exist
-      if (
-        !target.program.body.some(
-          (targetNode) =>
-            targetNode.type === sourceNode.type &&
-            targetNode.source.value === sourceNode.source.value &&
-            targetNode.specifiers.some((s) => s.local.name === sourceNode.specifiers[0].local.name)
-        )
-      ) {
-        const lastImport = target.program.body.findLastIndex((n) => n.type === 'ImportDeclaration');
-        target.program.body.splice(lastImport + 1, 0, sourceNode);
-      }
-    } else if (sourceNode.type === 'VariableDeclaration') {
-      // Copy over variable declarations, making sure they're inserted after any imports
-      if (
-        !target.program.body.some(
-          (targetNode) =>
-            targetNode.type === sourceNode.type &&
-            targetNode.declarations.some(
-              (d) =>
-                'name' in d.id &&
-                'name' in sourceNode.declarations[0].id &&
-                d.id.name === sourceNode.declarations[0].id.name
-            )
-        )
-      ) {
-        const lastImport = target.program.body.findLastIndex((n) => n.type === 'ImportDeclaration');
-        target.program.body.splice(lastImport + 1, 0, sourceNode);
-      }
-    } else if (sourceNode.type === 'ExportDefaultDeclaration') {
-      // Merge workspace array, which is the default export on both sides but may or may not be
-      // wrapped in a defineWorkspace call
-      const exportDefault = target.program.body.find((n) => n.type === 'ExportDefaultDeclaration');
-      if (
-        exportDefault &&
-        sourceNode.declaration.type === 'CallExpression' &&
-        sourceNode.declaration.arguments.length > 0 &&
-        sourceNode.declaration.arguments[0].type === 'ArrayExpression' &&
-        sourceNode.declaration.arguments[0].elements.length > 0
-      ) {
-        const { elements } = sourceNode.declaration.arguments[0];
-        if (exportDefault.declaration.type === 'ArrayExpression') {
-          exportDefault.declaration.elements.push(...elements);
-          updated = true;
-        } else if (
-          exportDefault.declaration.type === 'CallExpression' &&
-          exportDefault.declaration.callee.type === 'Identifier' &&
-          exportDefault.declaration.callee.name === 'defineWorkspace' &&
-          exportDefault.declaration.arguments[0]?.type === 'ArrayExpression'
-        ) {
-          exportDefault.declaration.arguments[0].elements.push(...elements);
-          updated = true;
-        }
-      }
-    }
-  }
-  return updated;
-};

code/addons/test/templates/typings.d.ts

@@ -0,0 +1 @@
+declare const BROWSER_CONFIG: object;

code/addons/test/templates/vitest.config.template.ts

@@ -10,15 +10,20 @@ const dirname =
 
 // More info at: https://storybook.js.org/docs/writing-tests/test-addon
 export default defineConfig({
-  plugins: [
-    // The plugin will run tests for the stories defined in your Storybook config
-    // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
-    storybookTest({ configDir: path.join(dirname, 'CONFIG_DIR') }),
-  ],
   test: {
-    name: 'storybook',
-    // @ts-expect-error (not defined, will be replaced)
-    browser: BROWSER_CONFIG,
-    setupFiles: ['SETUP_FILE'],
+    workspace: [
+      {
+        plugins: [
+          // The plugin will run tests for the stories defined in your Storybook config
+          // See options at: https://storybook.js.org/docs/writing-tests/test-addon#storybooktest
+          storybookTest({ configDir: path.join(dirname, 'CONFIG_DIR') }),
+        ],
+        test: {
+          name: 'storybook',
+          browser: BROWSER_CONFIG,
+          setupFiles: ['SETUP_FILE'],
+        },
+      },
+    ],
   },
 });

code/addons/test/templates/vitest.workspace.template.ts

@@ -20,7 +20,6 @@ export default defineWorkspace([
     ],
     test: {
       name: 'storybook',
-      // @ts-expect-error (not defined, will be replaced)
       browser: BROWSER_CONFIG,
       setupFiles: ['SETUP_FILE'],
     },