@codeyam/codeyam-cli 0.1.0-staging.8aea589 → 0.1.0-staging.8e7b1bd

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

@@ -1,4 +1,35 @@
-import { joinParenthesesAndArrays, splitOutsideParenthesesAndArrays, functionArguments, cleanOutBoundary, fillInDirectSchemaGapsAndUnknowns, } from "../../../../../packages/ai/index.js";
+import { joinParenthesesAndArrays, splitOutsideParenthesesAndArrays, functionArguments, cleanOutBoundary, fillInDirectSchemaGapsAndUnknowns, removeDuplicateFunctionCalls, } from "../../../../../packages/ai/index.js";
+/**
+ * Converts a variable qualifier to a safe identifier for use in function/const names.
+ * Handles indexed format like "data[1]" -> "data_1"
+ */
+function safeVariableQualifier(qualifier) {
+    // Replace brackets with underscores for valid JS identifier
+    // "data[1]" -> "data_1"
+    return qualifier.replace(/\[(\d+)\]/g, '_$1');
+}
+/**
+ * Extract property names that are jsx-components and should be preserved from original.
+ * These are paths like "MockName.Provider()" where the value or functionCallReturnValue is 'jsx-component'.
+ */
+function getJsxComponentProperties(mockName, schema) {
+    const jsxProperties = [];
+    for (const [path, value] of Object.entries(schema ?? {})) {
+        // Look for paths like "MockName.PropertyName()" that are jsx-components
+        // Skip paths that have deeper nesting (e.g., "MockName.Foo.Bar()")
+        const match = path.match(new RegExp(`^${mockName}\\.([^.()]+)\\(\\)$`));
+        if (!match)
+            continue;
+        const propertyName = match[1];
+        // Check if this is a jsx-component (either directly or via functionCallReturnValue)
+        const isJsxComponent = value === 'jsx-component' ||
+            schema[path + '.functionCallReturnValue'] === 'jsx-component';
+        if (isJsxComponent && !jsxProperties.includes(propertyName)) {
+            jsxProperties.push(propertyName);
+        }
+    }
+    return jsxProperties;
+}
 /**
  * Check if a string is a literal (quoted string or number).
  * Returns the unquoted value if it's a literal, null otherwise.
@@ -109,15 +140,54 @@ function isValidKey(key) {
     const keyWithOutArguments = key.split('(')[0];
     return !/\s/.test(keyWithOutArguments);
 }
-export default function constructMockCode(mockName, dependencySchemas) {
-    const mockNameParts = splitOutsideParenthesesAndArrays(mockName);
+export default function constructMockCode(mockName, dependencySchemas, entityType) {
+    // Check for variable-qualified mock name (e.g., "entityDiffFetcher <- useFetcher")
+    // Format: "variableName <- functionName" which reads naturally as "variableName receives from functionName"
+    // Also handles indexed format: "data[1] <- useFetch" for disambiguating reused variable names
+    const variableQualifierMatch = mockName.match(/^([a-zA-Z_][a-zA-Z0-9_]*(?:\[\d+\])?)\s*<-\s*/);
+    let variableQualifier = variableQualifierMatch
+        ? variableQualifierMatch[1]
+        : null;
+    const baseMockName = variableQualifier
+        ? mockName.slice(variableQualifierMatch[0].length)
+        : mockName;
+    const mockNameParts = splitOutsideParenthesesAndArrays(baseMockName);
     let relevantReturnValueSchema;
     let dataStructurePath;
     let dataStructureValue;
