chore/COMPASS-9767 Add editable interactions stress test story

src/components/diagram.stories.tsx

@@ -6,6 +6,7 @@ import { EMPLOYEES_TO_EMPLOYEES_EDGE, ORDERS_TO_EMPLOYEES_EDGE } from '@/mocks/d
 import { DiagramStressTestDecorator } from '@/mocks/decorators/diagram-stress-test.decorator';
 import { DiagramConnectableDecorator } from '@/mocks/decorators/diagram-connectable.decorator';
 import { DiagramEditableInteractionsDecorator } from '@/mocks/decorators/diagram-editable-interactions.decorator';
+import { DiagramEditableStressTestDecorator } from '@/mocks/decorators/diagram-editable-stress-test.decorator';
 
 const diagram: Meta<typeof Diagram> = {
   title: 'Diagram',
@@ -91,3 +92,13 @@ export const DiagramStressTest: Story = {
     nodes: [],
   },
 };
+
+export const DiagramEditableStressTest: Story = {
+  decorators: [DiagramEditableStressTestDecorator],
+  args: {
+    title: 'MongoDB Diagram',
+    isDarkMode: true,
+    edges: [],
+    nodes: [],
+  },
+};

src/mocks/decorators/diagram-editable-interactions.decorator.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useState, MouseEvent as ReactMouseEvent } from 'react';
+import { useCallback, useEffect, useRef, useState, MouseEvent as ReactMouseEvent } from 'react';
 import { Decorator } from '@storybook/react';
 
 import { DiagramProps, FieldId, NodeField, NodeProps } from '@/types';
@@ -49,8 +49,51 @@ function addFieldToNode(existingFields: NodeField[], parentFieldPath: string[])
   return fields;
 }
 
