@codeyam/codeyam-cli 0.1.0-staging.1669d45 → 0.1.0-staging.1a2737b

package/background/src/lib/virtualized/project/constructMockCode.js

@@ -225,6 +225,18 @@ function isValidKey(key) {
         return false;
     return !/\s/.test(keyWithOutArguments);
 }
+/**
+ * Known hooks that return tuples [value, setter] instead of arrays.
+ * These should NOT use the .map() pattern even when the schema has generic array access ([]).
+ * Instead, they should return [data, () => {}] where data is from scenarios().
+ */
+const TUPLE_RETURNING_HOOKS = new Set([
+    'useAtom', // Jotai
+    'useState', // React
+    'useReducer', // React
+    'useRecoilState', // Recoil
+    'useImmerAtom', // Jotai with Immer
+]);
 export default function constructMockCode(mockName, dependencySchemas, entityType, _canonicalKey, // DEPRECATED: No longer used, kept for API compatibility
 options) {
     // Check if mockName is a call signature (e.g., "useFetcher<User>()", "db.select(query)")
@@ -245,7 +257,7 @@ options) {
     let dataStructureValue;
     let foundEntityWithSignature = false;
     let signatureSchema;
-    for (const filePath in dependencySchemas) {
+    entitySearch: for (const filePath in dependencySchemas) {
         for (const entityName in dependencySchemas[filePath]) {
             // Match entity by base name (without generics/args)
             const entityBaseName = entityName.split(/[<(]/)[0];
@@ -277,7 +289,7 @@ options) {
                 // However, we still need to remove duplicate function calls that create invalid syntax
                 removeDuplicateFunctionCalls(relevantReturnValueSchema);
                 dataStructureValue = relevantReturnValueSchema?.[dataStructurePath];
-                break;
+                break entitySearch;
             }
         }
     }
@@ -399,6 +411,87 @@ options) {
     const pathDepth = splitOutsideParenthesesAndArrays(dataStructurePath).length;
     const isRootArray = dataStructureValue === 'array' &&
         (dataStructurePath === 'returnValue' || pathDepth <= mockNameParts.length);
+    // OPTIMIZATION: Early return for tuple-returning hooks (useAtom, useState, etc.)
+    // These hooks have simple [value, setter] return patterns that don't need the full
+    // 9216-key schema processing. Check if this is a tuple-returning hook and generate
+    // the mock code directly without iterating over all schema keys.
+    if (TUPLE_RETURNING_HOOKS.has(baseMockName) && isFunction) {
+        // Check if schema has generic array pattern (indicates tuple return like [value, setter])
+        const schemaKeys = Object.keys(relevantReturnValueSchema ?? {});
+        const hasGenericArrayInSchema = schemaKeys.some((k) => k.includes('.functionCallReturnValue[]') ||
+            k === `${dataKey}.functionCallReturnValue[]` ||
+            k === 'returnValue[]');
+        // Check for differentiated tuple indices (e.g., functionCallReturnValue[2], [3]) which would NOT be a standard tuple
+        // We only check indices immediately after functionCallReturnValue, not nested indices like signature[2]
+        const tupleHasDifferentiatedIndices = schemaKeys.some((k) => {
+            // Look for .functionCallReturnValue[N] where N >= 2
+            const match = k.match(/\.functionCallReturnValue\[(\d+)\]/);
+            if (!match)
+                return false;
+            const idx = parseInt(match[1], 10);
+            return idx >= 2;
+        });
+        const isTupleReturningHook = hasGenericArrayInSchema && !tupleHasDifferentiatedIndices;
+        if (isTupleReturningHook) {
+            // Find all call patterns for this hook (e.g., useAtom(quoteFilterAtom), useAtom(supplierAtom))
+            const hookCallPatterns = Object.keys(relevantReturnValueSchema ?? {})
+                .filter((k) => {
+                const regex = new RegExp(`^${baseMockName}\\([^)]+\\)$`);
+                return regex.test(k);
+            })
+                .map((k) => {
+                const match = k.match(/\(([^)]+)\)/);
+                return { key: k, arg: match?.[1] ?? '' };
+            });
+            let tupleReturnCode;
+            if (hookCallPatterns.length > 1) {
+                // Multiple patterns - generate conditional dispatch
+                const conditions = hookCallPatterns
+                    .map(({ key, arg }) => `if (argLabel === '${arg}' || argStr.includes('${arg}')) {\n    return [scenarios().data()?.["${key}"]?.[0] ?? [], () => {}];\n  }`)
+                    .join('\n  ');
+                const fallbackKey = hookCallPatterns[0]?.key ?? dataKey;
+                tupleReturnCode = `(() => {
+  // Dynamic dispatch for tuple-returning hook with multiple argument patterns
+  const argLabel = args[0]?.debugLabel ?? '';
+  const argStr = args[0]?.toString?.() ?? String(args[0] ?? '');
+  ${conditions}
+  return [scenarios().data()?.["${fallbackKey}"]?.[0] ?? [], () => {}];
+})()`;
+            }
+            else {
+                // Single or no patterns - use dynamic dispatch
+                const fallbackKey = hookCallPatterns[0]?.key ?? `${baseMockName}()`;
+                tupleReturnCode = `(() => {
+  // Dynamic dispatch for tuple-returning hook
+  const argLabel = args[0]?.debugLabel ?? '';
+  const argStr = args[0]?.toString?.() ?? '';
+  const allData = scenarios().data() ?? {};
+  if (argLabel) {
+    const labelKey = '${baseMockName}(' + argLabel + ')';
+    if (allData[labelKey]) {
+      return [allData[labelKey]?.[0] ?? [], () => {}];
+    }
+  }
+  const keys = Object.keys(allData).filter(k => k.startsWith('${baseMockName}('));
+  for (const key of keys) {
+    const keyArg = key.slice(${baseMockName.length + 1}, -1);
+    if (argStr.includes(keyArg)) {
+      return [allData[key]?.[0] ?? [], () => {}];
+    }
+  }
+  return [allData[keys[0] ?? '${fallbackKey}']?.[0] ?? [], () => {}];
+})()`;
+            }
+            const safeFunctionName = options?.uniqueFunctionSuffix
+                ? `${baseMockName}_${options.uniqueFunctionSuffix}`
+                : options?.keepOriginalFunctionName
+                    ? baseMockName
+                    : mockNameIsCallSignature && derivedFunctionName
+                        ? derivedFunctionName
+                        : baseMockName;
+            return `function ${safeFunctionName}(...args) {\n  return ${tupleReturnCode};\n}`;
+        }
+    }
     const returnValueParts = {
         name: dataStructureName,
         isArray: isRootArray,
@@ -679,7 +772,102 @@ options) {
                     return ARRAY_PROTOTYPE_METHODS.has(methodName);
                 });
             let returnValueContents = '';
