feat: test buildNodesAndEdgesFromDSLComponents (#940)

### What problem does this PR solve?
feat: test buildNodesAndEdgesFromDSLComponents #918

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

web/jest-setup.ts

@@ -0,0 +1,2 @@
+import '@testing-library/jest-dom';
+import 'umi/test-setup';

web/jest.config.ts

@@ -0,0 +1,33 @@
+import { Config, configUmiAlias, createConfig } from 'umi/test';
+
+export default async () => {
+  return (await configUmiAlias({
+    ...createConfig({
+      target: 'browser',
+      jsTransformer: 'esbuild',
+      // config opts for esbuild , it will pass to esbuild directly
+      jsTransformerOpts: { jsx: 'automatic' },
+    }),
+    setupFilesAfterEnv: ['<rootDir>/jest-setup.ts'],
+    collectCoverageFrom: [
+      '**/*.{ts,tsx,js,jsx}',
+      '!.umi/**',
+      '!.umi-test/**',
+      '!.umi-production/**',
+      '!.umirc.{js,ts}',
+      '!.umirc.*.{js,ts}',
+      '!jest.config.{js,ts}',
+      '!coverage/**',
+      '!dist/**',
+      '!config/**',
+      '!mock/**',
+    ],
+    // if you require some es-module npm package, please uncomment below line and insert your package name
+    // transformIgnorePatterns: ['node_modules/(?!.*(lodash-es|your-es-pkg-name)/)']
+    coverageThreshold: {
+      global: {
+        lines: 1,
+      },
+    },
+  })) as Config.InitialOptions;
+};

web/package.json

@@ -7,7 +7,8 @@
     "postinstall": "umi setup",
     "lint": "umi lint --eslint-only",
     "setup": "umi setup",
-    "start": "npm run dev"
+    "start": "npm run dev",
+    "test": "jest --no-cache --coverage"
   },
   "dependencies": {
     "@ant-design/icons": "^5.2.6",
@@ -18,6 +19,7 @@
     "antd": "^5.12.7",
     "axios": "^1.6.3",
     "classnames": "^2.5.1",
+    "dagre": "^0.8.5",
     "dayjs": "^1.11.10",
     "eventsource-parser": "^1.1.2",
     "i18next": "^23.7.16",
@@ -45,20 +47,28 @@
   },
   "devDependencies": {
     "@react-dev-inspector/umi4-plugin": "^2.0.1",
+    "@testing-library/jest-dom": "^6.4.5",
+    "@testing-library/react": "^15.0.7",
+    "@types/dagre": "^0.7.52",
+    "@types/jest": "^29.5.12",
     "@types/lodash": "^4.14.202",
     "@types/react": "^18.0.33",
     "@types/react-copy-to-clipboard": "^5.0.7",
     "@types/react-dom": "^18.0.11",
     "@types/react-syntax-highlighter": "^15.5.11",
+    "@types/testing-library__jest-dom": "^6.0.0",
     "@types/uuid": "^9.0.8",
     "@types/webpack-env": "^1.18.4",
     "@umijs/lint": "^4.1.1",
     "@umijs/plugins": "^4.1.0",
     "cross-env": "^7.0.3",
+    "jest": "^29.7.0",
+    "jest-environment-jsdom": "^29.7.0",
     "prettier": "^3.2.4",
     "prettier-plugin-organize-imports": "^3.2.4",
     "prettier-plugin-packagejson": "^2.4.9",
     "react-dev-inspector": "^2.0.1",
+    "ts-node": "^10.9.2",
     "typescript": "^5.0.3",
     "umi-plugin-icons": "^0.1.1"
   }

web/src/pages/flow/canvas/index.tsx

