CLI: Improve add command & add tests

code/lib/cli/src/add.test.ts

@@ -0,0 +1,148 @@
+import { describe, expect, test, vi } from 'vitest';
+import { add, getVersionSpecifier } from './add';
+
+const MockedConfig = vi.hoisted(() => {
+  return {
+    appendValueToArray: vi.fn(),
+  };
+});
+const MockedPackageManager = vi.hoisted(() => {
+  return {
+    retrievePackageJson: vi.fn(() => ({})),
+    latestVersion: vi.fn(() => '1.0.0'),
+    addDependencies: vi.fn(() => {}),
+    type: 'npm',
+  };
+});
+const MockedPostInstall = vi.hoisted(() => {
+  return {
+    postinstallAddon: vi.fn(),
+  };
+});
+const MockedConsole = {
+  log: vi.fn(),
+  warn: vi.fn(),
+  error: vi.fn(),
+} as any as Console;
+
+vi.mock('@storybook/csf-tools', () => {
+  return {
+    readConfig: vi.fn(() => MockedConfig),
+    writeConfig: vi.fn(),
+  };
+});
+vi.mock('./postinstallAddon', () => {
+  return MockedPostInstall;
+});
+vi.mock('@storybook/core-common', () => {
+  return {
+    getStorybookInfo: vi.fn(() => ({ mainConfig: {}, configDir: '' })),
+    serverRequire: vi.fn(() => ({})),
+    JsPackageManagerFactory: {
+      getPackageManager: vi.fn(() => MockedPackageManager),
+    },
+    getCoercedStorybookVersion: vi.fn(() => '8.0.0'),
+    versions: {
+      '@storybook/addon-docs': '^8.0.0',
+    },
+  };
+});
+
+describe('getVersionSpecifier', (it) => {
+  test.each([
+    ['@storybook/addon-docs', ['@storybook/addon-docs', undefined]],
+    ['@storybook/addon-docs@7.0.1', ['@storybook/addon-docs', '7.0.1']],
+    ['@storybook/addon-docs@7.0.1-beta.1', ['@storybook/addon-docs', '7.0.1-beta.1']],
+    ['@storybook/addon-docs@~7.0.1-beta.1', ['@storybook/addon-docs', '~7.0.1-beta.1']],
+    ['@storybook/addon-docs@^7.0.1-beta.1', ['@storybook/addon-docs', '^7.0.1-beta.1']],
+    ['@storybook/addon-docs@next', ['@storybook/addon-docs', 'next']],
+  ])('%s => %s', (input, expected) => {
+    const result = getVersionSpecifier(input);
+    expect(result[0]).toEqual(expected[0]);
+    expect(result[1]).toEqual(expected[1]);
+  });
+});
+
+describe('add', () => {
+  const testData = [
+    { input: 'aa', expected: 'aa@^1.0.0' }, // resolves to the latest version
+    { input: 'aa@4', expected: 'aa@^4' },
+    { input: 'aa@4.1.0', expected: 'aa@^4.1.0' },
+    { input: 'aa@^4', expected: 'aa@^4' },
+    { input: 'aa@~4', expected: 'aa@~4' },
+    { input: 'aa@4.1.0-alpha.1', expected: 'aa@^4.1.0-alpha.1' },
+    { input: 'aa@next', expected: 'aa@next' },
+
+    { input: '@org/aa', expected: '@org/aa@^1.0.0' },
+    { input: '@org/aa@4', expected: '@org/aa@^4' },
+    { input: '@org/aa@4.1.0', expected: '@org/aa@^4.1.0' },
+    { input: '@org/aa@4.1.0-alpha.1', expected: '@org/aa@^4.1.0-alpha.1' },
+    { input: '@org/aa@next', expected: '@org/aa@next' },
+
+    { input: '@storybook/addon-docs@~4', expected: '@storybook/addon-docs@~4' },
+    { input: '@storybook/addon-docs@next', expected: '@storybook/addon-docs@next' },
+    { input: '@storybook/addon-docs', expected: '@storybook/addon-docs@^8.0.0' }, // takes it from the versions file
+  ];
+
+  test.each(testData)('$input', async ({ input, expected }) => {
+    const [packageName] = getVersionSpecifier(input);
+
+    await add(input, { packageManager: 'npm', skipPostinstall: true }, MockedConsole);
+
+    expect(MockedConfig.appendValueToArray).toHaveBeenCalledWith(
+      expect.arrayContaining(['addons']),
+      packageName
+    );
+
+    expect(MockedPackageManager.addDependencies).toHaveBeenCalledWith(
+      { installAsDevDependencies: true },
+      [expected]
+    );
+  });
+});
+
+describe('add (extra)', () => {
+  test('not warning when installing the correct version of storybook', async () => {
+    await add(
+      '@storybook/addon-docs',
+      { packageManager: 'npm', skipPostinstall: true },
+      MockedConsole
+    );
+
+    expect(MockedConsole.warn).not.toHaveBeenCalledWith(
+      expect.stringContaining(`is not the same as the version of Storybook you are using.`)
+    );
+  });
+  test('not warning when installing unrelated package', async () => {
+    await add('aa', { packageManager: 'npm', skipPostinstall: true }, MockedConsole);
+
+    expect(MockedConsole.warn).not.toHaveBeenCalledWith(
+      expect.stringContaining(`is not the same as the version of Storybook you are using.`)
+    );
+  });
+  test('warning when installing a core addon mismatching version of storybook', async () => {
+    await add(
+      '@storybook/addon-docs@2.0.0',
+      { packageManager: 'npm', skipPostinstall: true },
+      MockedConsole
+    );
+
+    expect(MockedConsole.warn).toHaveBeenCalledWith(
+      expect.stringContaining(
+        `The version of @storybook/addon-docs you are installing is not the same as the version of Storybook you are using. This may lead to unexpected behavior.`
+      )
+    );
+  });
+
+  test('postInstall', async () => {
+    await add(
+      '@storybook/addon-docs',
+      { packageManager: 'npm', skipPostinstall: false },
+      MockedConsole
+    );
+
+    expect(MockedPostInstall.postinstallAddon).toHaveBeenCalledWith('@storybook/addon-docs', {
+      packageManager: 'npm',
+    });
+  });
+});

