fix edge map code and add tests

pufflyai · Sep 11, 2023 · 3e5a2f4 · 3e5a2f4
1 parent 1abcf36
commit 3e5a2f4
Show file tree

Hide file tree

Showing 5 changed files with 143 additions and 6 deletions.
diff --git a/packages/@pufflig/ps-chains/src/engines/dataflow/__snapshots__/runFlow.test.ts.snap b/packages/@pufflig/ps-chains/src/engines/dataflow/__snapshots__/runFlow.test.ts.snap
@@ -279,6 +279,22 @@ exports[`running a node updates the input of its children 1`] = `
 }
 `;
 
+exports[`running a node with several connections to a child should update the child correctly 1`] = `
+{
+  "n1": {
+    "input": {
+      "data": "TEST DATA",
+    },
+  },
+  "n2": {
+    "input": {
+      "data1": "TEST DATA",
+      "data2": "TEST DATA",
+    },
+  },
+}
+`;
+
 exports[`should run even when lifecycle methods are missing 1`] = `
 {
   "n1": {

diff --git a/packages/@pufflig/ps-chains/src/engines/dataflow/runFlow.test.ts b/packages/@pufflig/ps-chains/src/engines/dataflow/runFlow.test.ts
@@ -6,6 +6,7 @@ import {
   multiInput,
   multiInputWithOutput,
   multistep,
+  severalConnections,
   simpleExec,
   simpleExecWithData,
   simpleExistingState,
@@ -468,3 +469,27 @@ test("a node can run its children multiple times 2", async () => {
   expect(onNodeRunError).toHaveBeenCalledTimes(0);
   expect(res).toMatchSnapshot();
 });
+
+/**
+ * (1.1) -> (2.1)
+ *       -> (2.2)
+ */
+test("running a node with several connections to a child should update the child correctly", async () => {
+  const onNodeInputUpdate = jest.fn();
+  const onNodeRunComplete = jest.fn();
+  const onNodeRunError = jest.fn();
+  const res = await runFlow(
+    severalConnections,
+    "n1",
+    { data: "TEST DATA" },
+    {
+      onNodeRunComplete,
+      onNodeInputUpdate,
+      onNodeRunError,
+    }
+  );
+  expect(onNodeInputUpdate).toHaveBeenCalledTimes(3);
+  expect(onNodeRunComplete).toHaveBeenCalledTimes(1);
+  expect(onNodeRunError).toHaveBeenCalledTimes(0);
+  expect(res).toMatchSnapshot();
+});
diff --git a/packages/@pufflig/ps-chains/src/engines/dataflow/utils/utils.test.ts b/packages/@pufflig/ps-chains/src/engines/dataflow/utils/utils.test.ts
@@ -1,5 +1,5 @@
 import { simpleDataNode } from "../../../mocks/nodes";
-import { applyDefaultInputs } from "./utils";
+import { applyDefaultInputs, getEdgeMap, mapOutputToInput } from "./utils";
 
 test("apply default inputs without overwriting them", async () => {
   const res = applyDefaultInputs(
@@ -32,3 +32,59 @@ test("apply default inputs without overwriting them", async () => {
     }
   `);
 });
+
+test("return a map of edges between two nodes", async () => {
+  const res = getEdgeMap([
+    {
+      id: "1",
+      source: "1",
+      sourceHandle: "1",
+      target: "2",
+      targetHandle: "1",
+    },
+    {
+      id: "2",
+      source: "1",
+      sourceHandle: "2",
+      target: "2",
+      targetHandle: "2",
+    },
+    {
+      id: "3",
+      source: "1",
+      sourceHandle: "1",
+      target: "2",
+      targetHandle: "3",
+    },
+  ]);
+
+  expect(res).toMatchInlineSnapshot(`
+    {
+      "1": [
+        "1",
+        "3",
+      ],
+      "2": [
+        "2",
+      ],
+    }
+  `);
+});
+
+test("map the outputs of a node to the inputs of a node", async () => {
+  const res = mapOutputToInput(
+    { 1: "a", 2: "b" },
+    {
+      "1": ["1", "3"],
+      "2": ["2"],
+    }
+  );
+
+  expect(res).toMatchInlineSnapshot(`
+    {
+      "1": "a",
+      "2": "b",
+      "3": "a",
+    }
+  `);
+});
diff --git a/packages/@pufflig/ps-chains/src/engines/dataflow/utils/utils.ts b/packages/@pufflig/ps-chains/src/engines/dataflow/utils/utils.ts
@@ -31,10 +31,14 @@ export function applyDefaultInputs(nodeStateData: Record<string, ParamValue> | u
  * @param edges chain edges
  * @returns record of source and target handles
  */
-export function getEdgeMap(edges: FlowEdge[]): Record<string, string> {
-  const res: Record<string, string> = {};
+export function getEdgeMap(edges: FlowEdge[]): Record<string, string[]> {
+  const res: Record<string, string[]> = {};
   for (const e of edges) {
-    res[e.sourceHandle] = e.targetHandle;
+    if (!res[e.sourceHandle]) {
+      res[e.sourceHandle] = [e.targetHandle];
+    } else {
+      res[e.sourceHandle].push(e.targetHandle);
+    }
   }
   return res;
 }
@@ -45,10 +49,10 @@ export function getEdgeMap(edges: FlowEdge[]): Record<string, string> {
  * @param map
  * @returns
  */
-export function mapOutputToInput(output: Record<string, ParamValue>, map: Record<string, string>) {
+export function mapOutputToInput(output: Record<string, ParamValue>, map: Record<string, string[]>) {
   const res: Record<string, ParamValue> = {};
   for (const [key, value] of Object.entries(map)) {
-    res[value] = output[key];
+    value.forEach((v) => (res[v] = output[key]));
   }
   return res;
 }
diff --git a/packages/@pufflig/ps-chains/src/mocks/chains.ts b/packages/@pufflig/ps-chains/src/mocks/chains.ts
@@ -694,3 +694,39 @@ export const loopWithJoin: Flow = {
     },
   },
 };
+
+export const severalConnections: Flow = {
+  nodeTypes: {
+    simple_node: simpleExecNode,
+    multi_node: multiInputDataNode,
+  },
+  definition: {
+    edges: {
+      e1: {
+        id: "e1",
+        source: "n1",
+        target: "n2",
+        sourceHandle: "data",
+        targetHandle: "data1",
+      },
+      e2: {
+        id: "e2",
+        source: "n1",
+        target: "n2",
+        sourceHandle: "data",
+        targetHandle: "data2",
+      },
+    },
+    nodes: {
+      n1: {
+        id: "n1",
+        type: "simple_node",
+      },
+      n2: {
+        id: "n2",
+        type: "multi_node",
+      },
+    },
+  },
+  state: {},
+};