Vitest: Implement add command for vitest addon

code/addons/vitest/package.json

@@ -75,6 +75,8 @@
   },
   "devDependencies": {
     "@vitest/browser": "^2.0.0",
+    "tinyrainbow": "^1.2.0",
+    "ts-dedent": "^2.2.0",
     "vitest": "^2.0.0"
   },
   "peerDependencies": {

code/addons/vitest/src/postinstall.ts

@@ -1,3 +1,169 @@
-export default async function postinstall(context: any) {
-  console.log('[addon-vitest] postinstall with', context);
+import { existsSync } from 'node:fs';
+import * as fs from 'node:fs/promises';
+import { writeFile } from 'node:fs/promises';
+import { join, resolve } from 'node:path';
+
+import {
+  JsPackageManagerFactory,
+  extractProperFrameworkName,
+  loadAllPresets,
+  loadMainConfig,
+} from 'storybook/internal/common';
+import { logger } from 'storybook/internal/node-logger';
+
+import c from 'tinyrainbow';
+import dedent from 'ts-dedent';
+
+import { type PostinstallOptions } from '../../../lib/cli-storybook/src/add';
+
+export default async function postInstall(options: PostinstallOptions) {
+  const packageManager = JsPackageManagerFactory.getPackageManager({
+    force: options.packageManager,
+  });
+
+  const info = await getFrameworkInfo(options);
+
+  if (info.builderPackageName !== '@storybook/builder-vite') {
+    logger.info('Only @storybook/builder-vite is supported');
+    return;
+  }
+
+  const configFile = resolve('vitest.config.ts');
+  if (existsSync(configFile)) {
+    logger.info('You already have a vitest.config.ts file. Check our docs:');
+    return;
+  }
+
+  const annotationsImport = [
+    '@storybook/nextjs',
+    '@storybook/experimental-nextjs-vite',
+    '@storybook/sveltekit',
+  ].includes(info.frameworkPackageName)
+    ? info.frameworkPackageName
+    : ['@storybook/react', '@storybook/svelte', '@storybook/vue3'].includes(
+          info.rendererPackageName
+        )
+      ? info.rendererPackageName
+      : null;
+
+  if (!annotationsImport) {
+    logger.info('Your framework is not yet supported for the vitest addon.');
+    return;
+  }
+
+  const vitestInfo = getVitestPluginInfo(info.frameworkPackageName);
+
+  const packages = ['vitest@latest', '@vitest/browser@latest', 'playwright@latest'];
+  logger.info(c.bold('Installing packages...'));
+  logger.info(packages.join(', '));
+  await packageManager.addDependencies({ installAsDevDependencies: true }, packages);
+
+  await packageManager.executeCommand({
+    command: 'npx',
+    args: ['playwright', 'install', 'chromium', '--with-deps'],
+  });
+
+  await writeFile(
+    resolve(options.configDir, 'vitest.setup.ts'),
+    dedent`
+      import { beforeAll } from 'vitest'
+      import { setProjectAnnotations } from '${annotationsImport}'
+      import * as projectAnnotations from './preview'
+
+      const project = setProjectAnnotations(projectAnnotations)
+
+      beforeAll(project.beforeAll)
+    `
+  );
+  await writeFile(
+    configFile,
+    dedent`
+      import { defineConfig } from "vitest/config";
+      import { storybookTest } from "@storybook/experimental-addon-vitest/plugin";
+      ${vitestInfo.frameworkPluginImport ? vitestInfo.frameworkPluginImport + '\n' : ''}
+      export default defineConfig({
+        plugins: [
+          storybookTest(),${vitestInfo.frameworkPluginCall ? '\n' + vitestInfo.frameworkPluginCall : ''}
+        ],
+        test: {
+          include: ['**/*.{stories}.?(m)[jt]s?(x)'],
+          browser: {
+            enabled: true,
+            name: 'chromium',
+            provider: 'playwright',
+            headless: true,
+            screenshotFailures: false,
+          },
+          isolate: false,
+          setupFiles: ['./.storybook/vitest.setup.ts'],
+        },
+      })
+    `
+  );
+}
+
+const getVitestPluginInfo = (framework: string) => {
+  let frameworkPluginImport = '';
+  let frameworkPluginCall = '';
+
+  if (framework === '@storybook/nextjs') {
+    frameworkPluginImport = "import vitePluginNext from 'vite-plugin-storybook-nextjs'";
+    frameworkPluginCall = 'vitePluginNext()';
+  }
+
+  if (framework === '@storybook/sveltekit') {
+    frameworkPluginImport = "import { storybookSveltekitPlugin } from '@storybook/sveltekit/vite'";
+    frameworkPluginCall = 'storybookSveltekitPlugin()';
+  }
+
+  return { frameworkPluginImport, frameworkPluginCall };
+};
+
+async function getFrameworkInfo({ configDir, packageManager: pkgMgr }: PostinstallOptions) {
+  const packageManager = JsPackageManagerFactory.getPackageManager({ force: pkgMgr });
+  const packageJson = await packageManager.retrievePackageJson();
+
+  const config = await loadMainConfig({ configDir, noCache: true });
+  const { framework } = config;
+
+  const frameworkName = typeof framework === 'string' ? framework : framework?.name;
+
+  if (!frameworkName) {
+    throw new Error('Could not detect your storybook framework.');
+  }
+
+  const frameworkPackageName = extractProperFrameworkName(frameworkName);
+
+  const presets = await loadAllPresets({
+    corePresets: [join(frameworkName, 'preset')],
+    overridePresets: [
+      require.resolve('@storybook/core/core-server/presets/common-override-preset'),
+    ],
+    configDir,
+    packageJson,
+    isCritical: true,
+  });
+
+  const core = await presets.apply('core', {});
+
+  const { builder, renderer } = core;
+
+  if (!builder || !renderer) {
+    throw new Error('Could not detect your storybook framework.');
+  }
+
+  const builderPackageJson = await fs.readFile(
+    `${typeof builder === 'string' ? builder : builder.name}/package.json`,
+    'utf8'
+  );
+  const builderPackageName = JSON.parse(builderPackageJson).name;
+
+  const rendererPackageJson = await fs.readFile(`${renderer}/package.json`, 'utf8');
+  const rendererPackageName = JSON.parse(rendererPackageJson).name;
+
+  return {
+    frameworkPackageName,
+    builderPackageName,
+    rendererPackageName,
+  };
 }

code/lib/cli-storybook/src/add.ts

@@ -17,6 +17,7 @@ import { postinstallAddon } from './postinstallAddon';
 
 export interface PostinstallOptions {
   packageManager: PackageManagerName;
+  configDir: string;
 }
 
 /**
@@ -139,7 +140,7 @@ export async function add(
   await writeConfig(main);
 
   if (!skipPostinstall && isCoreAddon(addonName)) {
-    await postinstallAddon(addonName, { packageManager: packageManager.type });
+    await postinstallAddon(addonName, { packageManager: packageManager.type, configDir });
   }
 }
 function isValidVersion(version: string) {

code/yarn.lock

@@ -6133,6 +6133,8 @@ __metadata:
   dependencies:
     "@storybook/csf": "npm:^0.1.11"
     "@vitest/browser": "npm:^2.0.0"
+    tinyrainbow: "npm:^1.2.0"
+    ts-dedent: "npm:^2.2.0"
     vitest: "npm:^2.0.0"
   peerDependencies:
     "@vitest/browser": ^2.0.0