-            if (!returnsFunctionArgs &&
+            // Check if this is a known tuple-returning hook (useAtom, useState, etc.)
+            // These should return [value, setter] tuples, not arrays or data paths
+            // Check isGenericArray from current context OR from schema for root level calls
+            // (at root level, isGenericArray might not be set yet but the schema contains [] pattern)
+            const hasGenericArrayInSchema = root &&
+                TUPLE_RETURNING_HOOKS.has(baseMockName) &&
+                Object.keys(relevantReturnValueSchema ?? {}).some((k) => k.includes('.functionCallReturnValue[]'));
+            // Check if there are array indices beyond what a standard 2-element tuple would have
+            // For tuple-returning hooks, [0] and [1] are expected (value and setter)
+            // Only consider it "differentiated" if there are indices >= 2 (e.g., [2], [3])
+            const tupleHasDifferentiatedIndices = nested?.some((n) => {
+                const indexMatch = n.name.match(/^\[(\d+)\]$/);
+                if (!indexMatch)
+                    return false;
+                const index = parseInt(indexMatch[1], 10);
+                return index >= 2;
+            });
+            const isTupleReturningHook = TUPLE_RETURNING_HOOKS.has(baseMockName) &&
+                (isGenericArray || hasGenericArrayInSchema) &&
+                !tupleHasDifferentiatedIndices;
+            // Debug logging for tuple-returning hooks
+            if (TUPLE_RETURNING_HOOKS.has(baseMockName) && root) {
+                const schemaKeys = Object.keys(relevantReturnValueSchema ?? {});
+                const hasArrayPattern = schemaKeys.some((k) => k.includes('.functionCallReturnValue[]'));
+                console.log(`CodeYam: Tuple hook check for ${baseMockName} (root):`, `hasGenericArrayInSchema=${hasGenericArrayInSchema}`, `hasArrayPattern=${hasArrayPattern}`, `tupleHasDifferentiatedIndices=${tupleHasDifferentiatedIndices}`, `isTupleReturningHook=${isTupleReturningHook}`, `schemaKeysSample=${schemaKeys.slice(0, 5).join(', ')}`);
+            }
+            if (isTupleReturningHook) {
+                // Tuple-returning hooks should return [value, setter] tuple
+                // The value is the first element from scenarios data, setter is a no-op
+                // Default to [] when data is undefined to prevent errors like ".includes is not a function"
+                // Check if there are multiple call patterns for this hook in the schema
+                // (e.g., useAtom(quoteFilterAtom) and useAtom(supplierAtom))
+                const hookCallPatterns = Object.keys(relevantReturnValueSchema ?? {})
+                    .filter((k) => {
+                    // Match patterns like "useAtom(someArg)" but not nested paths like "useAtom(x).foo"
+                    const regex = new RegExp(`^${baseMockName}\\([^)]+\\)$`);
+                    return regex.test(k);
+                })
+                    .map((k) => {
+                    // Extract the argument from the key like "useAtom(quoteFilterAtom)" -> "quoteFilterAtom"
+                    const match = k.match(/\(([^)]+)\)/);
+                    return { key: k, arg: match?.[1] ?? '' };
+                });
+                if (hookCallPatterns.length > 1) {
+                    // Multiple patterns - generate conditional dispatch based on first argument
+                    // For Jotai atoms, we use debugLabel; for others, we try to match the argument string
+                    const conditions = hookCallPatterns
+                        .map(({ key, arg }) => `if (argLabel === '${arg}' || argStr.includes('${arg}')) {\n    return [scenarios().data()?.["${key}"]?.[0] ?? [], () => {}];\n  }`)
+                        .join('\n  ');
+                    // Use the first pattern as fallback
+                    const fallbackKey = hookCallPatterns[0]?.key ?? dataKey;
+                    returnValueContents = `(() => {
+  // Dynamic dispatch for tuple-returning hook with multiple argument patterns
+  const argLabel = args[0]?.debugLabel ?? '';
+  const argStr = args[0]?.toString?.() ?? String(args[0] ?? '');
+  ${conditions}
+  // Fallback to first pattern
+  return [scenarios().data()?.["${fallbackKey}"]?.[0] ?? [], () => {}];
+})()`;
+                }
+                else {
+                    // Single pattern or no patterns - use dynamic dispatch to handle case where
+                    // the mock is used with different atoms than what was captured in the schema.
+                    // Use the first argument to construct the data key dynamically.
+                    const fallbackKey = hookCallPatterns[0]?.key ?? `${baseMockName}()`;
+                    returnValueContents = `(() => {
+  // Dynamic dispatch for tuple-returning hook
+  // Try to construct key from argument's debugLabel (Jotai atoms) or toString
+  const argLabel = args[0]?.debugLabel ?? '';
+  const argStr = args[0]?.toString?.() ?? '';
+  const allData = scenarios().data() ?? {};
+
+  // Try to find a matching key using debugLabel first
+  if (argLabel) {
+    const labelKey = '${baseMockName}(' + argLabel + ')';
+    if (allData[labelKey]) {
+      return [allData[labelKey]?.[0] ?? [], () => {}];
+    }
+  }
+
+  // Try to find any matching key that contains part of the argument string
+  const keys = Object.keys(allData).filter(k => k.startsWith('${baseMockName}('));
+  for (const key of keys) {
+    const keyArg = key.slice(${baseMockName.length + 1}, -1);
+    if (argStr.includes(keyArg)) {
+      return [allData[key]?.[0] ?? [], () => {}];
+    }
+  }
+
+  // Fallback to first matching key or default
+  const fallback = keys[0] ?? '${fallbackKey}';
+  return [allData[fallback]?.[0] ?? [], () => {}];
+})()`;
+                }
+            }
+            else if (!returnsFunctionArgs &&
                 nestedContent.length === 0 &&
                 dataPaths.length === 1) {
                 returnValueContents = dataPaths[0];
@@ -769,11 +957,19 @@ options) {
                                     // Valid patterns where arrow function is NOT bare:
                                     // 1. Property value: "key: (arg) => ..." - ends with ':'
                                     // 2. Function argument: ".map((arg) => ..." - ends with '('
+                                    // 3. Method call: "?.map" followed directly by the arrow function
+                                    //    In this case, the '(' is consumed by the arrow function regex match,
+                                    //    so beforeTrimmed ends with the method name (e.g., 'map'), not '('.
+                                    //    We detect this by checking if beforeTrimmed ends with an identifier
+                                    //    that could be a method name (preceded by '.' or '?.').
                                     // NOTE: We don't include ',' because "{ prop, () => {} }" is invalid
                                     // (can't distinguish function argument from object property context)
                                     const isPropertyValue = beforeTrimmed.endsWith(':');
                                     const isFunctionArg = beforeTrimmed.endsWith('(');
-                                    const hasPropertyName = isPropertyValue || isFunctionArg;
+                                    // Check if before ends with a method call pattern like ".map" or "?.map"
+                                    // The '(' after the method name is consumed by the arrow function regex
+                                    const isMethodCallArg = /\??\.\w+$/.test(beforeTrimmed);
+                                    const hasPropertyName = isPropertyValue || isFunctionArg || isMethodCallArg;
                                     if (!hasPropertyName) {
                                         // This is a bare arrow function - filter it out
                                         // Found arrow function start, need to find its end
@@ -974,6 +1170,16 @@ options) {
                                     returnValueContents = `${arrayBasePath}?.map((__item__, __idx__) => (\n${indent(trimmedContent)}\n))`;
                                 }
                                 else {
+                                    // When generating object-wrapped .map(), ensure original item data is preserved.
+                                    // If no data spread was included (e.g., because this is a plain array property,
+                                    // not a function return), add ...__item__ to spread the original item properties.
+                                    // Without this, the .map() would create new objects with only nested function
+                                    // properties, losing data like filePath, frontmatter, body, etc.
+                                    const hasDataSpread = mappedContents.includes('...scenarios()') ||
+                                        mappedContents.includes('...__item__');
+                                    if (!hasDataSpread) {
+                                        mappedContents = `...__item__,\n${mappedContents}`;
+                                    }
                                     returnValueContents = `${arrayBasePath}?.map((__item__, __idx__) => ({\n${indent(mappedContents)}\n}))`;
                                 }
                             } // Close the empty content check else block