code/lib/cli/src/add.ts

@@ -1,44 +1,32 @@
 import {
   getStorybookInfo,
   serverRequire,
-  getCoercedStorybookVersion,
-  isCorePackage,
   JsPackageManagerFactory,
+  getCoercedStorybookVersion,
   type PackageManagerName,
+  versions,
 } from '@storybook/core-common';
 import { readConfig, writeConfig } from '@storybook/csf-tools';
 import { isAbsolute, join } from 'path';
 import SemVer from 'semver';
 import dedent from 'ts-dedent';
+import { postinstallAddon } from './postinstallAddon';
 
-const logger = console;
-
-interface PostinstallOptions {
+export interface PostinstallOptions {
   packageManager: PackageManagerName;
 }
 
-const postinstallAddon = async (addonName: string, options: PostinstallOptions) => {
-  try {
-    const modulePath = require.resolve(`${addonName}/postinstall`, { paths: [process.cwd()] });
-
-    const postinstall = require(modulePath);
-
-    try {
-      logger.log(`Running postinstall script for ${addonName}`);
-      await postinstall(options);
-    } catch (e) {
-      logger.error(`Error running postinstall script for ${addonName}`);
-      logger.error(e);
-    }
-  } catch (e) {
-    // no postinstall script
-  }
-};
-
-const getVersionSpecifier = (addon: string) => {
-  const groups = /^(...*)@(.*)$/.exec(addon);
+/**
+ * Extract the addon name and version specifier from the input string
+ * @param addon - the input string
+ * @returns [addonName, versionSpecifier]
+ * @example
+ * getVersionSpecifier('@storybook/addon-docs@7.0.1') => ['@storybook/addon-docs', '7.0.1']
+ */
+export const getVersionSpecifier = (addon: string) => {
+  const groups = /^(@{0,1}[^@]+)(?:@(.+))?$/.exec(addon);
   if (groups) {
-    return [groups[0], groups[2]] as const;
+    return [groups[1], groups[2]] as const;
   }
   return [addon, undefined] as const;
 };
@@ -58,6 +46,8 @@ const checkInstalled = (addonName: string, main: any) => {
   return !!existingAddon;
 };
 
+const isCoreAddon = (addonName: string) => Object.hasOwn(versions, addonName);
+
 /**
  * Install the given addon package and add it to main.js
  *
@@ -71,9 +61,11 @@ const checkInstalled = (addonName: string, main: any) => {
  */
 export async function add(
   addon: string,
-  options: { packageManager: PackageManagerName; skipPostinstall: boolean }
+  options: { packageManager: PackageManagerName; skipPostinstall: boolean },
+  logger = console
 ) {
   const { packageManager: pkgMgr } = options;
+  const [addonName, inputVersion] = getVersionSpecifier(addon);
 
   const packageManager = JsPackageManagerFactory.getPackageManager({ force: pkgMgr });
   const packageJson = await packageManager.retrievePackageJson();
@@ -85,43 +77,52 @@ export async function add(
     `);
   }
 
-  if (checkInstalled(addon, requireMain(configDir))) {
-    throw new Error(dedent`
-      Addon ${addon} is already installed; we skipped adding it to your ${mainConfig}.
-    `);
-  }
-
-  const [addonName, versionSpecifier] = getVersionSpecifier(addon);
-
   if (!mainConfig) {
     logger.error('Unable to find storybook main.js config');
     return;
   }
+
+  if (checkInstalled(addonName, requireMain(configDir))) {
+    throw new Error(dedent`
+      Addon ${addonName} is already installed; we skipped adding it to your ${mainConfig}.
+    `);
+  }
+
   const main = await readConfig(mainConfig);
   logger.log(`Verifying ${addonName}`);
-  const latestVersion = await packageManager.latestVersion(addonName);
-  if (!latestVersion) {
-    logger.error(`Unknown addon ${addonName}`);
-  }
 
-  // add to package.json
-  const isStorybookAddon = addonName.startsWith('@storybook/');
-  const isAddonFromCore = isCorePackage(addonName);
   const storybookVersion = await getCoercedStorybookVersion(packageManager);
-  const version = versionSpecifier || (isAddonFromCore ? storybookVersion : latestVersion);
 
-  const addonWithVersion = SemVer.valid(version)
+  let version = inputVersion;
+
+  if (!version && isCoreAddon(addonName) && storybookVersion) {
+    version = storybookVersion;
+  }
+  if (!version) {
+    version = await packageManager.latestVersion(addonName);
+  }
+
+  if (isCoreAddon(addonName) && version !== storybookVersion) {
+    logger.warn(
+      `The version of ${addonName} you are installing is not the same as the version of Storybook you are using. This may lead to unexpected behavior.`
+    );
+  }
+
+  const addonWithVersion = isValidVersion(version)
     ? `${addonName}@^${version}`
     : `${addonName}@${version}`;
+
   logger.log(`Installing ${addonWithVersion}`);
   await packageManager.addDependencies({ installAsDevDependencies: true }, [addonWithVersion]);
 
-  // add to main.js
   logger.log(`Adding '${addon}' to main.js addons field.`);
   main.appendValueToArray(['addons'], addonName);
   await writeConfig(main);
 
-  if (!options.skipPostinstall && isStorybookAddon) {
+  if (!options.skipPostinstall && isCoreAddon(addonName)) {
     await postinstallAddon(addonName, { packageManager: packageManager.type });
   }
 }
+function isValidVersion(version: string) {
+  return SemVer.valid(version) || version.match(/^\d+$/);
+}

code/lib/cli/src/postinstallAddon.ts

@@ -0,0 +1,19 @@
+import type { PostinstallOptions } from './add';
+
+export const postinstallAddon = async (addonName: string, options: PostinstallOptions) => {
+  try {
+    const modulePath = require.resolve(`${addonName}/postinstall`, { paths: [process.cwd()] });
+
+    const postinstall = require(modulePath);
+
+    try {
+      console.log(`Running postinstall script for ${addonName}`);
+      await postinstall(options);
+    } catch (e) {
+      console.error(`Error running postinstall script for ${addonName}`);
+      console.error(e);
+    }
+  } catch (e) {
+    // no postinstall script
+  }
+};