-export const DiagramEditableInteractionsDecorator: Decorator<DiagramProps> = (Story, context) => {
-  const [nodes, setNodes] = useState<NodeProps[]>(context.args.nodes);
+let idAccumulator: string[];
+let lastDepth = 0;
+// Used to build a string array id based on field depth.
+function idFromDepthAccumulator(name: string, depth?: number) {
+  if (!depth) {
+    idAccumulator = [name];
+  } else if (depth > lastDepth) {
+    idAccumulator.push(name);
+  } else if (depth === lastDepth) {
+    idAccumulator[idAccumulator.length - 1] = name;
+  } else {
+    idAccumulator = idAccumulator.slice(0, depth);
+    idAccumulator[depth] = name;
+  }
+  lastDepth = depth ?? 0;
+  return [...idAccumulator];
+}
+function editableNodesFromNodes(nodes: NodeProps[]): NodeProps[] {
+  return nodes.map(node => ({
+    ...node,
+    type: 'collection',
+    fields: node.fields.map(field => ({
+      ...field,
+      selectable: true,
+      id: idFromDepthAccumulator(field.name, field.depth),
+    })),
+  }));
+}
+
+export const useEditableNodes = (initialNodes: NodeProps[]) => {
+  const [nodes, setNodes] = useState<NodeProps[]>([]);
+
+  const hasInitialized = useRef(false);
+  useEffect(() => {
+    if (hasInitialized.current) {
+      return;
+    }
+
+    if (!initialNodes || initialNodes.length === 0) {
+      return;
+    }
+
+    hasInitialized.current = true;
+    setNodes(editableNodesFromNodes(initialNodes));
+  }, [initialNodes]);
 
   const onFieldClick = useCallback(
     (
@@ -61,18 +104,32 @@ export const DiagramEditableInteractionsDecorator: Decorator<DiagramProps> = (St
       },
     ) => {
       setNodes(nodes =>
-        nodes.map(node => ({
-          ...node,
-          fields: node.fields.map(field => ({
-            ...field,
-            selected:
+        nodes.map(node => {
+          let nodeFieldDidChange = false;
+          const fields = node.fields.map(field => {
+            const selected =
               params.nodeId === node.id &&
               !!field.id &&
               typeof field.id !== 'string' &&
               typeof params.id !== 'string' &&
-              stringArrayCompare(params.id, field.id),
-          })),
-        })),
+              stringArrayCompare(params.id, field.id);
+            if (field.selected !== selected) {
+              nodeFieldDidChange = true;
+            }
+            return {
+              ...field,
+              selected,
+            };
+          });
+
+          if (!nodeFieldDidChange) {
+            return node;
+          }
+          return {
+            ...node,
+            fields,
+          };
+        }),
       );
     },
     [],
@@ -107,14 +164,17 @@ export const DiagramEditableInteractionsDecorator: Decorator<DiagramProps> = (St
     [],
   );
 
+  return { nodes, onFieldClick, onAddFieldToNodeClick, onAddFieldToObjectFieldClick };
+};
+
+export const DiagramEditableInteractionsDecorator: Decorator<DiagramProps> = (Story, context) => {
+  const editableArgs = useEditableNodes(context.args.nodes || []);
+
   return Story({
     ...context,
     args: {
       ...context.args,
-      nodes,
-      onFieldClick,
-      onAddFieldToNodeClick,
-      onAddFieldToObjectFieldClick,
+      ...editableArgs,
     },
   });
 };

src/mocks/decorators/diagram-editable-stress-test.decorator.tsx

@@ -0,0 +1,19 @@
+import { Decorator } from '@storybook/react';
+
+import { DiagramProps } from '@/types';
+import { useEditableNodes } from '@/mocks/decorators/diagram-editable-interactions.decorator';
+import { useStressTestNodesAndEdges } from '@/mocks/decorators/diagram-stress-test.decorator';
+
+export const DiagramEditableStressTestDecorator: Decorator<DiagramProps> = (Story, context) => {
+  const { nodes: initialNodes, edges } = useStressTestNodesAndEdges(100);
+  const editableArgs = useEditableNodes(initialNodes);
+
+  return Story({
+    ...context,
+    args: {
+      ...context.args,
+      ...editableArgs,
+      edges: editableArgs.nodes.length > 0 ? edges : [],
+    },
+  });
+};

src/mocks/decorators/diagram-stress-test.decorator.tsx

@@ -17,45 +17,81 @@ const names = [
   'api_keys',
 ];
 
-export const DiagramStressTestDecorator: Decorator<DiagramProps> = (Story, context) => {
+const types = ['string', 'number', 'boolean', 'date', 'object', 'array'];
+
+let previousWasObject = false;
+let previousDepth = 0;
+function getRandomTypeAndDepth(i: number) {
+  if (i === 0) {
+    previousWasObject = false;
+    previousDepth = 0;
+  }
+  const type = types[Math.floor(Math.random() * types.length)];
+
+  const depth = previousWasObject
+    ? Math.random() > 0.25
+      ? previousDepth + 1
+      : previousDepth
+    : previousDepth > 0 && Math.random() > 0.2
+      ? previousDepth
+      : 0;
+
+  previousWasObject = type === 'object';
+
+  previousDepth = depth || 0;
+
+  return {
+    type,
+    depth,
+  };
+}
+
+const generateNodes = (count: number): NodeProps[] => {
+  return Array.from(Array(count).keys()).map((nodeIndex: number) => ({
+    id: `node_${nodeIndex}`,
+    type: 'table',
+    position: {
+      x: 0,
+      y: 0,
+    },
+    title: names[Math.floor(Math.random() * names.length)],
+    fields: Array.from(Array(1 + (nodeIndex % 9)).keys()).map((fieldIndex: number) => ({
+      name: `${names[Math.floor(Math.random() * names.length)]}-${fieldIndex}`,
+      ...getRandomTypeAndDepth(fieldIndex),
+    })),
+  }));
+};
+
+export const useStressTestNodesAndEdges = (nodeCount: number) => {
   const [nodes, setNodes] = useState<NodeProps[]>([]);
   const [edges, setEdges] = useState<EdgeProps[]>([]);
 
   useEffect(() => {
-    const nodes = generateNodes(100);
-    const edges = generateEdges(nodes);
+    const newNodes = generateNodes(nodeCount);
+    const newEdges: EdgeProps[] = newNodes.map(node => ({
+      id: `edge_${node.id}`,
+      source: newNodes[Math.floor(Math.random() * newNodes.length)].id,
+      target: newNodes[Math.floor(Math.random() * newNodes.length)].id,
+      markerStart: 'many',
+      markerEnd: 'one',
+    }));
 
-    applyLayout<NodeProps, EdgeProps>(nodes, edges, 'STAR').then(result => {
+    let applyUpdate = true;
+    applyLayout<NodeProps, EdgeProps>(newNodes, newEdges, 'STAR').then(result => {
+      if (!applyUpdate) return;
       setNodes(result.nodes);
       setEdges(result.edges);
     });
-  }, []);
+    return () => {
+      applyUpdate = false;
+    };
+  }, [nodeCount]);
 
-  const generateEdges = (nodes: NodeProps[]): EdgeProps[] => {
-    return nodes.map(node => ({
-      id: `edge_${node.id}`,
-      source: nodes[Math.floor(Math.random() * nodes.length)].id,
-      target: nodes[Math.floor(Math.random() * nodes.length)].id,
-      markerStart: 'many',
-      markerEnd: 'one',
-    }));
-  };
+  return { nodes, edges };
+};
 
-  const generateNodes = (count: number): NodeProps[] => {
-    return Array.from(Array(count).keys()).map(i => ({
-      id: `node_${i}`,
-      type: 'table',
-      position: {
-        x: 0,
-        y: 0,
-      },
-      title: names[Math.floor(Math.random() * names.length)],
-      fields: Array.from(Array(1 + (i % 9)).keys()).map(_ => ({
-        name: names[Math.floor(Math.random() * names.length)],
-        type: 'varchar',
-      })),
-    }));
-  };
+export const DiagramStressTestDecorator: Decorator<DiagramProps> = (Story, context) => {
+  const { nodes, edges } = useStressTestNodesAndEdges(100);
 
   return Story({
     ...context,