@@ -1035,6 +1241,13 @@ options) {
                             mappedContents = filterOutArrowFunctions(mappedContents);
                             mappedContents = filterOutBareObjects(mappedContents, true);
                             mappedContents = cleanupContent(mappedContents);
+                            // Same as needsWrapper branch: ensure item data is preserved in .map()
+                            const hasDataSpreadInner = mappedContents.includes('...scenarios()') ||
+                                mappedContents.includes('...__item__');
+                            if (!hasDataSpreadInner && mappedContents.trim().length > 0) {
+                                // Insert ...__item__ after the opening brace
+                                mappedContents = mappedContents.replace(/^\s*\{/, '{\n...__item__,');
+                            }
                             returnValueContents = `${arrayBasePath}?.map((__item__, __idx__) => (\n${indent(mappedContents)}\n))`;
                         }
                     }
@@ -1043,7 +1256,36 @@ options) {
                     }
                 }
                 else {
-                    returnValueContents = `{\n${indent(levelContents)}\n}`;
+                    // When we have a single data path and nested content that creates an object structure,
+                    // and we're NOT at the root level, we need to handle the case where the parent data
+                    // value is null or undefined. Without this check, `{ ...null, prop: null?.["prop"] }`
+                    // creates `{ prop: undefined }` instead of `null`, causing errors like
+                    // "Cannot read properties of undefined (reading 'some')" when code does
+                    // data?.prop.some(...) because data is an object with prop: undefined, not null.
+                    // We only apply this to non-root cases because root-level mocks are expected to exist.
+                    // We also skip structural elements (like [0] inside arrays) because the null check
+                    // syntax doesn't work inside .map() callbacks where structural elements are used.
+                    // We also skip array index elements ([0], [1], etc.) because they represent tuple/array
+                    // elements, not properties that could be null.
+                    // We also only apply this when we're inside a function return value context - i.e.,
+                    // when the data path contains a function call pattern like ?.["someFunction(...)"].
+                    // This prevents adding null checks to intermediate objects in chains like supabase.auth.
+                    const hasNestedContent = nestedContent.filter(Boolean).length > 0;
+                    const isArrayIndexElement = name.match(/^\[\d*\]$/);
+                    // Check if data path contains a function call pattern, indicating we're inside a function return value
+                    const isInsideFunctionReturnValue = dataPaths.length === 1 &&
+                        dataPaths[0].match(/\?\.\["\w+\([^"]*\)"\]/);
+                    if (!root &&
+                        !returnValue.isStructural &&
+                        !isArrayIndexElement &&
+                        isInsideFunctionReturnValue &&
+                        hasNestedContent) {
+                        // Wrap with null check: if parent is null/undefined, return it directly; otherwise create object
+                        returnValueContents = `${dataPaths[0]} == null ? ${dataPaths[0]} : {\n${indent(levelContents)}\n}`;
+                    }
+                    else {
+                        returnValueContents = `{\n${indent(levelContents)}\n}`;
+                    }
                 }
             }
             if (root) {
@@ -1158,7 +1400,31 @@ options) {
         return content;
     };
     // Create the return value structure