+    let foundEntityWithSignature = false;
+    let signatureSchema;
     for (const filePath in dependencySchemas) {
         for (const entityName in dependencySchemas[filePath]) {
-            if (entityName !== mockNameParts[0])
+            // For variable-qualified mocks, look for the full qualified name first
+            // Format: "variableName <- functionName" (e.g., "entityDiffFetcher <- useFetcher")
+            const targetEntityName = variableQualifier
+                ? `${variableQualifier} <- ${baseMockName}`
+                : mockNameParts[0];
+            // Check for direct match
+            let matches = entityName === targetEntityName || entityName === mockNameParts[0];
+            // If no direct match and no qualifier was provided, check if the entity
+            // is stored under a variable-qualified key (e.g., "stateBadge <- getStateBadge")
+            // This handles the case where gatherDataForMocks stored the entity with a variable
+            // qualifier but writeScenarioComponents called constructMockCode without one.
+            if (!matches && !variableQualifier) {
+                const qualifiedKeyMatch = entityName.match(new RegExp(`^([a-zA-Z_][a-zA-Z0-9_]*)\\s*<-\\s*${mockNameParts[0]}$`));
+                if (qualifiedKeyMatch) {
+                    matches = true;
+                    // Extract the variable qualifier from the entity name so we can use
+                    // it for the data lookup key later
+                    variableQualifier = qualifiedKeyMatch[1];
+                }
+            }
+            if (!matches)
                 continue;
+            // Track if we found the entity and it has a signature (is a function)
+            const entitySchema = dependencySchemas[filePath][entityName];
+            if (entitySchema?.signatureSchema &&
+                Object.keys(entitySchema.signatureSchema).length > 0) {
+                foundEntityWithSignature = true;
+                signatureSchema = entitySchema.signatureSchema;
+            }
             dataStructurePath = Object.keys(dependencySchemas[filePath][entityName]?.returnValueSchema ?? {}).find((path) => {
                 if (path === 'returnValue')
                     return true;
@@ -131,16 +201,61 @@ export default function constructMockCode(mockName, dependencySchemas) {
                 relevantReturnValueSchema = fillInDirectSchemaGapsAndUnknowns({
                     schema: dependencySchemas[filePath][entityName]?.returnValueSchema,
                 });
-                // NOTE: Temporarily disabled clearAttributesFromMapping as it deletes
+                // NOTE: clearAttributesFromMapping is disabled because it deletes
                 // method calls on arrays (like .eq() after functionCallReturnValue: 'array')
-                // clearAttributesFromMapping(relevantReturnValueSchema)
+                // However, we still need to remove duplicate function calls that create invalid syntax
+                removeDuplicateFunctionCalls(relevantReturnValueSchema);
                 dataStructureValue = relevantReturnValueSchema?.[dataStructurePath];
                 break;
             }
         }
     }
-    if (!relevantReturnValueSchema || !dataStructurePath)
+    // Check if the return value schema only contains function type markers
+    // (e.g., "validateInputs()": "function") without actual return data
+    // (no functionCallReturnValue entries)
+    const hasMeaningfulReturnData = (schema) => {
+        if (!schema)
+            return false;
+        const keys = Object.keys(schema);
+        // Check if there are any functionCallReturnValue entries
+        // which indicate actual mock data, not just type markers
+        return keys.some((key) => key.includes('functionCallReturnValue') || key === 'returnValue');
+    };
+    // If we found the entity with a signature but no return value schema,
+    // OR the schema only has function type markers without actual return data,
+    // generate an empty mock function. This handles validation functions like
+    // validateInputs that are marked for mocking but have no meaningful return value.
+    // EXCEPTION: If entityType is 'class', we still need to generate the class.
+    // EXCEPTION: If no signature was found, let the normal code path handle it.
+    if (!relevantReturnValueSchema || !dataStructurePath) {
+        if (foundEntityWithSignature && signatureSchema) {
+            // Count the number of arguments from signature schema
+            const argCount = Object.keys(signatureSchema).filter((key) => key.startsWith('signature[')).length;
+            const args = Array.from({ length: argCount }, (_, i) => `arg${i + 1}`);
+            const argsString = args.join(', ');
+            // Generate empty mock function
+            return `function ${mockName}(${argsString}) {
+  // Empty mock - original function mocked out
+}`;
+        }
         return;
+    }
+    // Check if the schema only has function type markers without actual return data.
+    // Only apply this when we have a signature (i.e., it's explicitly a function)
+    // and it's not a class. This handles validation functions that should be
+    // empty mocks rather than returning meaningless "function" strings.
+    if (foundEntityWithSignature &&
+        signatureSchema &&
+        entityType !== 'class' &&
+        !hasMeaningfulReturnData(relevantReturnValueSchema)) {
+        const argCount = Object.keys(signatureSchema).filter((key) => key.startsWith('signature[')).length;
+        const args = Array.from({ length: argCount }, (_, i) => `arg${i + 1}`);
+        const argsString = args.join(', ');
+        // Generate empty mock function
+        return `function ${mockName}(${argsString}) {
+  // Empty mock - original function mocked out
+}`;
+    }
     const dataStructureName = dataStructurePath === 'returnValue'
         ? mockName
         : joinParenthesesAndArrays(splitOutsideParenthesesAndArrays(dataStructurePath).slice(0, mockNameParts.length));
@@ -165,14 +280,24 @@ export default function constructMockCode(mockName, dependencySchemas) {
         if (root) {
             basePaths[0] = 'scenarios().data()';
         }
-        const optionalAccess = (name, stripArguments = false) => {
+        const optionalAccess = (name, stripArguments = false, isRootAccess = false) => {
             if (stripArguments && name.includes('(')) {
                 name = name.split('(')[0] + '()';
             }
             // Strip type parameters like <typeof loader> from function names
             // so "useLoaderData<typeof loader>()" becomes "useLoaderData()"
             name = cleanOutTypes(name);
-            if (name.match(/\[\d+\]/)) {
+            // For variable-qualified mocks at root level, use the "variableName <- functionName" format
+            // This matches the format generated by gatherDataForMocks and is more intuitive for LLMs
+            // e.g., "useFetcher()" becomes "entityDiffFetcher <- useFetcher"
+            if (isRootAccess && variableQualifier) {
+                // Remove trailing () from the function name to match gatherDataForMocks key format
+                const baseName = name.replace(/\(\)$/, '');
+                name = `${variableQualifier} <- ${baseName}`;
+            }
+            // Only use unquoted array access syntax for pure array indices like [0], [1]
+            // NOT for variable-qualified names like "fetcher[1] <- useFetcher" which need quotes
+            if (name.match(/^\[\d+\]$/) && !name.includes(' <- ')) {
                 return `?.${name}`;
             }
             return `?.["${name.replace(/\n/g, '\\n')}"]`;
@@ -187,7 +312,13 @@ export default function constructMockCode(mockName, dependencySchemas) {
                 return [];
             }
             const addReturnValueFunctionAccessor = (dataPath) => {
-                if (returnValue.returnsFunctionArgs) {
+                // Add function call accessor if:
+                // - There are actual arguments, OR
+                // - This is a callable (not a method that returns an array directly)
+                // For methods like getAll() that return arrays, the data is at ["getAll()"] not ["getAll()"]["()"]
+                if (returnValue.returnsFunctionArgs &&
+                    (returnValue.returnsFunctionArgs.length > 0 ||
+                        !returnValue.returnsFunctionArray)) {
                     if (returnValue.isArray) {
                         dataPath = `${dataPath}${optionalAccess('[0]')}`;
                     }
@@ -198,7 +329,8 @@ export default function constructMockCode(mockName, dependencySchemas) {
             return basePaths
                 .flatMap((basePath) => {
                 if (!returnValue.args) {
-                    let dataPath = `${basePath}${optionalAccess(safeString(returnValue.name))}`;
+                    // Pass root as both stripArguments and isRootAccess for non-function paths
+                    let dataPath = `${basePath}${optionalAccess(safeString(returnValue.name), false, root)}`;
                     dataPath = addReturnValueFunctionAccessor(dataPath);
                     return dataPath;
                 }
@@ -206,7 +338,8 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     const signature = returnValue.name.match(/\[\d*\]$/)
                         ? returnValue.name
                         : `${returnValue.name}(${argSet.join(', ')})`;
-                    let dataPath = `${basePath}${optionalAccess(safeString(signature), root)}`;
+                    // Pass root as both stripArguments and isRootAccess
+                    let dataPath = `${basePath}${optionalAccess(safeString(signature), root, root)}`;
                     dataPath = addReturnValueFunctionAccessor(dataPath);
                     return dataPath;
                 });
@@ -214,14 +347,57 @@ export default function constructMockCode(mockName, dependencySchemas) {
                 .filter(Boolean);
         };
         const constructContent = (dataPaths) => {
-            const { name, args, nested, isArray, isGenericArray, returnsFunctionArgs, returnsFunctionArray, isAsyncFunction, } = returnValue;
+            const { name, args, nested, isArray, isGenericArray, returnsFunctionArgs, returnsFunctionArray, isAsyncFunction, hasNoReturnData, } = returnValue;
             const nestedContent = (nested ?? []).map((nestedItem) => {
                 const nestedContent = constructReturnValueString(nestedItem, dataPaths);
                 return nestedContent;
             });
+            // Array prototype methods that should be ignored when building mocks
+            // (these work on any array - we don't need to mock them)
+            const ARRAY_PROTOTYPE_METHODS = new Set([
+                'at',
+                'concat',
+                'copyWithin',
+                'entries',
+                'every',
+                'fill',
+                'filter',
+                'find',
+                'findIndex',
+                'findLast',
+                'findLastIndex',
+                'flat',
+                'flatMap',
+                'forEach',
+                'includes',
+                'indexOf',
+                'join',
+                'keys',
+                'lastIndexOf',
+                'map',
+                'pop',
+                'push',
+                'reduce',
+                'reduceRight',
+                'reverse',
+                'shift',
+                'slice',
+                'some',
+                'sort',
+                'splice',
+                'toLocaleString',
+                'toReversed',
+                'toSorted',
+                'toSpliced',
+                'toString',
+                'unshift',
+                'values',
+                'with',
+                'length',
+            ]);
             const levelContentItems = [];
             // Add spread for data paths when:
-            // - Not a function returning an array
+            // - Not a function returning an array, OR function returns array with custom methods
             // - Has function args OR not an array
             // - Not structural, OR is a structural [0] element with nested content
             //   (array elements with nested functions need their other properties spread)
@@ -229,7 +405,19 @@ export default function constructMockCode(mockName, dependencySchemas) {
                 returnValue.name === '[0]' &&
                 nested &&
                 nested.length > 0;
-            if (!returnsFunctionArray &&
+            // Check if function returns array but has custom methods (not just array prototype methods)
+            // Excludes array indices like [0], [1], [] which are not methods
+            const hasCustomMethodsOnArrayReturn = returnsFunctionArray &&
+                nested &&
+                nested.length > 0 &&
+                nested.some((n) => {
+                    // Skip array indices - they're not methods
+                    if (n.name.match(/^\[\d*\]$/))
+                        return false;
+                    const methodName = n.name.replace(/[<(].*$/, '');
+                    return !ARRAY_PROTOTYPE_METHODS.has(methodName);
+                });
+            if ((!returnsFunctionArray || hasCustomMethodsOnArrayReturn) &&
                 (returnsFunctionArgs || !isArray) &&
                 (!returnValue.isStructural || isStructuralArrayElementWithNested)) {
                 levelContentItems.push(...dataPaths.map((path) => `...${path}`));
@@ -242,7 +430,16 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     .join(', ');
                 let funcContents = '';
                 if (returnsFunctionArray) {
-                    if (levelContents.length === 0) {
+                    if (hasNoReturnData) {
+                        // Function has no return data (only signatures) - return empty array
+                        funcContents = 'return []';
+                    }
+                    else if (levelContents.length === 0 && dataPaths.length === 1) {
+                        // When returning an array with no nested content, return the data path directly
+                        // (the data path points to the array in scenario data)
+                        funcContents = `return ${dataPaths[0]}`;
+                    }
+                    else if (levelContents.length === 0) {
                         funcContents = 'return []';
                     }
                     else {
@@ -250,8 +447,19 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     }
                 }
                 else {
+                    // Check if function has no actual return data (only signatures)
+                    const hasNestedItems = nested && nested.length > 0;
+                    const hasActualNestedContent = nestedContent.filter(Boolean).length > 0;
                     if (levelContentItems.length === 1 && dataPaths.length === 1) {
-                        funcContents = `return ${dataPaths[0]}`;
+                        if (hasNoReturnData ||
+                            (hasNestedItems && !hasActualNestedContent)) {
+                            // Function has no return data (only signatures) - return empty array
+                            funcContents = 'return []';
+                        }
+                        else {
+                            // Has return data - return data path
+                            funcContents = `return ${dataPaths[0]}`;
+                        }
                     }
                     else {
                         funcContents = `return {\n${indent(levelContents)}\n}`;
@@ -262,12 +470,30 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     return levelContents;
                 }
             }
+            // Check if all nested items are array prototype methods
+            const hasOnlyArrayPrototypeMethods = returnsFunctionArray &&
+                nested &&
+                nested.length > 0 &&
+                nested.every((n) => {
+                    // Extract method name before type params and args (e.g., "reduce<...>(...)" -> "reduce")
+                    const methodName = n.name.replace(/[<(].*$/, '');
+                    return ARRAY_PROTOTYPE_METHODS.has(methodName);
+                });
             let returnValueContents = '';
             if (!returnsFunctionArgs &&
                 nestedContent.length === 0 &&
                 dataPaths.length === 1) {
                 returnValueContents = dataPaths[0];
             }
+            else if (returnsFunctionArray &&
+                dataPaths.length === 1 &&
+                !root &&
+                (!nested || nested.length === 0 || hasOnlyArrayPrototypeMethods)) {
+                // Nested method returns an array - return the data path directly
+                // Only if nested content is empty or only contains array prototype methods
+                // (like .reduce() which work on any array - we don't need to mock them)
+                returnValueContents = dataPaths[0];
+            }
             else {
                 if (isArray) {
                     // When GENERIC array (using []) has nested content (like functions that need wrapping),
@@ -427,7 +653,9 @@ export default function constructMockCode(mockName, dependencySchemas) {
                 parts.splice(i, 1);
             }
         }
-        if (parts[0].split('(')[0] !== mockName)
+        // For variable-qualified mocks like "revalidator <- useRevalidator",
+        // compare against baseMockName ("useRevalidator"), not the full mockName
+        if (parts[0].split('(')[0] !== baseMockName)
             continue;
         // Include paths with functionCallReturnValue OR function-typed paths that need mocking
         const hasFunctionCallReturnValue = parts.some((p) => isFunctionCallReturnValue(p));
@@ -438,10 +666,22 @@ export default function constructMockCode(mockName, dependencySchemas) {
         // Exclude paths with functionCallReturnValue - those are handled by existing logic
         // Also check if there's a functionCallReturnValue path that extends this path
         const hasSignaturePath = parts.some((p) => p.startsWith('signature['));
+        // Skip signature paths entirely - they describe INPUT structure (function arguments),
+        // not OUTPUT (return values). This applies even if they contain functionCallReturnValue
+        // 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}`));
+        // 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
+        const isJsxComponent = value === 'jsx-component' ||
+            relevantReturnValueSchema[key + '.functionCallReturnValue'] ===
+                'jsx-component';
+        if (isJsxComponent)
+            continue;
         const isFunctionPath = ['function', 'async-function'].includes(value) &&
             parts[parts.length - 1].endsWith(')') &&
-            !hasSignaturePath &&
             !hasAnyFunctionCallReturnValue &&
             !hasExtendedFunctionCallReturnValue;
         if (!hasFunctionCallReturnValue && !isFunctionPath)
@@ -523,6 +763,23 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     if (returnValueSection.returnsFunctionArgs && value === 'array') {
                         returnValueSection.returnsFunctionArray = true;
                     }
+                    // Check if this function call has actual return data in the schema
+                    // If the path only has .signature entries (no .functionCallReturnValue), mark as no return data
+                    // Use parts[i] (original part) instead of modified 'part' to include array indices like [1]
+                    const functionCallPath = joinParenthesesAndArrays([
+                        ...parts.slice(0, i),
+                        parts[i], // Use original part which may include [1](args) together
+                    ]);
+                    // Check if there's a .functionCallReturnValue AFTER this function call path
+                    const hasFunctionCallReturnValue = Object.keys(relevantReturnValueSchema).some((k) => {
+                        if (!k.startsWith(functionCallPath))
+                            return false;
+                        const afterPath = k.slice(functionCallPath.length);
+                        return afterPath.startsWith('.functionCallReturnValue');
+                    });
+                    if (!hasFunctionCallReturnValue) {
+                        returnValueSection.hasNoReturnData = true;
+                    }
                 }
                 if (nextIsDifferentiatedArray || hasNestedFunction) {
                     continue;
@@ -576,6 +833,11 @@ export default function constructMockCode(mockName, dependencySchemas) {
                         if (['function', 'async-function'].includes(nextValue)) {
                             continue; // Continue to process the function property
                         }
+                        // If the next part is an object with nested content, continue processing
+                        // This handles paths like functionCallReturnValue.selectedOptions.elementOptions[]
+                        if (nextValue === 'object' || nextValue === 'array') {
+                            continue;
+                        }
                     }
                     break;
                 }
@@ -583,8 +845,29 @@ export default function constructMockCode(mockName, dependencySchemas) {
             else if (isKeyWord) {
                 continue;
             }
+            // Handle parts that end with [] (e.g., "elementOptions[]")
+            // These should be treated as the base part + array accessor
+            const endsWithArrayAccessor = part.endsWith('[]') && part !== '[]';
+            const basePart = endsWithArrayAccessor ? part.slice(0, -2) : part;
             const withoutArgs = (p) => p.name.split('(')[0];
-            let relevantPart = returnValueSection.nested?.find((p) => withoutArgs(p) === withoutArgs({ name: part }));
+            let relevantPart = returnValueSection.nested?.find((p) => withoutArgs(p) === withoutArgs({ name: basePart }));
+            // If part ends with [], update the existing part to be an array
+            if (endsWithArrayAccessor && relevantPart) {
+                relevantPart.isArray = true;
+                relevantPart.isGenericArray = true;
+                // Create structural [0] element for nested content
+                relevantPart.nested || (relevantPart.nested = []);
+                if (!relevantPart.nested.find((p) => p.name === '[0]')) {
+                    relevantPart.nested.push({
+                        name: '[0]',
+                        nested: [],
+                        isStructural: true,
+                    });
+                }
+                // Move to the [0] element for processing subsequent parts
+                returnValueSection = relevantPart.nested.find((p) => p.name === '[0]');
+                continue;
+            }
             if (!relevantPart) {
                 const hasArguments = part.includes('(') && part.endsWith(')');
                 if (part === '[]') {
@@ -605,6 +888,29 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     }
                     continue;
                 }
+                // Handle parts that end with [] (e.g., "elementOptions[]") when base part doesn't exist
+                if (endsWithArrayAccessor) {
+                    // Create the base part with array properties
+                    relevantPart = {
+                        name: basePart,
+                        isArray: true,
+                        isGenericArray: true,
+                        nested: [],
+                    };
+                    returnValueSection.nested || (returnValueSection.nested = []);
+                    returnValueSection.nested.push(relevantPart);
+                    // Create structural [0] element for nested content
+                    if (partValue === 'object' && !isLastPart) {
+                        relevantPart.nested.push({
+                            name: '[0]',
+                            nested: [],
+                            isStructural: true,
+                        });
+                        // Move to the [0] element for processing subsequent parts
+                        returnValueSection = relevantPart.nested.find((p) => p.name === '[0]');
+                    }
+                    continue;
+                }
                 if (isLastPart &&
                     ['function', 'async-function'].includes(value) &&
                     !hasArguments) {
@@ -640,14 +946,39 @@ export default function constructMockCode(mockName, dependencySchemas) {
                     returnValueSection.nested.push(relevantPart);
                 }
             }
-            else if (funcArgs(part).join(',') !== funcArgs(relevantPart.name).join(',')) {
+            else if (relevantPart.args &&
+                funcArgs(part).join(',') !== funcArgs(relevantPart.name).join(',')) {
                 const existingArgs = relevantPart.args.find((args) => args.join(',') === funcArgs(part).join(','));
                 if (!existingArgs) {
                     relevantPart.args || (relevantPart.args = []);
                     relevantPart.args.push(funcArgs(part));
                 }
             }
+            // If nextPart is [], update existing part to be a generic array
+            if (relevantPart && nextPart === '[]') {
+                relevantPart.isArray = true;
+                relevantPart.isGenericArray = true;
+            }
             if (!hasNestedFunction) {
+                // Before breaking, check if this function returns an array
+                // by looking for a functionCallReturnValue: 'array' entry in the schema
+                if (relevantPart && part.endsWith(')')) {
+                    const functionReturnPath = joinParenthesesAndArrays([
+                        ...parts.slice(0, i + 1),
+                        RETURN_VALUE,
+                    ]);
+                    const returnType = relevantReturnValueSchema[functionReturnPath];
+                    if (returnType === 'array') {
+                        // Only set returnsFunctionArray if the array contains plain data,
+                        // NOT if it contains functions (which need wrapping)
+                        const arrayElementPath = `${functionReturnPath}[]`;
+                        const arrayElementType = relevantReturnValueSchema[arrayElementPath];
+                        // If array element is a function, we need to wrap it - don't set returnsFunctionArray
+                        if (arrayElementType !== 'function') {
+                            relevantPart.returnsFunctionArray = true;
+                        }
+                    }
+                }
                 break;
             }
             returnValueSection = relevantPart;
@@ -680,10 +1011,75 @@ export default function constructMockCode(mockName, dependencySchemas) {
         return `const ${mockNameParts[0]} = {\n${indent(returnParts.join(',\n'))}\n};`;
     }
     else if (isFunction) {
-        return `const ${mockName}ReturnValue = ${contents};\n\n${isRootAsyncFunction ? 'async ' : ''}function ${mockName}() {\n${indent(`return ${mockName}ReturnValue;`)}\n}`;
+        // For headers() and cookies() from next/headers, add common iterator methods
+        // These are needed when the mock is passed to functions that use .entries(), .keys(), etc.
+        // (e.g., Object.fromEntries(headers.entries()) in buildLegacyHeaders)
+        const needsIteratorMethods = mockName === 'headers' || mockName === 'cookies';
+        let enhancedContents = contents;
+        if (needsIteratorMethods && contents.trim().startsWith('{')) {
+            // Add iterator methods that operate on the scenario data
+            const iteratorMethods = `,
+  entries: () => Object.entries(scenarios().data()?.["${mockName}()"] || {}),
+  keys: () => Object.keys(scenarios().data()?.["${mockName}()"] || {}),
+  values: () => Object.values(scenarios().data()?.["${mockName}()"] || {}),
+  forEach: (fn) => Object.entries(scenarios().data()?.["${mockName}()"] || {}).forEach(([k, v]) => fn(v, k)),
+  has: (key) => Object.prototype.hasOwnProperty.call(scenarios().data()?.["${mockName}()"] || {}, key)`;
+            // Insert before the closing brace (handle trailing whitespace)
+            enhancedContents = contents.replace(/\}\s*$/, iteratorMethods + '\n}');
+        }
+        // Classes need to be mocked as classes, not functions.
+        // They're used with `new ClassName(args)` which requires a proper constructor.
+        // A function mock like `function ClassName() { return value; }` doesn't work because
+        // `new ClassName("arg")` wouldn't create the expected instance.
+        // For Error subclasses (detected by name ending in "Error"), extend Error for proper error handling.
+        if (entityType === 'class') {
+            const isErrorSubclass = mockName.endsWith('Error');
+            const baseClass = isErrorSubclass ? 'Error' : 'Object';
+            const superCall = isErrorSubclass ? 'super(message);' : '';
+            const nameAssignment = isErrorSubclass
+                ? `this.name = '${mockName}';`
+                : '';
+            return `class ${mockName}${isErrorSubclass ? ' extends Error' : ''} {
+  constructor(message) {
+    ${superCall}
+    ${nameAssignment}
+    Object.assign(this, scenarios().data()?.["${mockName}()"] || {});
+  }
+}`;
+        }
+        // For variable-qualified mocks, use underscore separator in the function name
+        // e.g., "entityDiffFetcher <- useFetcher" becomes "useFetcher_entityDiffFetcher"
+        // For indexed qualifiers like "data[1]", convert to safe identifier: "useFetcher_data_1"
+        const safeFunctionName = variableQualifier
+            ? `${baseMockName}_${safeVariableQualifier(variableQualifier)}`
+            : mockName;
+        return `const ${safeFunctionName}ReturnValue = ${enhancedContents};\n\n${isRootAsyncFunction ? 'async ' : ''}function ${safeFunctionName}() {\n${indent(`return ${safeFunctionName}ReturnValue;`)}\n}`;
     }
     else {
-        return `const ${mockName} = ${contents}`;
+        // For variable-qualified mocks, use underscore separator in the const name
+        // For indexed qualifiers like "data[1]", convert to safe identifier: "useFetcher_data_1"
+        const safeName = variableQualifier
+            ? `${baseMockName}_${safeVariableQualifier(variableQualifier)}`
+            : mockName;
+        // Get any jsx-component properties that need to be preserved from the original
+        const jsxProperties = getJsxComponentProperties(mockName, relevantReturnValueSchema);
+        // If there are jsx-component properties, add them as references to the original
+        if (jsxProperties.length > 0) {
+            const jsxPreservations = jsxProperties
+                .map((prop) => `${prop}: ${mockName}__cyOriginal.${prop}`)
+                .join(',\n');
+            if (contents.trim().startsWith('{')) {
+                // Contents is already an object - insert jsx preservations before the closing brace
+                const enhancedContents = contents.replace(/\}\s*$/, `,\n${indent(jsxPreservations)}\n}`);
+                return `const ${safeName} = ${enhancedContents}`;
+            }
+            else {
+                // Contents is a simple accessor (e.g., scenarios().data()?.["MockName"])
+                // Wrap it in an object with the jsx preservations
+                return `const ${safeName} = {\n${indent(`...${contents}`)},\n${indent(jsxPreservations)}\n}`;
+            }
+        }
+        return `const ${safeName} = ${contents}`;
     }
 }
 //# sourceMappingURL=constructMockCode.js.map