@@ -8,7 +8,6 @@ import ReactFlow, {
   OnConnect,
   OnEdgesChange,
   OnNodesChange,
-  Position,
   addEdge,
   applyEdgeChanges,
   applyNodeChanges,
@@ -19,47 +18,24 @@ import { NodeContextMenu, useHandleNodeContextMenu } from './context-menu';
 
 import FlowDrawer from '../flow-drawer';
 import { useHandleDrop, useShowDrawer } from '../hooks';
+import { initialEdges, initialNodes } from '../mock';
+import { getLayoutedElements } from '../utils';
 import { TextUpdaterNode } from './node';
 
 const nodeTypes = { textUpdater: TextUpdaterNode };
 
-const initialNodes = [
-  {
-    sourcePosition: Position.Left,
-    targetPosition: Position.Right,
-    id: 'node-1',
-    type: 'textUpdater',
-    position: { x: 400, y: 100 },
-    data: { label: 123 },
-  },
-  {
-    sourcePosition: Position.Right,
-    targetPosition: Position.Left,
-    id: '1',
-    data: { label: 'Hello' },
-    position: { x: 0, y: 50 },
-    type: 'input',
-  },
-  {
-    sourcePosition: Position.Right,
-    targetPosition: Position.Left,
-    id: '2',
-    data: { label: 'World' },
-    position: { x: 200, y: 50 },
-  },
-];
-
-const initialEdges = [
-  { id: '1-2', source: '1', target: '2', label: 'to the', type: 'step' },
-];
-
 interface IProps {
   sideWidth: number;
 }
 
 function FlowCanvas({ sideWidth }: IProps) {
-  const [nodes, setNodes] = useState<Node[]>(initialNodes);
-  const [edges, setEdges] = useState<Edge[]>(initialEdges);
+  const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(
+    initialNodes,
+    initialEdges,
+    'LR',
+  );
+  const [nodes, setNodes] = useState<Node[]>(layoutedNodes);
+  const [edges, setEdges] = useState<Edge[]>(layoutedEdges);
   const { ref, menu, onNodeContextMenu, onPaneClick } =
     useHandleNodeContextMenu(sideWidth);
   const { drawerVisible, hideDrawer, showDrawer } = useShowDrawer();

web/src/pages/flow/mock.tsx

@@ -3,6 +3,7 @@ import {
   RocketOutlined,
   SendOutlined,
 } from '@ant-design/icons';
+import { Position } from 'reactflow';
 
 export const componentList = [
   { name: 'Begin', icon: <SendOutlined />, description: '' },
@@ -10,6 +11,39 @@ export const componentList = [
   { name: 'Generate', icon: <MergeCellsOutlined />, description: '' },
 ];
 
+export const initialNodes = [
+  {
+    sourcePosition: Position.Left,
+    targetPosition: Position.Right,
+    id: 'node-1',
+    type: 'textUpdater',
+    position: { x: 0, y: 0 },
+    // position: { x: 400, y: 100 },
+    data: { label: 123 },
+  },
+  {
+    sourcePosition: Position.Right,
+    targetPosition: Position.Left,
+    id: '1',
+    data: { label: 'Hello' },
+    position: { x: 0, y: 0 },
+    // position: { x: 0, y: 50 },
+    type: 'input',
+  },
+  {
+    sourcePosition: Position.Right,
+    targetPosition: Position.Left,
+    id: '2',
+    data: { label: 'World' },
+    position: { x: 0, y: 0 },
+    // position: { x: 200, y: 50 },
+  },
+];
+
+export const initialEdges = [
+  { id: '1-2', source: '1', target: '2', label: 'to the', type: 'step' },
+];
+
 export const dsl = {
   components: {
     begin: {
@@ -17,8 +51,8 @@ export const dsl = {
         component_name: 'Begin',
         params: {},
       },
-      downstream: ['Answer:China'],
-      upstream: [],
+      downstream: ['Answer:China'], // other edge target is downstream, edge source is current node id
+      upstream: [], // edge source is upstream, edge target is current node id
     },
     'Answer:China': {
       obj: {

web/src/pages/flow/utils.test.ts

@@ -0,0 +1,30 @@
+import { dsl } from './mock';
+import { buildNodesAndEdgesFromDSLComponents } from './utils';
+
+test('buildNodesAndEdgesFromDSLComponents', () => {
+  const { edges, nodes } = buildNodesAndEdgesFromDSLComponents(dsl.components);
+
+  expect(nodes.length).toEqual(4);
+  expect(edges.length).toEqual(4);
+
+  expect(edges).toEqual(
+    expect.arrayContaining([
+      expect.objectContaining({
+        source: 'begin',
+        target: 'Answer:China',
+      }),
+      expect.objectContaining({
+        source: 'Answer:China',
+        target: 'Retrieval:China',
+      }),
+      expect.objectContaining({
+        source: 'Retrieval:China',
+        target: 'Generate:China',
+      }),
+      expect.objectContaining({
+        source: 'Generate:China',
+        target: 'Answer:China',
+      }),
+    ]),
+  );
+});

web/src/pages/flow/utils.ts

@@ -1,10 +1,32 @@
 import { DSLComponents } from '@/interfaces/database/flow';
+import dagre from 'dagre';
 import { Edge, Node, Position } from 'reactflow';
 import { v4 as uuidv4 } from 'uuid';
 
-export const buildNodesFromDSLComponents = (data: DSLComponents) => {
+const buildEdges = (
+  operatorIds: string[],
+  currentId: string,
+  allEdges: Edge[],
+  isUpstream = false,
+) => {
+  operatorIds.forEach((cur) => {
+    const source = isUpstream ? cur : currentId;
+    const target = isUpstream ? currentId : cur;
+    if (!allEdges.some((e) => e.source === source && e.target === target)) {
+      allEdges.push({
+        id: uuidv4(),
+        label: '',
+        type: 'step',
+        source: source,
+        target: target,
+      });
+    }
+  });
+};
+
+export const buildNodesAndEdgesFromDSLComponents = (data: DSLComponents) => {
   const nodes: Node[] = [];
-  const edges: Edge[] = [];
+  let edges: Edge[] = [];
 
   Object.entries(data).forEach(([key, value]) => {
     const downstream = [...value.downstream];
@@ -23,22 +45,51 @@ export const buildNodesFromDSLComponents = (data: DSLComponents) => {
       targetPosition: Position.Right,
     });
 
-    // intermediate node
-    // The first and last nodes do not need to be considered
-    if (upstream.length > 0 && downstream.length > 0) {
-      for (let i = 0; i < upstream.length; i++) {
-        const up = upstream[i];
-        for (let j = 0; j < downstream.length; j++) {
-          const down = downstream[j];
-          edges.push({
-            id: uuidv4(),
-            label: '',
-            type: 'step',
-            source: up,
-            target: down,
-          });
-        }
-      }
-    }
+    buildEdges(upstream, key, edges, true);
+    buildEdges(downstream, key, edges, false);
+  });
+
+  return { nodes, edges };
+};
+
+const dagreGraph = new dagre.graphlib.Graph();
+dagreGraph.setDefaultEdgeLabel(() => ({}));
+
+const nodeWidth = 172;
+const nodeHeight = 36;
+
+export const getLayoutedElements = (
+  nodes: Node[],
+  edges: Edge[],
+  direction = 'TB',
+) => {
+  const isHorizontal = direction === 'LR';
+  dagreGraph.setGraph({ rankdir: direction });
+
+  nodes.forEach((node) => {
+    dagreGraph.setNode(node.id, { width: nodeWidth, height: nodeHeight });
+  });
+
+  edges.forEach((edge) => {
+    dagreGraph.setEdge(edge.source, edge.target);
+  });
+
+  dagre.layout(dagreGraph);
+
+  nodes.forEach((node) => {
+    const nodeWithPosition = dagreGraph.node(node.id);
+    node.targetPosition = isHorizontal ? Position.Left : Position.Top;
+    node.sourcePosition = isHorizontal ? Position.Right : Position.Bottom;
+
+    // We are shifting the dagre node position (anchor=center center) to the top left
+    // so it matches the React Flow node anchor point (top left).
+    node.position = {
+      x: nodeWithPosition.x - nodeWidth / 2,
+      y: nodeWithPosition.y - nodeHeight / 2,
+    };
+
+    return node;
   });
+
+  return { nodes, edges };
 };

web/tsconfig.json

@@ -1,6 +1,6 @@
 {
   "extends": "./src/.umi/tsconfig.json",
   "@@/*": [
-    "src/.umi/*"
+    "src/.umi/*",
   ],
 }