From f7994816410f15805bbfe20d708ca495f3f38edf Mon Sep 17 00:00:00 2001
From: Marcos Passos <marcos@croct.com>
Date: Tue, 10 Feb 2026 19:25:08 -0300
Subject: [PATCH] Fix Storyblok init codemod

---
 .../javascript/storyblokInitCodemod.ts        | 45 ++++-----
 .../javascript/storyblokInitCodemod.test.ts   | 96 +++++++++++++++++++
 2 files changed, 119 insertions(+), 22 deletions(-)

diff --git a/src/application/project/code/transformation/javascript/storyblokInitCodemod.ts b/src/application/project/code/transformation/javascript/storyblokInitCodemod.ts
index 251f060..d8c5d5a 100644
--- a/src/application/project/code/transformation/javascript/storyblokInitCodemod.ts
+++ b/src/application/project/code/transformation/javascript/storyblokInitCodemod.ts
@@ -45,33 +45,34 @@ export class StoryblokInitCodemod implements Codemod<t.File, StoryblokInitCodemo
             CallExpression: path => {
                 const {callee} = path.node;
 
-                // Match direct function calls: targetFunction(...)
-                if (t.isIdentifier(callee) && callee.name === 'storyblokInit') {
-                    path.node.arguments = [
-                        t.callExpression(
-                            t.identifier(wrapperName),
-                            path.node.arguments,
-                        ),
-                    ];
-
-                    modified = true;
+                const isDirectCall = t.isIdentifier(callee) && callee.name === 'storyblokInit';
+                const isMemberCall = t.isMemberExpression(callee)
+                    && t.isIdentifier(callee.property)
+                    && callee.property.name === 'storyblokInit';
+
+                if (!isDirectCall && !isMemberCall) {
+                    return;
                 }
 
-                // Match member expression calls: obj.targetFunction(...)
+                const args = path.node.arguments;
+
                 if (
-                    t.isMemberExpression(callee)
-                    && t.isIdentifier(callee.property)
-                    && callee.property.name === 'storyblokInit'
+                    args.length === 1
+                    && t.isCallExpression(args[0])
+                    && t.isIdentifier(args[0].callee)
+                    && args[0].callee.name === wrapperName
                 ) {
-                    path.node.arguments = [
-                        t.callExpression(
-                            t.identifier(wrapperName),
-                            path.node.arguments,
-                        ),
-                    ];
-
-                    modified = true;
+                    return;
                 }
+
+                path.node.arguments = [
+                    t.callExpression(
+                        t.identifier(wrapperName),
+                        args,
+                    ),
+                ];
+
+                modified = true;
             },
         });
 
diff --git a/test/application/project/code/transformation/javascript/storyblokInitCodemod.test.ts b/test/application/project/code/transformation/javascript/storyblokInitCodemod.test.ts
index 7ebd38a..0046e74 100644
--- a/test/application/project/code/transformation/javascript/storyblokInitCodemod.test.ts
+++ b/test/application/project/code/transformation/javascript/storyblokInitCodemod.test.ts
@@ -199,6 +199,102 @@ describe('StoryblokInitCodemod', () => {
         ].join('\n'));
     });
 
+    it('should not wrap arguments that are already wrapped', async () => {
+        const transformer = createTransformer();
+
+        const input = [
+            "import { withCroct } from '@croct/storyblok';",
+            "import { storyblokInit } from '@storyblok/js';",
+            '',
+            'storyblokInit(withCroct({ accessToken: "token" }));',
+        ].join('\n');
+
+        const {result, modified} = await transformer.apply(input, {
+            name: 'withCroct',
+            module: '@croct/storyblok',
+        });
+
+        expect(modified).toBe(false);
+        expect(result).toBe(input);
+    });
+
+    it('should not wrap member expression arguments that are already wrapped', async () => {
+        const transformer = createTransformer();
+
+        const input = [
+            "import { withCroct } from '@croct/storyblok';",
+            "import * as sb from '@storyblok/js';",
+            '',
+            'sb.storyblokInit(withCroct({ accessToken: "token" }));',
+        ].join('\n');
+
+        const {result, modified} = await transformer.apply(input, {
+            name: 'withCroct',
+            module: '@croct/storyblok',
+        });
+
+        expect(modified).toBe(false);
+        expect(result).toBe(input);
+    });
+
+    it('should not wrap arguments that are already wrapped with an alias', async () => {
+        const transformer = createTransformer();
+
+        const input = [
+            "import { withCroct as croctWrapper } from '@croct/storyblok';",
+            "import { storyblokInit } from '@storyblok/js';",
+            '',
+            'storyblokInit(croctWrapper({ accessToken: "token" }));',
+        ].join('\n');
+
+        const {result, modified} = await transformer.apply(input, {
+            name: 'withCroct',
+            module: '@croct/storyblok',
+        });
+
+        expect(modified).toBe(false);
+        expect(result).toBe(input);
+    });
+
+    it('should not wrap arguments that are already wrapped with no arguments', async () => {
+        const transformer = createTransformer();
+
+        const input = [
+            "import { withCroct } from '@croct/storyblok';",
+            "import { storyblokInit } from '@storyblok/js';",
+            '',
+            'storyblokInit(withCroct());',
+        ].join('\n');
+
+        const {result, modified} = await transformer.apply(input, {
+            name: 'withCroct',
+            module: '@croct/storyblok',
+        });
+
+        expect(modified).toBe(false);
+        expect(result).toBe(input);
+    });
+
+    it('should wrap when the argument is a call to a different function', async () => {
+        const transformer = createTransformer();
+
+        const input = [
+            "import { storyblokInit } from '@storyblok/js';",
+            '',
+            'storyblokInit(otherWrapper({ accessToken: "token" }));',
+        ].join('\n');
+
+        const {result, modified} = await transformer.apply(input, {
+            name: 'withCroct',
+            module: '@croct/storyblok',
+        });
+
+        expect(modified).toBe(true);
+        expect(result).toContain(
+            'storyblokInit(withCroct(otherWrapper({ accessToken: "token" })));',
+        );
+    });
+
     it('should add empty statement before import when first statement is not an import', async () => {
         const transformer = createTransformer();