@codeyam/codeyam-cli 0.1.0-staging.d0ad4ae → 0.1.0-staging.d4f25c3

package/analyzer-template/project/constructMockCode.ts

@@ -294,6 +294,19 @@ function isValidKey(key: string) {
   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: string,
   dependencySchemas: DeepReadonly<DataStructure['dependencySchemas']>,
@@ -328,8 +341,9 @@ export default function constructMockCode(
   let dataStructureValue: string | undefined;
   let foundEntityWithSignature = false;
   let signatureSchema: DataStructure['signatureSchema'] | undefined;
+  let baseSchemaHasMethodCalls = false;
 
-  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];
@@ -365,15 +379,60 @@ export default function constructMockCode(
       });
 
       if (dataStructurePath) {
+        // Start with the base entity's return value schema
+        const baseReturnValueSchema =
+          dependencySchemas[filePath][entityName]?.returnValueSchema;
+        const mergedSchema = {
+          ...baseReturnValueSchema,
+        };
+
+        // Check if the base schema has method-call entries (e.g., .map().functionCallReturnValue)
+        // When it does, the scenario data is stored as an object with method keys, and
+        // array prototype methods need mock implementations. When it doesn't, the data
+        // is a raw array and native methods like .includes() work directly.
+        if (baseReturnValueSchema) {
+          baseSchemaHasMethodCalls = Object.keys(baseReturnValueSchema).some(
+            (k) =>
+              k.startsWith(baseMockName + '.') &&
+              k.includes('(') &&
+              k.includes('.functionCallReturnValue'),
+          );
+        }
+
+        // Merge in method-call dependencies that are separate entries.
+        // e.g., "activityTypes.find((a) => a.value === type)" is a separate dependency
+        // for a .find() call on activityTypes. We need to include these with a
+        // .functionCallReturnValue path so constructMockCode generates callable mock methods.
+        for (const otherEntityName in dependencySchemas[filePath]) {
+          if (otherEntityName === entityName) continue;
+          if (
+            otherEntityName.startsWith(baseMockName + '.') &&
+            otherEntityName.includes('(')
+          ) {
+            // Add a functionCallReturnValue entry for this method call.
+            // This ensures constructMockCode treats it as a function that returns data,
+            // generating a proper mock method with data lookup.
+            const fcrvPath = `${otherEntityName}.functionCallReturnValue`;
+            if (!mergedSchema[fcrvPath]) {
+              // Infer the return type from the method-call dependency's schema
+              const otherSchema =
+                dependencySchemas[filePath][otherEntityName]?.returnValueSchema;
+              // Look for element type (baseMockName[]) or fall back to 'unknown'
+              const elementType = otherSchema?.[`${baseMockName}[]`];
+              mergedSchema[fcrvPath] = elementType || 'unknown';
+            }
+          }
+        }
+
         relevantReturnValueSchema = fillInDirectSchemaGapsAndUnknowns({
-          schema: dependencySchemas[filePath][entityName]?.returnValueSchema,
+          schema: mergedSchema,
         });
         // NOTE: clearAttributesFromMapping is disabled because it deletes
         // method calls on arrays (like .eq() after functionCallReturnValue: 'array')
         // However, we still need to remove duplicate function calls that create invalid syntax
         removeDuplicateFunctionCalls(relevantReturnValueSchema);
         dataStructureValue = relevantReturnValueSchema?.[dataStructurePath];
-        break;
+        break entitySearch;
       }
     }
   }
