Merge pull request #30060 from storybookjs/valentin/improve-addon-a11…

…y-addon-test-automigration Automigration: Improve setup file transformation and version range handling for a11y migration
storybookjs · Dec 15, 2024 · 4192dde · 4192dde
2 parents 22bcecd + b9550a9
commit 4192dde
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 10 deletions.
diff --git a/code/lib/cli-storybook/src/automigrate/fixes/addon-a11y-addon-test.test.ts b/code/lib/cli-storybook/src/automigrate/fixes/addon-a11y-addon-test.test.ts
@@ -141,7 +141,8 @@ describe('addonA11yAddonTest', () => {
       `;
       vi.mocked(readFileSync).mockReturnValue(source);
 
-      const transformedCode = transformSetupFile(setupFile);
+      const s = readFileSync(setupFile, 'utf8');
+      const transformedCode = transformSetupFile(s);
       expect(transformedCode).toMatchInlineSnapshot(`
         "import * as a11yAddonAnnotations from "@storybook/addon-a11y/preview";
         import { beforeAll } from 'vitest';
@@ -169,7 +170,35 @@ describe('addonA11yAddonTest', () => {
       `;
       vi.mocked(readFileSync).mockReturnValue(source);
 
-      const transformedCode = transformSetupFile(setupFile);
+      const s = readFileSync(setupFile, 'utf8');
+      const transformedCode = transformSetupFile(s);
+      expect(transformedCode).toMatchInlineSnapshot(`
+        "import * as a11yAddonAnnotations from "@storybook/addon-a11y/preview";
+        import { beforeAll } from 'vitest';
+        import { setProjectAnnotations } from 'storybook';
+        import * as projectAnnotations from './preview';
+
+        const project = setProjectAnnotations([a11yAddonAnnotations, projectAnnotations]);
+
+        beforeAll(project.beforeAll);"
+      `);
+    });
+
+    it('should transform setup file correctly - project annotation is not an array', () => {
+      const setupFile = '/path/to/vitest.setup.ts';
+      const source = dedent`
+        import { beforeAll } from 'vitest';
+        import { setProjectAnnotations } from 'storybook';
+        import * as projectAnnotations from './preview';
+
+        const project = setProjectAnnotations(projectAnnotations);
+
+        beforeAll(project.beforeAll);
+      `;
+      vi.mocked(readFileSync).mockReturnValue(source);
+
+      const s = readFileSync(setupFile, 'utf8');
+      const transformedCode = transformSetupFile(s);
       expect(transformedCode).toMatchInlineSnapshot(`
         "import * as a11yAddonAnnotations from "@storybook/addon-a11y/preview";
         import { beforeAll } from 'vitest';

diff --git a/code/lib/cli-storybook/src/automigrate/fixes/addon-a11y-addon-test.ts b/code/lib/cli-storybook/src/automigrate/fixes/addon-a11y-addon-test.ts
@@ -22,8 +22,7 @@ interface AddonA11yAddonTestOptions {
  */
 export const addonA11yAddonTest: Fix<AddonA11yAddonTestOptions> = {
   id: 'addonA11yAddonTest',
-  // TODO: Change to the correct version after testing
-  versionRange: ['<8.5.0', '*'],
+  versionRange: ['<8.5.0', '>=8.5.0'],
 
   promptType(result) {
     if (result.setupFile === null) {
@@ -53,7 +52,11 @@ export const addonA11yAddonTest: Fix<AddonA11yAddonTestOptions> = {
 
     try {
       if (vitestSetupFile) {
-        const transformedSetupCode = transformSetupFile(vitestSetupFile);
+        const source = readFileSync(vitestSetupFile, 'utf8');
+        if (source.includes('@storybook/addon-a11y')) {
+          return null;
+        }
+        const transformedSetupCode = transformSetupFile(source);
         return {
           setupFile: vitestSetupFile,
           transformedSetupCode,
@@ -124,8 +127,7 @@ export const addonA11yAddonTest: Fix<AddonA11yAddonTestOptions> = {
   },
 };
 
-export function transformSetupFile(setupFile: string) {
-  const source = readFileSync(setupFile, 'utf8');
+export function transformSetupFile(source: string) {
   const j = jscodeshift.withParser('ts');
 
   const root = j(source);
@@ -148,9 +150,14 @@ export function transformSetupFile(setupFile: string) {
     throw new Error('Could not find setProjectAnnotations call in vitest.setup file');
   }
 
-  // Add a11yAddonAnnotations to the annotations array
-  setProjectAnnotationsCall.find(j.ArrayExpression).forEach((p) => {
-    p.value.elements.unshift(j.identifier('a11yAddonAnnotations'));
+  // Add a11yAddonAnnotations to the annotations array or create a new array if argument is a string
+  setProjectAnnotationsCall.forEach((p) => {
+    if (p.value.arguments.length === 1 && p.value.arguments[0].type === 'ArrayExpression') {
+      p.value.arguments[0].elements.unshift(j.identifier('a11yAddonAnnotations'));
+    } else if (p.value.arguments.length === 1 && p.value.arguments[0].type === 'Identifier') {
+      const arg = p.value.arguments[0];
+      p.value.arguments[0] = j.arrayExpression([j.identifier('a11yAddonAnnotations'), arg]);
+    }
   });
 
   // Add the import declaration at the top