@codeyam/codeyam-cli 0.1.0-staging.28f73cf → 0.1.0-staging.76566f9

package/analyzer-template/project/constructMockCode.ts

@@ -4,8 +4,9 @@ import {
   functionArguments,
   cleanOutBoundary,
   fillInDirectSchemaGapsAndUnknowns,
+  removeDuplicateFunctionCalls,
 } from '~codeyam/ai';
-import { DataStructure, DeepReadonly } from '~codeyam/types';
+import { DataStructure, DeepReadonly, EntityType } from '~codeyam/types';
 
 interface ReturnValuePart {
   name: string;
@@ -23,6 +24,136 @@ interface ReturnValuePart {
   argVariants?: Map<string, ReturnValuePart[]>;
 }
 
+/**
+ * Converts a call signature to a valid JavaScript identifier (function name).
+ * The original signature is preserved for data access - this only creates the function name.
+ *
+ * Examples:
+ * - "useAuth()" → "useAuth"
+ * - "db.select(usersQuery)" → "db_select_usersQuery"
+ * - "db.select(postsQuery)" → "db_select_postsQuery"
+ * - "useFetcher<User>()" → "useFetcher_User"
+ * - "useFetcher<{ data: UserData | null }>()" → "useFetcher_data_UserData_null"
+ * - "eq('user_id', value)" → "eq_user_id_value"
+ * - "from('workouts')" → "from_workouts"
+ */
+function callSignatureToFunctionName(signature: string): string {
+  // Extract components from the signature
+  const components: string[] = [];
+
+  // 1. Extract function path (parts separated by dots outside parens/brackets)
+  const pathMatch = signature.match(/^([^<(]+)/);
+  if (pathMatch) {
+    const path = pathMatch[1];
+    // Split on dots but preserve the parts
+    components.push(...path.split('.').filter(Boolean));
+  }
+
+  // 2. Extract generic type parameters (content between < and >)
+  const genericMatch = signature.match(/<([^>]+)>/);
+  if (genericMatch) {
+    const genericContent = genericMatch[1];
+    // Extract meaningful identifiers from generic type
+    // Handle complex types like "{ data: UserData | null }"
+    const typeIdentifiers = genericContent
+      .replace(/[{}:;,]/g, ' ') // Remove structural chars
+      .replace(/\|/g, ' ') // Handle union types
+      .split(/\s+/)
+      .filter(Boolean)
+      .filter((s) => /^[A-Za-z_][A-Za-z0-9_]*$/.test(s)) // Only valid identifiers
+      .filter(
+        (s) =>
+          ![
+            'null',
+            'undefined',
+            'void',
+            'never',
+            'any',
+            'unknown',
+            'data',
+            'typeof',
+          ].includes(s),
+      ); // Skip common non-meaningful keywords
+
+    if (typeIdentifiers.length > 0) {
+      components.push(...typeIdentifiers.slice(0, 2)); // Limit to first 2 for reasonable length
+    }
+  }
+
+  // 3. Extract function arguments (first 2 for disambiguation)
+  const argsMatch = signature.match(/\(([^)]*)\)/);
+  if (argsMatch && argsMatch[1]) {
+    const argsContent = argsMatch[1].trim();
+    if (argsContent) {
+      const args = argsContent.split(',').map((arg) => arg.trim());
+      for (const arg of args.slice(0, 2)) {
+        // For quoted strings, extract the content
+        const stringMatch = arg.match(/^['"`](.+)['"`]$/);
+        if (stringMatch) {
+          // Split on dots for string paths like 'users.id'
+          const parts = stringMatch[1].split('.').filter(Boolean);
+          components.push(...parts);
+        } else if (/^[A-Za-z_][A-Za-z0-9_]*$/.test(arg)) {
+          // Valid identifier - use as-is
+          components.push(arg);
+        } else if (/^\d+$/.test(arg)) {
+          // Number - use as-is
+          components.push(arg);
+        }
+        // Skip complex expressions
+      }
+    }
+  }
+
+  // Build the function name from components
+  const functionName = components
+    .join('_')
+    .replace(/[^a-zA-Z0-9_]/g, '_') // Sanitize special chars
+    .replace(/_+/g, '_') // Collapse multiple underscores
+    .replace(/^_|_$/g, ''); // Trim underscores
+
+  return functionName || 'mock';
+}
+
+/**
+ * Check if a mock name is a call signature (e.g., "useFetcher<User>()", "db.select(query)")
+ */
+function isCallSignature(mockName: string): boolean {
+  // Call signatures contain parentheses (function calls)
+  return mockName.includes('(');
+}
+
+/**
+ * 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: string,
+  schema: Record<string, string>,
+): string[] {
+  const jsxProperties: string[] = [];
+
+  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.
@@ -163,16 +294,55 @@ function isValidKey(key: string) {
 export default function constructMockCode(
   mockName: string,
   dependencySchemas: DeepReadonly<DataStructure['dependencySchemas']>,
+  entityType?: EntityType,
+  _canonicalKey?: string, // DEPRECATED: No longer used, kept for API compatibility
+  options?: {
+    keepOriginalFunctionName?: boolean;
+    uniqueFunctionSuffix?: string;
+  },
 ) {
-  const mockNameParts = splitOutsideParenthesesAndArrays(mockName);
+  // Check if mockName is a call signature (e.g., "useFetcher<User>()", "db.select(query)")
+  const mockNameIsCallSignature = isCallSignature(mockName);
+
+  // For call signatures, use the original signature for data access but generate
+  // a valid JS function name from it
+  const derivedFunctionName = mockNameIsCallSignature
+    ? callSignatureToFunctionName(mockName)
+    : null;
+
+  // The baseMockName is the function name without type params and args
+  // e.g., "useFetcher<User>()" -> "useFetcher", "db.select(query)" -> "db"
+  const baseMockName = mockName.split(/[<(]/)[0];
+
+  // The data key is the mockName (call signature) for data access
+  let dataKey: string;
+
+  const mockNameParts = splitOutsideParenthesesAndArrays(baseMockName);
 
   let relevantReturnValueSchema: DataStructure['returnValueSchema'] | undefined;
   let dataStructurePath: string | undefined;
   let dataStructureValue: string | undefined;
+  let foundEntityWithSignature = false;
+  let signatureSchema: DataStructure['signatureSchema'] | undefined;
 
   for (const filePath in dependencySchemas) {
     for (const entityName in dependencySchemas[filePath]) {
-      if (entityName !== mockNameParts[0]) continue;
+      // Match entity by base name (without generics/args)
+      const entityBaseName = entityName.split(/[<(]/)[0];
+      const matches =
+        entityBaseName === baseMockName || entityName === mockNameParts[0];
+
+      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 ?? {},
@@ -194,16 +364,114 @@ export default function constructMockCode(
         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) return;
+  // Check if the entity is used as a function (called with ()) vs an object/namespace.
+  // Look for paths in the schema that start with "baseMockName(" or "baseMockName<" indicating function calls.
+  // The "<" handles generic type parameters like useLoaderData<T>().
+  // Also check dataStructurePath === 'returnValue' which indicates a function return value.
+  const entityIsFunction =
+    foundEntityWithSignature ||
+    dataStructurePath === 'returnValue' ||
+    Object.keys(relevantReturnValueSchema ?? {}).some(
+      (key) =>
+        key.startsWith(`${baseMockName}(`) ||
+        key.startsWith(`${baseMockName}<`),
+    );
+
+  // Calculate the data key - use the call signature (mockName) for data access
+  // For simple names without parentheses:
+  // - Append () ONLY if the entity is a function/hook (detected above)
+  // - Don't append () for object/namespace mocks like "supabase"
+  if (mockNameIsCallSignature || mockName.includes('(')) {
+    dataKey = mockName;
+  } else if (entityIsFunction) {
+    // Entity is a function/hook - append () to match call signature format
+    dataKey = `${mockName}()`;
+  } else {
+    // Entity is an object/namespace - use bare name as key
+    dataKey = mockName;
+  }
+
+  // LOGGING: Track data key generation for debugging key mismatch issues
+  console.log(
+    'CODEYAM DEBUG [constructMockCode]: Generated dataKey:',
+    JSON.stringify({
+      mockName,
+      dataKey,
+      mockNameIsCallSignature,
+      entityIsFunction,
+      baseMockName,
+    }),
+  );
+
+  // 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: DataStructure['returnValueSchema'] | undefined,
+  ): boolean => {
+    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'
@@ -246,7 +514,11 @@ export default function constructMockCode(
       basePaths[0] = 'scenarios().data()';
     }
 
-    const optionalAccess = (name: string, stripArguments = false) => {
+    const optionalAccess = (
+      name: string,
+      stripArguments = false,
+      isRootAccess = false,
+    ) => {
       if (stripArguments && name.includes('(')) {
         name = name.split('(')[0] + '()';
       }
@@ -255,7 +527,17 @@ export default function constructMockCode(
       // so "useLoaderData<typeof loader>()" becomes "useLoaderData()"
       name = cleanOutTypes(name);
 
-      if (name.match(/\[\d+\]/)) {
+      // For root data access, use the dataKey (original call signature or canonical key)
+      // This preserves the original call signature for LLM clarity
+      if (isRootAccess) {
+        // For call signature format, use the original mockName as the data key
+        // e.g., scenarios().data()?.["useFetcher<User>()"]
+        // e.g., scenarios().data()?.["db.select(usersQuery)"]
+        return `?.["${dataKey.replace(/\n/g, '\\n')}"]`;
+      }
+
+      // Only use unquoted array access syntax for pure array indices like [0], [1]
+      if (name.match(/^\[\d+\]$/)) {
         return `?.${name}`;
       }
 
@@ -294,7 +576,8 @@ export default function constructMockCode(
       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;
           }
@@ -303,7 +586,8 @@ export default function constructMockCode(
               ? 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;
@@ -492,10 +776,19 @@ export default function constructMockCode(
           // When GENERIC array (using []) has nested content (like functions that need wrapping),
           // use .map() to transform ALL elements instead of just creating [0]
           // For DIFFERENTIATED arrays (using [0], [1], etc.), keep the static array structure
+          //
+          // IMPORTANT: If the nested content contains differentiated indices like [0], [1],
+          // we MUST use static array pattern, not .map(). The presence of differentiated
+          // indices means the array elements have different types/structures, so .map()
+          // would generate invalid code trying to treat them uniformly.
+          const hasDifferentiatedIndices =
+            nested &&
+            nested.some((n) => n.name.match(/^\[\d+\]$/) && n.name !== '[0]');
           if (
             isGenericArray &&
             nestedContent.length > 0 &&
-            dataPaths.length > 0
+            dataPaths.length > 0 &&
+            !hasDifferentiatedIndices
           ) {
             // Get the array base path (without the [0])
             const arrayBasePath = dataPaths[0].replace(/\?\.\[0\]$/, '');
@@ -528,6 +821,13 @@ export default function constructMockCode(
       if (args && args.length > 0) {
         if (!isValidKey(name)) return;
 
+        // Skip array index patterns like [], [0], [1] when they have args
+        // These represent function calls on array elements, not property keys
+        // e.g., customSizes[].(args) means each array element is callable, not a property named "[]"
+        if (name.match(/^\[\d*\]$/)) {
+          return;
+        }
+
         const mostArgs = args.sort(
           (a: string[], b: string[]) => b.length - a.length,
         )[0];
@@ -587,8 +887,18 @@ export default function constructMockCode(
             // Use all paths for fallback (existing behavior)
             fallbackContent = `return ${returnValueContents}`;
           } else {
-            // Use all paths as fallback for unknown arguments
-            fallbackContent = `return ${returnValueContents}`;
+            // No explicit fallback paths - return the first literal's value as default
+            // Returning spread of all values is dangerous because if values are primitives (strings),
+            // spreading them creates objects with numeric keys like {0:'a', 1:'b', ...}
+            // which causes "Objects are not valid as React child" errors
+            const firstLiteralValue = literalKeys[0];
+            const firstGroupPaths = argGroups.get(firstLiteralValue);
+            if (firstGroupPaths && firstGroupPaths.length === 1) {
+              fallbackContent = `return ${firstGroupPaths[0]}`;
+            } else {
+              // Multiple paths for first literal - return undefined as safe fallback
+              fallbackContent = `return undefined`;
+            }
           }
 
           const funcContents =
@@ -682,7 +992,9 @@ export default function constructMockCode(
       }
     }
 
-    if (parts[0].split('(')[0] !== mockName) continue;
+    // Compare against baseMockName (without generics/args), not the full mockName
+    // e.g., for "useFetcher<User>()", baseMockName is "useFetcher"
+    if (parts[0].split('(')[0] !== baseMockName) continue;
 
     // Include paths with functionCallReturnValue OR function-typed paths that need mocking
     const hasFunctionCallReturnValue = parts.some((p) =>
@@ -697,13 +1009,61 @@ export default function constructMockCode(
     // 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;
+
+    // Skip paths that bypass .functionCallReturnValue when there's a corresponding path with it.
+    // Example: If we have both:
+    //   - trpc.customer.useQuery(...).data (incorrect - no .functionCallReturnValue)
+    //   - trpc.customer.useQuery(...).functionCallReturnValue.data (correct)
+    // We should skip the first path because the second one properly captures the return value.
+    // This can happen when the analyzer sees both the raw property access and the return value structure.
+    if (!hasFunctionCallReturnValue) {
+      // This path has no .functionCallReturnValue. Check if any function call in this path
+      // has a corresponding .functionCallReturnValue path in the schema.
+      let shouldSkipKey = false;
+      for (let i = 0; i < parts.length - 1; i++) {
+        const part = parts[i];
+        if (part.endsWith(')') && !isFunctionCallReturnValue(parts[i + 1])) {
+          // 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'),
+          );
+          if (hasProperFunctionCallPath) {
+            // Skip this path - the .functionCallReturnValue path will handle it correctly
+            shouldSkipKey = true;
+            break;
+          }
+        }
+      }
+      if (shouldSkipKey) continue;
+    }
+
     const isFunctionPath =
       ['function', 'async-function'].includes(value) &&
       parts[parts.length - 1].endsWith(')') &&
-      !hasSignaturePath &&
       !hasAnyFunctionCallReturnValue &&
       !hasExtendedFunctionCallReturnValue;
     if (!hasFunctionCallReturnValue && !isFunctionPath) continue;
@@ -879,6 +1239,11 @@ export default function constructMockCode(
             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;
         }
@@ -886,11 +1251,35 @@ export default function constructMockCode(
         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: ReturnValuePart) => p.name.split('(')[0];
       let relevantPart = returnValueSection.nested?.find(
         (p: ReturnValuePart) =>
-          withoutArgs(p) === withoutArgs({ name: part } as ReturnValuePart),
+          withoutArgs(p) === withoutArgs({ name: basePart } as ReturnValuePart),
       );
+
+      // 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 ||= [];
+        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 === '[]') {
@@ -914,6 +1303,33 @@ export default function constructMockCode(
           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.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) &&
@@ -954,19 +1370,31 @@ export default function constructMockCode(
         ) {
           returnValueSection.nested.push(relevantPart);
         }
-      } else if (
-        funcArgs(part).join(',') !== funcArgs(relevantPart.name).join(',')
-      ) {
-        const existingArgs = relevantPart.args.find(
-          (args) => args.join(',') === funcArgs(part).join(','),
-        );
+      } else {
+        // Add args to existing entry if current part has function arguments
+        // This handles the case where bare `t` is processed first (creating {name: 't', args: undefined})
+        // and then `t("common.close")` is processed - we need to add its args to the existing entry
+        const currentArgs = funcArgs(part);
+        const hasNewArgs = currentArgs.length > 0 || part.includes('(');
+
+        if (hasNewArgs) {
+          const existingArgs = relevantPart.args?.find(
+            (args) => args.join(',') === currentArgs.join(','),
+          );
 
-        if (!existingArgs) {
-          relevantPart.args ||= [];
-          relevantPart.args.push(funcArgs(part));
+          if (!existingArgs) {
+            relevantPart.args ||= [];
+            relevantPart.args.push(currentArgs);
+          }
         }
       }
 
+      // 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
@@ -1035,21 +1463,150 @@ export default function constructMockCode(
     // 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';
+      baseMockName === 'headers' || baseMockName === 'cookies';
     let enhancedContents = contents;
     if (needsIteratorMethods && contents.trim().startsWith('{')) {
       // Add iterator methods that operate on the scenario data
+      // Use the dataKey (original call signature or canonical key)
       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)`;
+  entries: () => Object.entries(scenarios().data()?.["${dataKey}"] || {}),
+  keys: () => Object.keys(scenarios().data()?.["${dataKey}"] || {}),
+  values: () => Object.values(scenarios().data()?.["${dataKey}"] || {}),
+  forEach: (fn) => Object.entries(scenarios().data()?.["${dataKey}"] || {}).forEach(([k, v]) => fn(v, k)),
+  has: (key) => Object.prototype.hasOwnProperty.call(scenarios().data()?.["${dataKey}"] || {}, key)`;
       // Insert before the closing brace (handle trailing whitespace)
       enhancedContents = contents.replace(/\}\s*$/, iteratorMethods + '\n}');
     }
-    return `const ${mockName}ReturnValue = ${enhancedContents};\n\n${isRootAsyncFunction ? 'async ' : ''}function ${mockName}() {\n${indent(`return ${mockName}ReturnValue;`)}\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 = baseMockName.endsWith('Error');
+      const superCall = isErrorSubclass ? 'super(message);' : '';
+      const nameAssignment = isErrorSubclass
+        ? `this.name = '${baseMockName}';`
+        : '';
+      // Use the safe function name for the class definition
+      const className = mockNameIsCallSignature
+        ? derivedFunctionName
+        : baseMockName;
+
+      return `class ${className}${isErrorSubclass ? ' extends Error' : ''} {
+  constructor(message) {
+    ${superCall}
+    ${nameAssignment}
+    Object.assign(this, scenarios().data()?.["${dataKey}"] || {});
+  }
+}`;
+    }
+
+    // Generate safe function name:
+    // 1. For call signatures: use derivedFunctionName
+    //    e.g., "useFetcher<User>()" becomes "useFetcher_User"
+    //    e.g., "db.select(usersQuery)" becomes "db_select_usersQuery"
+    // 2. With uniqueFunctionSuffix option: append suffix for unique naming
+    //    e.g., baseMockName = "useFetcher", suffix = "entityDiffFetcher" -> "useFetcher_entityDiffFetcher"
+    // 3. EXCEPTION: When keepOriginalFunctionName is true (for single-call cases), use the base name
+    let safeFunctionName: string;
+    if (options?.keepOriginalFunctionName) {
+      safeFunctionName = baseMockName;
+    } else if (options?.uniqueFunctionSuffix) {
+      safeFunctionName = `${baseMockName}_${options.uniqueFunctionSuffix}`;
+    } else if (mockNameIsCallSignature && derivedFunctionName) {
+      safeFunctionName = derivedFunctionName;
+    } else {
+      safeFunctionName = baseMockName;
+    }
+
+    // Check if this function returns a function (detected by double-call pattern: mockName(args)())
+    // This happens when the schema has keys like "wrapThrows(() => JSON.parse(savedFilters))()"
+    // where the function call is immediately followed by another call.
+    // Example usage: const result = wrapThrows(() => JSON.parse(x))(); // double call
+    const isHigherOrderFunction = Object.keys(
+      relevantReturnValueSchema ?? {},
+    ).some((key) => {
+      if (!key.startsWith(baseMockName)) return false;
+
+      // Find the first ( after baseMockName (the start of the function call)
+      const firstOpenParen = key.indexOf('(', baseMockName.length);
+      if (firstOpenParen === -1) return false;
+
+      // Skip if the ( is not immediately after the mock name
+      // (there might be type params like func<T>() - handle by checking for < or ()
+      const between = key.slice(baseMockName.length, firstOpenParen);
+      if (between.length > 0 && !between.startsWith('<')) return false;
+
+      // Find the matching ) for the first ( using depth counting
+      let depth = 1;
+      let i = firstOpenParen + 1;
+      while (i < key.length && depth > 0) {
+        if (key[i] === '(') depth++;
+        if (key[i] === ')') depth--;
+        i++;
+      }
+
+      if (depth !== 0) return false; // Unbalanced parentheses
+
+      // Now i points just after the matching )
+      // Check if there's another ( immediately (indicating double call)
+      const remaining = key.slice(i);
+      if (remaining.startsWith('(')) return true;
+
+      return false;
+    });
+
+    // For higher-order functions, wrap the return in an arrow function
+    // so that mockFunc(arg)() works correctly (outer call returns a function, inner call gets the data)
+    const returnValue = isHigherOrderFunction
+      ? `() => ${safeFunctionName}ReturnValue`
+      : `${safeFunctionName}ReturnValue`;
+
+    return `const ${safeFunctionName}ReturnValue = ${enhancedContents};\n\n${isRootAsyncFunction ? 'async ' : ''}function ${safeFunctionName}() {\n${indent(`return ${returnValue};`)}\n}`;
   } else {
-    return `const ${mockName} = ${contents}`;
+    // Generate safe const name:
+    // 1. For call signatures: use derivedFunctionName
+    // 2. With uniqueFunctionSuffix option: append suffix for unique naming
+    // 3. EXCEPTION: When keepOriginalFunctionName is true (for single-call cases), use the base name
+    let safeName: string;
+    if (options?.keepOriginalFunctionName) {
+      safeName = baseMockName;
+    } else if (options?.uniqueFunctionSuffix) {
+      safeName = `${baseMockName}_${options.uniqueFunctionSuffix}`;
+    } else if (mockNameIsCallSignature && derivedFunctionName) {
+      safeName = derivedFunctionName;
+    } else {
+      safeName = baseMockName;
+    }
+
+    // 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}`;
   }
 }