@@ -449,7 +508,10 @@ export default function constructMockCode(
       const argsString = args.join(', ');
 
       // Generate empty mock function
-      return `function ${mockName}(${argsString}) {
+      // Use baseMockName (not mockName) because mockName may contain a full call
+      // signature with argument expressions (e.g., "logSignOutAction(sessionUser.id, ...)")
+      // which would produce invalid syntax as function parameter names.
+      return `function ${baseMockName}(${argsString}) {
   // Empty mock - original function mocked out
 }`;
     }
@@ -497,13 +559,13 @@ export default function constructMockCode(
 
     if (firstArgIsFunctionOrUnknown && returnsFunction) {
       // HOC pattern detected - return the first argument
-      return `function ${mockName}(${argsString}) {
+      return `function ${baseMockName}(${argsString}) {
   return arg1;
 }`;
     }
 
     // Generate empty mock function
-    return `function ${mockName}(${argsString}) {
+    return `function ${baseMockName}(${argsString}) {
   // Empty mock - original function mocked out
 }`;
   }
@@ -530,6 +592,99 @@ export default function constructMockCode(
     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: string;
+      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: ReturnValuePart = {
     name: dataStructureName,
     isArray: isRootArray,
@@ -591,10 +746,6 @@ export default function constructMockCode(
       }
 
       const addReturnValueFunctionAccessor = (dataPath: string) => {
-        // 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 ||
@@ -603,7 +754,17 @@ export default function constructMockCode(
           if (returnValue.isArray) {
             dataPath = `${dataPath}${optionalAccess('[0]')}`;
           }
-          dataPath = `${dataPath}${optionalAccess(`(${safeString(returnValue.returnsFunctionArgs.join(', '))})`)}`;
+          // Only add the function call accessor ?.["(args)"] when there are actual
+          // arguments. When returnsFunctionArgs is empty [] (function-returns-function
+          // with no specific arg patterns), skip the ?.["()"] because:
+          // 1. preprocessSchemaForMocks collapses nested functionCallReturnValue chains
+          //    into flat entries (e.g., getTranslate() = string, not {(): string})
+          // 2. The mock data is a flat value, so ?.["()"] on a string returns undefined
+          // 3. constructContent still wraps the return in a function (via returnsFunctionArgs)
+          //    so the function-returns-function behavior is preserved without data nesting
+          if (returnValue.returnsFunctionArgs.length > 0) {
+            dataPath = `${dataPath}${optionalAccess(`(${safeString(returnValue.returnsFunctionArgs.join(', '))})`)}`;
+          }
         }
         return dataPath;
       };
@@ -644,29 +805,6 @@ export default function constructMockCode(
         hasNoReturnData,
       } = returnValue;
 
-      // When an array has differentiated indices ([0], [1], etc.), filter out any
-      // non-index items from nested. These non-index items come from generic [] paths
-      // like [].filter or [].sort, which describe element properties, not array elements.
-      // Including them would generate invalid syntax like "sort: ..." inside an array literal.
-      const hasDifferentiatedIndices =
-        isArray &&
-        nested &&
-        nested.some((n) => n.name.match(/^\[\d+\]$/) && n.name !== '[0]');
-      const filteredNested =
-        hasDifferentiatedIndices && nested
-          ? nested.filter((n) => n.name.match(/^\[\d+\]$/))
-          : nested;
-
-      const nestedContent: (string | undefined)[] = (filteredNested ?? []).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([
@@ -711,6 +849,48 @@ export default function constructMockCode(
         'length',
       ]);
 
+      // When an array has differentiated indices ([0], [1], etc.), filter out any
+      // non-index items from nested. These non-index items come from generic [] paths
+      // like [].filter or [].sort, which describe element properties, not array elements.
+      // Including them would generate invalid syntax like "sort: ..." inside an array literal.
+      const hasDifferentiatedIndices =
+        isArray &&
+        nested &&
+        nested.some((n) => n.name.match(/^\[\d+\]$/) && n.name !== '[0]');
+      let filteredNested =
+        hasDifferentiatedIndices && nested
+          ? nested.filter((n) => n.name.match(/^\[\d+\]$/))
+          : nested;
+
+      // When a variable IS an array (not a function returning an array),
+      // filter out array prototype methods like .includes(), .filter(), etc.
+      // ONLY when the base schema has no method-call entries. When the base
+      // schema has methods (e.g., .map().functionCallReturnValue), the scenario
+      // data is stored as an object with method-call keys, and ALL methods
+      // need mock implementations. When the base schema has no methods, the
+      // data is a raw array and native methods like .includes() work directly.
+      if (
+        isArray &&
+        !returnsFunctionArray &&
+        !baseSchemaHasMethodCalls &&
+        filteredNested
+      ) {
+        filteredNested = filteredNested.filter((n) => {
+          const methodName = n.name.replace(/[<(].*$/, '');
+          return !ARRAY_PROTOTYPE_METHODS.has(methodName);
+        });
+      }
+
+      const nestedContent: (string | undefined)[] = (filteredNested ?? []).map(
+        (nestedItem) => {
+          const nestedContent = constructReturnValueString(
+            nestedItem,
+            dataPaths,
+          );
+          return nestedContent;
+        },
+      );
+
       const levelContentItems = [];
       // Add spread for data paths when:
       // - Not a function returning an array, OR function returns array with custom methods
@@ -864,7 +1044,123 @@ export default function constructMockCode(
         });
 
       let returnValueContents = '';
-      if (
+
+      // 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
@@ -967,11 +1263,20 @@ export default function constructMockCode(
                   // 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
@@ -1182,6 +1487,17 @@ export default function constructMockCode(
                     .replace(/,\s*$/, '');
                   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
@@ -1192,12 +1508,14 @@ export default function constructMockCode(
               // Should become: "{ ...spread, props, filter: ..., sort: ... }"
               const trimmed = mappedContents.trim();
 
-              // Find first }, that is NOT inside a string literal
+              // Find first }, at depth 0 that is NOT inside a string literal
               // This prevents splitting keys like ?.["useQuery({ id }, { enabled })"]
+              // and also prevents finding }, inside nested arrow functions
               const findBraceCommaOutsideStrings = (
                 content: string,
               ): number => {
                 let i = 0;
+                let depth = 0; // Track brace depth to find the outer object's },
                 while (i < content.length - 1) {
                   // Skip over string literals
                   const strEnd = skipStringLiteral(content, i);
@@ -1205,9 +1523,24 @@ export default function constructMockCode(
                     i = strEnd;
                     continue;
                   }
-                  // Check for }, pattern
-                  if (content[i] === '}' && content[i + 1] === ',') {
-                    return i;
+                  // Track brace depth
+                  if (content[i] === '{') {
+                    depth++;
+                    i++;
+                    continue;
+                  }
+                  // Check for }, pattern at depth 1 (the outer object level)
+                  // We're looking for the outer object's closing brace, which is at depth 1
+                  // (we started at depth 0, opened { at depth 0 -> 1)
+                  if (content[i] === '}') {
+                    depth--;
+                    if (
+                      depth === 0 &&
+                      i + 1 < content.length &&
+                      content[i + 1] === ','
+                    ) {
+                      return i;
+                    }
                   }
                   i++;
                 }
@@ -1231,13 +1564,55 @@ export default function constructMockCode(
               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))`;
             }
           } else {
             returnValueContents = `[\n${indent(levelContents)}\n]`;
           }
         } 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}`;
+          }
         }
       }
 
@@ -1335,19 +1710,38 @@ export default function constructMockCode(
             fallbackContent;
           content = `${cleanOutTypes(name)}: ${isAsyncFunction ? 'async ' : ''}(${argsString}) => {\n${indent(funcContents)}\n}`;
         } else {
-          // No argument variants - use existing behavior
-          // But if there's nested content, we need to include it in the return object
-          // (similar to how argument variant branches handle this at line 1070-1072)
-          const hasNestedContent = validNestedContent.length > 0;
-          let funcReturnContents: string;
-          if (hasNestedContent && levelContentItems.length > 1) {
-            // Include both spread and nested content in the return
-            funcReturnContents = `{\n${indent(levelContents)}\n}`;
+          // No argument variants
+          // Check if this is an array method callback containing JSX.
+          // JSX can't be serialized to JSON, so the LLM generates [{}] as data.
+          // Instead of returning that unusable data, generate a passthrough that
+          // calls the real callback on the best available array data from siblings.
+          const containsJsx = dataPaths.some((p) => /<[A-Z]/.test(p));
+          const isArrayMethod = ARRAY_PROTOTYPE_METHODS.has(name);
+
+          if (containsJsx && isArrayMethod && dataPaths.length > 0) {
+            // Extract parent data path by removing the last ?.["..."] segment
+            const parentPath = dataPaths[0].replace(/\?\.\["[^"]*"\]$/, '');
+            const funcLines = [
+              `const _d = ${parentPath};`,
+              `const _a = Object.values(_d || {}).filter(v => Array.isArray(v) && v.length > 0 && v.some(i => i && typeof i === "object" && Object.keys(i).length > 0)).sort((a, b) => b.length - a.length);`,
+              `return _a[0] ? _a[0].${name}(${argsString}) : []`,
+            ];
+            const funcContents = funcLines.join('\n');
+            content = `${cleanOutTypes(name)}: ${isAsyncFunction ? 'async ' : ''}(${argsString}) => {\n${indent(funcContents)}\n}`;
           } else {
-            funcReturnContents = returnValueContents;
+            // But if there's nested content, we need to include it in the return object
+            // (similar to how argument variant branches handle this at line 1070-1072)
+            const hasNestedContent = validNestedContent.length > 0;
+            let funcReturnContents: string;
+            if (hasNestedContent && levelContentItems.length > 1) {
+              // Include both spread and nested content in the return
+              funcReturnContents = `{\n${indent(levelContents)}\n}`;
+            } else {
+              funcReturnContents = returnValueContents;
+            }
+            const funcContents = `return ${funcReturnContents}`;
+            content = `${cleanOutTypes(name)}: ${isAsyncFunction ? 'async ' : ''}(${argsString}) => {\n${indent(funcContents)}\n}`;
           }
-          const funcContents = `return ${funcReturnContents}`;
-          content = `${cleanOutTypes(name)}: ${isAsyncFunction ? 'async ' : ''}(${argsString}) => {\n${indent(funcContents)}\n}`;
         }
       } else {
         if (!isValidKey(name)) {
@@ -1373,34 +1767,91 @@ export default function constructMockCode(
   };
 
   // Create the return value structure
-  const sortedKeys = Object.keys(relevantReturnValueSchema ?? {}).sort(
-    (a: string, b: string) => {
-      const aParts = splitOutsideParenthesesAndArrays(a);
-      const bParts = splitOutsideParenthesesAndArrays(b);
-
-      const maxLength = Math.max(aParts.length, bParts.length);
-      for (let i = 0; i < maxLength; ++i) {
-        const aPart = aParts[i];
-        const bPart = bParts[i];
-
-        if (!aPart) return -1;
-        if (!bPart) return 1;
-
-        if (aPart === bPart) continue;
-
-        const aName = aPart.split('(')[0];
-        const bName = bPart.split('(')[0];
+  // 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: string, b: string) => {
+    const aParts = splitOutsideParenthesesAndArrays(a);
+    const bParts = splitOutsideParenthesesAndArrays(b);
+
+    const maxLength = Math.max(aParts.length, bParts.length);
+    for (let i = 0; i < maxLength; ++i) {
+      const aPart = aParts[i];
+      const bPart = bParts[i];
+
+      if (!aPart) return -1;
+      if (!bPart) return 1;
+
+      if (aPart === bPart) continue;
+
+      const aName = aPart.split('(')[0];
+      const bName = bPart.split('(')[0];
+
+      if (aName !== bName) {
+        return aName.localeCompare(bName);
+      } else {
+        return aPart.localeCompare(bPart);
+      }
+    }
 
-        if (aName !== bName) {
-          return aName.localeCompare(bName);
-        } else {
-          return aPart.localeCompare(bPart);
-        }
+    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<string>();
+  const functionCallsWithReturnValue = new Set<string>();
+
+  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);
       }
 
-      return 0;
-    },
-  );
+      // 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];
@@ -1458,9 +1909,10 @@ export default function constructMockCode(
     // 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
@@ -1489,11 +1941,10 @@ export default function constructMockCode(
           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;
@@ -1930,6 +2381,23 @@ export default function constructMockCode(
     }
   }
 
+  // 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) {
@@ -2002,16 +2470,26 @@ export default function constructMockCode(
       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});
   }
 }`;
     }
@@ -2081,7 +2559,31 @@ export default function constructMockCode(
 
     // 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:
     // 1. For call signatures: use derivedFunctionName