-    const sortedKeys = Object.keys(relevantReturnValueSchema ?? {}).sort((a, b) => {
+    // OPTIMIZATION: Filter keys to only those starting with baseMockName before sorting.
+    // This dramatically reduces processing time for large schemas (e.g., 9216 keys -> ~100 relevant keys).
+    // Without this filter, the loop would call splitOutsideParenthesesAndArrays on every key
+    // even though most are filtered out later by the baseMockName check.
+    const allSchemaKeys = Object.keys(relevantReturnValueSchema ?? {});
+    const relevantKeys = allSchemaKeys.filter((key) => {
+        // Fast prefix check - key must start with baseMockName followed by ( or < or .
+        // This matches: "useAtom()", "useAtom<T>()", "useAtom.something", but not "useAtomValue()"
+        if (key === baseMockName)
+            return true;
+        if (key.startsWith(baseMockName + '('))
+            return true;
+        if (key.startsWith(baseMockName + '<'))
+            return true;
+        if (key.startsWith(baseMockName + '.'))
+            return true;
+        // Also include 'returnValue' paths which are normalized later
+        if (key === 'returnValue' ||
+            key.startsWith('returnValue.') ||
+            key.startsWith('returnValue['))
+            return true;
+        return false;
+    });
+    const schemaKeyCount = relevantKeys.length;
+    const sortedKeys = relevantKeys.sort((a, b) => {
         const aParts = splitOutsideParenthesesAndArrays(a);
         const bParts = splitOutsideParenthesesAndArrays(b);
         const maxLength = Math.max(aParts.length, bParts.length);
@@ -1182,6 +1448,36 @@ options) {
         }
         return 0;
     });
+    // OPTIMIZATION: Pre-compute prefix indexes for O(1) lookups instead of O(n) scans.
+    // This reduces complexity from O(n²) to O(n) for large schemas (9k+ keys).
+    //
+    // 1. extendedReturnValuePrefixes: Set of all path prefixes that have a .functionCallReturnValue extension
+    //    Used by hasExtendedFunctionCallReturnValue check at line ~1754
+    // 2. functionCallsWithReturnValue: Set of function call paths where .functionCallReturnValue IMMEDIATELY follows
+    //    Used by hasProperFunctionCallPath check at line ~1787
+    //    IMPORTANT: Only includes paths where the function call is directly followed by .functionCallReturnValue
+    //    e.g., "a.b().functionCallReturnValue" -> adds "a.b()" but NOT "a" even if "a" ends with ")"
+    const extendedReturnValuePrefixes = new Set();
+    const functionCallsWithReturnValue = new Set();
+    for (const k of relevantKeys) {
+        const parts = splitOutsideParenthesesAndArrays(k);
+        const returnValueIndex = parts.findIndex((part) => part.startsWith(RETURN_VALUE));
+        if (returnValueIndex !== -1) {
+            // Add all prefixes of k up to (but not including) functionCallReturnValue
+            const prefix = joinParenthesesAndArrays(parts.slice(0, returnValueIndex));
+            extendedReturnValuePrefixes.add(prefix);
+            // ONLY add to functionCallsWithReturnValue if functionCallReturnValue IMMEDIATELY follows
+            if (prefix.endsWith(')')) {
+                functionCallsWithReturnValue.add(prefix);
+            }
+            // Also add intermediate prefixes for nested paths to extendedReturnValuePrefixes
+            // This helps hasExtendedFunctionCallReturnValue which checks key + '.'
+            for (let i = 1; i < returnValueIndex; i++) {
+                const partialPrefix = joinParenthesesAndArrays(parts.slice(0, i));
+                extendedReturnValuePrefixes.add(partialPrefix);
+            }
+        }
+    }
     for (const key of sortedKeys) {
         const value = relevantReturnValueSchema[key];
         const parts = splitOutsideParenthesesAndArrays(key);
@@ -1227,7 +1523,9 @@ options) {
         // nested inside (e.g., methods on array elements passed as arguments).
         if (hasSignaturePath)
             continue;
-        const hasExtendedFunctionCallReturnValue = Object.keys(relevantReturnValueSchema).some((k) => k.startsWith(key + '.') && k.includes(`.${RETURN_VALUE}`));
+        // OPTIMIZATION: Use pre-computed index instead of O(n) scan
+        // Old code: Object.keys(relevantReturnValueSchema).some((k) => k.startsWith(key + '.') && k.includes('.functionCallReturnValue'))
+        const hasExtendedFunctionCallReturnValue = extendedReturnValuePrefixes.has(key);
         // Skip JSX components - they look like function calls (e.g., Context.Provider())
         // but they're React components used in JSX, not functions that need mocking
         // Check both the value type and whether the functionCallReturnValue is jsx-component
@@ -1252,7 +1550,9 @@ options) {
                     // This part is a function call, and the next part is NOT .functionCallReturnValue
                     // Check if there's any path with .functionCallReturnValue for this function call
                     const functionCallPath = joinParenthesesAndArrays(parts.slice(0, i + 1));
-                    const hasProperFunctionCallPath = Object.keys(relevantReturnValueSchema).some((k) => k.startsWith(functionCallPath + '.functionCallReturnValue'));
+                    // OPTIMIZATION: Use pre-computed index instead of O(n) scan
+                    // Old code: Object.keys(relevantReturnValueSchema).some((k) => k.startsWith(functionCallPath + '.functionCallReturnValue'))
+                    const hasProperFunctionCallPath = functionCallsWithReturnValue.has(functionCallPath);
                     if (hasProperFunctionCallPath) {
                         // Skip this path - the .functionCallReturnValue path will handle it correctly
                         shouldSkipKey = true;
@@ -1617,6 +1917,20 @@ options) {
             returnValueSection = relevantPart;
         }
     }
+    // Post-processing: When the root functionCallReturnValue is typed as "function" but the
+    // return value also has nested properties (methods like .from(), .auth, etc.), it's actually
+    // an object, not a function to be called. Clear returnsFunctionArgs to prevent double-wrapping
+    // (adding an extra () => { return { ... } } wrapper and ["()"] data paths).
+    // This handles cases like Supabase's createClient() which returns an object with methods.
+    // Only applied to the root level - nested parts that are functions with methods (like
+    // useSearchParams()[1] which is a setter function with .set() and .delete()) should keep
+    // their returnsFunctionArgs since they genuinely ARE functions.
+    if (returnValueParts.returnsFunctionArgs &&
+        returnValueParts.returnsFunctionArgs.length === 0 &&
+        returnValueParts.nested &&
+        returnValueParts.nested.length > 0) {
+        returnValueParts.returnsFunctionArgs = undefined;
+    }
     const contents = constructReturnValueString(returnValueParts);
     if (mockNameParts.length > 1) {
         const originalLib = `${mockNameParts[0]}__cyOriginal`;
@@ -1678,15 +1992,24 @@ options) {
             const nameAssignment = isErrorSubclass
                 ? `this.name = '${baseMockName}';`
                 : '';
-            // Use the safe function name for the class definition
-            const className = mockNameIsCallSignature
-                ? derivedFunctionName
-                : baseMockName;
+            // Use the base class name for the class definition, not the call-signature-derived name.
+            // When mockName is "StatsCalculator(supabase)", baseMockName is "StatsCalculator"
+            // and derivedFunctionName would be "StatsCalculator_supabase" which is wrong.
+            // Classes are instantiated with `new ClassName(args)` so the name must match the original.
+            const className = baseMockName;
+            // Use the already-generated contents (which has proper function wrappers for methods)
+            // instead of raw scenarios().data() which would create non-callable string-keyed properties.
+            // For classes with methods like calculateStats(), the contents will have:
+            //   { calculateStats: (...args) => scenarios().data()?.["key"]?.["calculateStats(...)"], ... }
+            // which makes methods callable on the instance.
+            const classContents = enhancedContents.trim().startsWith('{')
+                ? enhancedContents
+                : `scenarios().data()?.${quotePropertyKey(dataKey)} || {}`;
             return `class ${className}${isErrorSubclass ? ' extends Error' : ''} {
   constructor(message) {
     ${superCall}
     ${nameAssignment}
-    Object.assign(this, scenarios().data()?.${quotePropertyKey(dataKey)} || {});
+    Object.assign(this, ${classContents});
   }
 }`;
         }
@@ -1754,7 +2077,29 @@ options) {
             : enhancedContents;
         // Inline the return value directly in the function to avoid module-level const
         // that would be evaluated before scenario context is ready
-        return `${isRootAsyncFunction ? 'async ' : ''}function ${safeFunctionName}(...args) {\n${indent(`return ${returnValue};`)}\n}`;
+        // Add fallback for simple data path returns to prevent undefined errors (e.g., createTheme)
+        // Only add fallback if returnValue is a simple data accessor (starts with scenarios().data())
+        // and doesn't already have nested structure (object literal, array, or method chains like .map())
+        const isSimpleDataPath = returnValue.startsWith('scenarios().data()') &&
+            !returnValue.trim().startsWith('{') &&
+            !returnValue.trim().startsWith('[') &&
+            !returnValue.includes('.map('); // Exclude method chains
+        const safeReturnValue = isSimpleDataPath
+            ? `${returnValue} ?? {}`
+            : returnValue;
+        const refName = `_${safeFunctionName}Ref`;
+        const assignment = `${refName}.current = ${safeReturnValue};`;
+        const ifBlock = `if (!${refName}.current) {\n${indent(assignment)}\n}`;
+        const body = `${ifBlock}\nreturn ${refName}.current;`;
+        return [
+            `// PATCHED: memoize to return stable reference (prevents infinite useEffect re-triggers)`,
+            `const ${refName} = {`,
+            `  current: null,`,
+            `};`,
+            `${isRootAsyncFunction ? 'async ' : ''}function ${safeFunctionName}(...args) {`,
+            indent(body),
+            `}`,
+        ].join('\n');
     }
     else {
         // Generate safe const name: