add more connection tests

Signed-off-by: Teo Koon Peng <[email protected]>

Author: koonpeng

diagram-editor/frontend/edges/buffer-edge.tsx

@@ -1,7 +1,7 @@
 import { type EdgeProps, StepEdge } from '@xyflow/react';
+import { memo } from 'react';
 import type { Edge } from '../types/react-flow';
 import type { SectionBufferSlotData } from './section-edge';
-import { memo } from 'react';
 
 export type BufferKeySlotData = {
   type: 'bufferKey';

diagram-editor/frontend/edges/create-edge.ts

@@ -2,6 +2,8 @@ import { MarkerType } from '@xyflow/react';
 import { v4 as uuidv4 } from 'uuid';
 import type { DiagramEditorEdge, EdgeOutputData } from '.';
 import type { BufferKeySlotData, BufferSeqSlotData } from './buffer-edge';
+import type { DefaultEdge } from './default-edge';
+import type { SectionBufferSlotData } from './section-edge';
 
 export function createBaseEdge(
   source: string,
@@ -14,20 +16,21 @@ export function createBaseEdge(
     target,
     markerEnd: {
       type: MarkerType.ArrowClosed,
-      width: 24,
-      height: 24,
+      width: 20,
+      height: 20,
     },
   };
 }
 
 export function createDefaultEdge(
   source: string,
   target: string,
+  inputSlot?: DefaultEdge['data']['input'],
 ): DiagramEditorEdge<'default'> {
   return {
     ...createBaseEdge(source, target),
     type: 'default',
-    data: { output: {}, input: { type: 'default' } },
+    data: { output: {}, input: inputSlot || { type: 'default' } },
   };
 }
 
@@ -106,7 +109,7 @@ export function createSplitRemainingEdge(
 export function createBufferEdge(
   source: string,
   target: string,
-  data: BufferKeySlotData | BufferSeqSlotData,
+  data: BufferKeySlotData | BufferSeqSlotData | SectionBufferSlotData,
 ): DiagramEditorEdge<'buffer'> {
   return {
     ...createBaseEdge(source, target),

diagram-editor/frontend/edges/default-edge.tsx

@@ -2,9 +2,9 @@ import { type EdgeProps, StepEdge } from '@xyflow/react';
 import { memo } from 'react';
 import type { DataEdge } from '.';
 
-export type DefaultOutputData = Record<string, never>;
+export type DefaultEdgeOutputData = Record<string, never>;
 
-export type DefaultEdge = DataEdge<DefaultOutputData, 'default'>;
+export type DefaultEdge = DataEdge<DefaultEdgeOutputData, 'default'>;
 
 export type DefaultEdgeCompProps = Omit<EdgeProps<DefaultEdge>, 'label'>;
 

diagram-editor/frontend/edges/index.ts

@@ -29,12 +29,12 @@ export type { SplitSeqEdge } from './split-seq-edge';
 export type { StreamOutEdge } from './stream-out-edge';
 export type { UnzipEdge } from './unzip-edge';
 
-export type DefaultInputData = { type: 'default' };
+export type DefaultInputSlotData = { type: 'default' };
 
 export type DataEdge<
   O extends Record<string, unknown>,
   S extends string,
-> = Edge<O, DefaultInputData | SectionInputSlotData, S>;
+> = Edge<O, DefaultInputSlotData | SectionInputSlotData, S>;
 
 export type EdgeMapping = {
   default: DefaultEdge;

diagram-editor/frontend/edges/stream-out-edge.tsx

@@ -1,15 +1,15 @@
 import { type EdgeProps, StepEdge } from '@xyflow/react';
 import { memo } from 'react';
 import type { Edge } from '../types/react-flow';
-import type { DefaultInputData } from '.';
+import type { DefaultInputSlotData } from '.';
 
 export type StreamOutOutputData = {
   name: string;
 };
 
 export type StreamOutEdge = Edge<
   StreamOutOutputData,
-  DefaultInputData,
+  DefaultInputSlotData,
   'streamOut'
 >;
 

diagram-editor/frontend/utils/connection.test.ts

@@ -411,4 +411,85 @@ describe('validate edges', () => {
       expect(result.valid).toBe(false);
     }
   });
+
+  test('buffer edges connecting to a section must have "sectionBuffer" input', () => {
+    const bufferNode = createOperationNode(
+      ROOT_NAMESPACE,
+      undefined,
+      { x: 0, y: 0 },
+      {
+        type: 'buffer',
+      },
+      'test_op_buffer',
+    );
+    const sectionNode = createOperationNode(
+      ROOT_NAMESPACE,
+      undefined,
+      { x: 0, y: 0 },
+      { type: 'section', builder: 'test_section' },
+      'test_op_section',
+    );
+
+    {
+      const reactFlow = new MockReactFlowAccessor(
+        [bufferNode, sectionNode],
+        [],
+      );
+      const edge = createBufferEdge(bufferNode.id, sectionNode.id, {
+        type: 'bufferSeq',
+        seq: 0,
+      });
+      const result = validateEdgeSimple(edge, reactFlow);
+      expect(result.valid).toBe(false);
+    }
+    {
+      const reactFlow = new MockReactFlowAccessor(
+        [bufferNode, sectionNode],
+        [],
+      );
+      const edge = createBufferEdge(bufferNode.id, sectionNode.id, {
+        type: 'sectionBuffer',
+        inputId: 'test',
+      });
+      const result = validateEdgeSimple(edge, reactFlow);
+      expect(result.valid).toBe(true);
+    }
+  });
+
+  test('data edges connecting to a section must have "sectionInput" input', () => {
+    const nodeNode = createOperationNode(
+      ROOT_NAMESPACE,
+      undefined,
+      { x: 0, y: 0 },
+      {
+        type: 'node',
+        builder: 'test_builder',
+        next: { builtin: 'dispose' },
+      },
+      'test_op_node',
+    );
+    const sectionNode = createOperationNode(
+      ROOT_NAMESPACE,
+      undefined,
+      { x: 0, y: 0 },
+      { type: 'section', builder: 'test_section' },
+      'test_op_section',
+    );
+
+    {
+      const reactFlow = new MockReactFlowAccessor([nodeNode, sectionNode], []);
+      const edge = createDefaultEdge(nodeNode.id, sectionNode.id);
+      const result = validateEdgeSimple(edge, reactFlow);
+      expect(result.valid).toBe(false);
+    }
+    {
+      const reactFlow = new MockReactFlowAccessor([nodeNode, sectionNode], []);
+      const edge = createDefaultEdge(nodeNode.id, sectionNode.id, {
+        type: 'sectionInput',
+        inputId: 'test',
+      });
+      const result = validateEdgeSimple(edge, reactFlow);
+      expect(result.valid).toBe(true);
+    }
+  });
 });