Array length map nesting

Author: jcobis

packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.spec.ts

@@ -256,7 +256,7 @@ describe('Script Generation', () => {
         collectionName: 'posts',
         documentCount: 1,
         arrayLengthMap: {
-          tags: 5,
+          tags: [5],
         },
       });
 
@@ -277,7 +277,7 @@ describe('Script Generation', () => {
         documentCount: 1,
         arrayLengthMap: {
           users: {
-            tags: 4,
+            tags: [4],
           },
         },
       });
@@ -289,55 +289,57 @@ describe('Script Generation', () => {
       }
     });
 
-    it('should handle hierarchical array length map', () => {
+    it('should handle zero-length arrays', () => {
       const schema = {
-        'departments[].employees[].skills[]': createFieldMapping('lorem.word'),
+        'tags[]': createFieldMapping('lorem.word'),
+        'categories[]': createFieldMapping('lorem.word'),
       };
 
       const result = generateScript(schema, {
         databaseName: 'testdb',
-        collectionName: 'company',
+        collectionName: 'posts',
         documentCount: 1,
         arrayLengthMap: {
-          departments: {
-            employees: {
-              skills: 3,
-            },
-          },
+          tags: [0],
+          categories: [2],
         },
       });
 
       expect(result.success).to.equal(true);
       if (result.success) {
-        expect(result.script).to.contain('Array.from({length: 3}');
+        // Should have tags array with length 0 (empty array)
+        expect(result.script).to.contain('Array.from({length: 0}');
+        // Should have categories array with length 2
+        expect(result.script).to.contain('Array.from({length: 2}');
+        // Verify both arrays are present
+        expect(result.script).to.contain('tags:');
+        expect(result.script).to.contain('categories:');
       }
     });
 
-    it('should handle zero-length arrays', () => {
+    it('should handle multi-dimensional arrays with custom lengths', () => {
       const schema = {
-        'tags[]': createFieldMapping('lorem.word'),
-        'categories[]': createFieldMapping('lorem.word'),
+        'matrix[][]': createFieldMapping('number.int'),
+        'cube[][][]': createFieldMapping('number.float'),
       };
 
       const result = generateScript(schema, {
         databaseName: 'testdb',
-        collectionName: 'posts',
+        collectionName: 'data',
         documentCount: 1,
         arrayLengthMap: {
-          tags: 0,
-          categories: 2,
+          matrix: [2, 5], // 2x5 matrix
+          cube: [3, 4, 2], // 3x4x2 cube
         },
       });
 
       expect(result.success).to.equal(true);
       if (result.success) {
-        // Should have tags array with length 0 (empty array)
-        expect(result.script).to.contain('Array.from({length: 0}');
-        // Should have categories array with length 2
-        expect(result.script).to.contain('Array.from({length: 2}');
-        // Verify both arrays are present
-        expect(result.script).to.contain('tags:');
-        expect(result.script).to.contain('categories:');
+        const expectedReturnBlock = `return {
+    matrix: Array.from({length: 2}, () => Array.from({length: 5}, () => faker.number.int())),
+    cube: Array.from({length: 3}, () => Array.from({length: 4}, () => Array.from({length: 2}, () => faker.number.float())))
+  };`;
+        expect(result.script).to.contain(expectedReturnBlock);
       }
     });
   });

packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts

@@ -12,7 +12,7 @@ export interface FieldMapping {
 
 // Hierarchical array length map that mirrors document structure
 export type ArrayLengthMap = {
-  [fieldName: string]: number | ArrayLengthMap;
+  [fieldName: string]: number[] | ArrayLengthMap;
 };
 
 export interface ScriptOptions {
@@ -351,24 +351,24 @@ function generateDocumentCode(
       }
     } else if ('type' in value && value.type === 'array') {
       // It's an array
-      const nestedArrayLengthMap =
-        typeof arrayLengthMap[fieldName] === 'object'
-          ? arrayLengthMap[fieldName]
-          : {};
       const arrayCode = generateArrayCode(
         value as ArrayStructure,
         indent + INDENT_SIZE,
         fieldName,
         arrayLengthMap,
-        nestedArrayLengthMap
+        0 // Start at dimension 0
       );
       rootLevelFields.push(`${fieldIndent}${fieldName}: ${arrayCode}`);
     } else {
       // It's a nested object: recursive call
+
+      // Type Validation/ Fallback for misconfigured array length map
+      const arrayDimensions = arrayLengthMap[fieldName];
       const nestedArrayLengthMap =
-        typeof arrayLengthMap[fieldName] === 'object'
-          ? arrayLengthMap[fieldName]
-          : arrayLengthMap;
+        typeof arrayDimensions === 'object' && !Array.isArray(arrayDimensions)
+          ? arrayDimensions
+          : {};
+
       const nestedCode = generateDocumentCode(
         value as DocumentStructure,
         indent + INDENT_SIZE,
@@ -393,33 +393,37 @@ function generateArrayCode(
   arrayStructure: ArrayStructure,
   indent: number = INDENT_SIZE,
   fieldName: string = '',
-  parentArrayLengthMap: ArrayLengthMap = {},
-  nestedArrayLengthMap: ArrayLengthMap = {}
+  arrayLengthMap: ArrayLengthMap = {},
+  dimensionIndex: number = 0
 ): string {
   const elementType = arrayStructure.elementType;
 
-  // Get array length from map or use default
-  const arrayLength =
-    typeof parentArrayLengthMap[fieldName] === 'number'
-      ? parentArrayLengthMap[fieldName]
-      : DEFAULT_ARRAY_LENGTH;
+  // Get array length for this dimension
+  const arrayDimensions = arrayLengthMap[fieldName];
+  const arrayLength = Array.isArray(arrayDimensions)
+    ? arrayDimensions[dimensionIndex] ?? DEFAULT_ARRAY_LENGTH
+    : DEFAULT_ARRAY_LENGTH; // Fallback for misconfigured array length map
 
   if ('mongoType' in elementType) {
     // Array of primitives
     const fakerCall = generateFakerCall(elementType as FieldMapping);
     return `Array.from({length: ${arrayLength}}, () => ${fakerCall})`;
   } else if ('type' in elementType && elementType.type === 'array') {
-    // Nested array (e.g., matrix[][])
+    // Nested array (e.g., matrix[][]) - keep same fieldName, increment dimension
     const nestedArrayCode = generateArrayCode(
       elementType as ArrayStructure,
       indent,
-      '', // No specific field name for nested arrays
-      nestedArrayLengthMap,
-      {}
+      fieldName,
+      arrayLengthMap,
+      dimensionIndex + 1 // Next dimension
     );
     return `Array.from({length: ${arrayLength}}, () => ${nestedArrayCode})`;
   } else {
     // Array of objects
+    const nestedArrayLengthMap =
+      typeof arrayDimensions === 'object' && !Array.isArray(arrayDimensions)
+        ? arrayDimensions
+        : {}; // Fallback for misconfigured array length map
     const objectCode = generateDocumentCode(
       elementType as DocumentStructure,
       indent,