tests

tests/test_async_flow.py

@@ -97,7 +97,7 @@ class TestAsyncFlow(unittest.TestCase):
 
         # We'll run the flow synchronously (which under the hood is asyncio.run())
         shared_storage = {}
-        flow.run(shared_storage)
+        asyncio.run(flow.run_async(shared_storage))
 
         self.assertEqual(shared_storage['current'], 6)
 
@@ -144,7 +144,7 @@ class TestAsyncFlow(unittest.TestCase):
         start_node - "negative_branch" >> negative_node
 
         flow = AsyncFlow(start_node)
-        flow.run(shared_storage)
+        asyncio.run(flow.run_async(shared_storage))
 
         self.assertEqual(shared_storage["path"], "positive", 
                          "Should have taken the positive branch")

tests/test_batch_flow.py

@@ -0,0 +1,184 @@
+import unittest
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent.parent))
+from minillmflow import Node, BatchFlow, Flow
+
+class DataProcessNode(Node):
+    def process(self, shared_storage, prep_result):
+        key = self.params.get('key')
+        data = shared_storage['input_data'][key]
+        if 'results' not in shared_storage:
+            shared_storage['results'] = {}
+        shared_storage['results'][key] = data * 2
+
+class ErrorProcessNode(Node):
+    def process(self, shared_storage, prep_result):
+        key = self.params.get('key')
+        if key == 'error_key':
+            raise ValueError(f"Error processing key: {key}")
+        if 'results' not in shared_storage:
+            shared_storage['results'] = {}
+        shared_storage['results'][key] = True
+
+class TestBatchFlow(unittest.TestCase):
+    def setUp(self):
+        self.process_node = DataProcessNode()
+        
+    def test_basic_batch_processing(self):
+        """Test basic batch processing with multiple keys"""
+        class SimpleTestBatchFlow(BatchFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {
+                'a': 1,
+                'b': 2,
+                'c': 3
+            }
+        }
+
+        flow = SimpleTestBatchFlow(start_node=self.process_node)
+        flow.run(shared_storage)
+
+        expected_results = {
+            'a': 2,
+            'b': 4,
+            'c': 6
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+    def test_empty_input(self):
+        """Test batch processing with empty input dictionary"""
+        class EmptyTestBatchFlow(BatchFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {}
+        }
+
+        flow = EmptyTestBatchFlow(start_node=self.process_node)
+        flow.run(shared_storage)
+
+        self.assertEqual(shared_storage.get('results', {}), {})
+
+    def test_single_item(self):
+        """Test batch processing with single item"""
+        class SingleItemBatchFlow(BatchFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {
+                'single': 5
+            }
+        }
+
+        flow = SingleItemBatchFlow(start_node=self.process_node)
+        flow.run(shared_storage)
+
+        expected_results = {
+            'single': 10
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+    def test_error_handling(self):
+        """Test error handling during batch processing"""
+        class ErrorTestBatchFlow(BatchFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        shared_storage = {
+            'input_data': {
+                'normal_key': 1,
+                'error_key': 2,
+                'another_key': 3
+            }
+        }
+
+        flow = ErrorTestBatchFlow(start_node=ErrorProcessNode())
+        
+        with self.assertRaises(ValueError):
+            flow.run(shared_storage)
+
+    def test_nested_flow(self):
+        """Test batch processing with nested flows"""
+        class InnerNode(Node):
+            def process(self, shared_storage, prep_result):
+                key = self.params.get('key')
+                if 'intermediate_results' not in shared_storage:
+                    shared_storage['intermediate_results'] = {}
+                shared_storage['intermediate_results'][key] = shared_storage['input_data'][key] + 1
+
+        class OuterNode(Node):
+            def process(self, shared_storage, prep_result):
+                key = self.params.get('key')
+                if 'results' not in shared_storage:
+                    shared_storage['results'] = {}
+                shared_storage['results'][key] = shared_storage['intermediate_results'][key] * 2
+
+        class NestedBatchFlow(BatchFlow):
+            def preprocess(self, shared_storage):
+                return [{'key': k} for k in shared_storage['input_data'].keys()]
+
+        # Create inner flow
+        inner_node = InnerNode()
+        outer_node = OuterNode()
+        inner_node >> outer_node
+
+        shared_storage = {
+            'input_data': {
+                'x': 1,
+                'y': 2
+            }
+        }
+
+        flow = NestedBatchFlow(start_node=inner_node)
+        flow.run(shared_storage)
+
+        expected_results = {
+            'x': 4,  # (1 + 1) * 2
+            'y': 6   # (2 + 1) * 2
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+    def test_custom_parameters(self):
+        """Test batch processing with additional custom parameters"""
+        class CustomParamNode(Node):
+            def process(self, shared_storage, prep_result):
+                key = self.params.get('key')
+                multiplier = self.params.get('multiplier', 1)
+                if 'results' not in shared_storage:
+                    shared_storage['results'] = {}
+                shared_storage['results'][key] = shared_storage['input_data'][key] * multiplier
+
+        class CustomParamBatchFlow(BatchFlow):
+            def preprocess(self, shared_storage):
+                return [{
+                    'key': k,
+                    'multiplier': i + 1
+                } for i, k in enumerate(shared_storage['input_data'].keys())]
+
+        shared_storage = {
+            'input_data': {
+                'a': 1,
+                'b': 2,
+                'c': 3
+            }
+        }
+
+        flow = CustomParamBatchFlow(start_node=CustomParamNode())
+        flow.run(shared_storage)
+
+        expected_results = {
+            'a': 1 * 1,  # first item, multiplier = 1
+            'b': 2 * 2,  # second item, multiplier = 2
+            'c': 3 * 3   # third item, multiplier = 3
+        }
+        self.assertEqual(shared_storage['results'], expected_results)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_batch_node.py

@@ -0,0 +1,162 @@
+import unittest
+import sys
+from pathlib import Path
+
+sys.path.append(str(Path(__file__).parent.parent))
+from minillmflow import Node, BatchNode, Flow
+
+class ArrayChunkNode(BatchNode):
+    def __init__(self, chunk_size=10):
+        super().__init__()
+        self.chunk_size = chunk_size
+    
+    def preprocess(self, shared_storage):
+        # Get array from shared storage and split into chunks
+        array = shared_storage.get('input_array', [])
+        chunks = []
+        for i in range(0, len(array), self.chunk_size):
+            end = min(i + self.chunk_size, len(array))
+            chunks.append((i, end))
+        return chunks
+    
+    def process(self, shared_storage, chunk_indices):
+        start, end = chunk_indices
+        array = shared_storage['input_array']
+        # Process the chunk and return its sum
+        chunk_sum = sum(array[start:end])
+        return chunk_sum
+        
+    def postprocess(self, shared_storage, prep_result, proc_result):
+        # Store chunk results in shared storage
+        shared_storage['chunk_results'] = proc_result
+        return "default"
+
+class SumReduceNode(Node):
+    def process(self, shared_storage, data):
+        # Get chunk results from shared storage and sum them
+        chunk_results = shared_storage.get('chunk_results', [])
+        total = sum(chunk_results)
+        shared_storage['total'] = total
+
+class TestBatchNode(unittest.TestCase):
+    def test_array_chunking(self):
+        """
+        Test that the array is correctly split into chunks
+        """
+        shared_storage = {
+            'input_array': list(range(25))  # [0,1,2,...,24]
+        }
+        
+        chunk_node = ArrayChunkNode(chunk_size=10)
+        chunks = chunk_node.preprocess(shared_storage)
+        
+        self.assertEqual(chunks, [(0, 10), (10, 20), (20, 25)])
+        
+    def test_map_reduce_sum(self):
+        """
+        Test a complete map-reduce pipeline that sums a large array:
+        1. Map: Split array into chunks and sum each chunk
+        2. Reduce: Sum all the chunk sums
+        """
+        # Create test array: [0,1,2,...,99]
+        array = list(range(100))
+        expected_sum = sum(array)  # 4950
+        
+        shared_storage = {
+            'input_array': array
+        }
+        
+        # Create nodes
+        chunk_node = ArrayChunkNode(chunk_size=10)
+        reduce_node = SumReduceNode()
+        
+        # Connect nodes
+        chunk_node >> reduce_node
+        
+        # Create and run pipeline
+        pipeline = Flow(start_node=chunk_node)
+        pipeline.run(shared_storage)
+        
+        self.assertEqual(shared_storage['total'], expected_sum)
+        
+    def test_uneven_chunks(self):
+        """
+        Test that the map-reduce works correctly with array lengths
+        that don't divide evenly by chunk_size
+        """
+        array = list(range(25))
+        expected_sum = sum(array)  # 300
+        
+        shared_storage = {
+            'input_array': array
+        }
+        
+        chunk_node = ArrayChunkNode(chunk_size=10)
+        reduce_node = SumReduceNode()
+        
+        chunk_node >> reduce_node
+        pipeline = Flow(start_node=chunk_node)
+        pipeline.run(shared_storage)
+        
+        self.assertEqual(shared_storage['total'], expected_sum)
+
+    def test_custom_chunk_size(self):
+        """
+        Test that the map-reduce works with different chunk sizes
+        """
+        array = list(range(100))
+        expected_sum = sum(array)
+        
+        shared_storage = {
+            'input_array': array
+        }
+        
+        # Use chunk_size=15 instead of default 10
+        chunk_node = ArrayChunkNode(chunk_size=15)
+        reduce_node = SumReduceNode()
+        
+        chunk_node >> reduce_node
+        pipeline = Flow(start_node=chunk_node)
+        pipeline.run(shared_storage)
+        
+        self.assertEqual(shared_storage['total'], expected_sum)
+        
+    def test_single_element_chunks(self):
+        """
+        Test extreme case where chunk_size=1
+        """
+        array = list(range(5))
+        expected_sum = sum(array)
+        
+        shared_storage = {
+            'input_array': array
+        }
+        
+        chunk_node = ArrayChunkNode(chunk_size=1)
+        reduce_node = SumReduceNode()
+        
+        chunk_node >> reduce_node
+        pipeline = Flow(start_node=chunk_node)
+        pipeline.run(shared_storage)
+        
+        self.assertEqual(shared_storage['total'], expected_sum)
+
+    def test_empty_array(self):
+        """
+        Test edge case of empty input array
+        """
+        shared_storage = {
+            'input_array': []
+        }
+        
+        chunk_node = ArrayChunkNode(chunk_size=10)
+        reduce_node = SumReduceNode()
+        
+        chunk_node >> reduce_node
+        pipeline = Flow(start_node=chunk_node)
+        pipeline.run(shared_storage)
+        
+        self.assertEqual(shared_storage['total'], 0)
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_flow_composition.py

@@ -2,10 +2,11 @@ import unittest
 import asyncio
 import sys
 from pathlib import Path
-
 sys.path.append(str(Path(__file__).parent.parent))
+
 from minillmflow import Node, Flow
 
+# Simple example Nodes
 class NumberNode(Node):
     def __init__(self, number):
         super().__init__()
@@ -30,106 +31,79 @@ class MultiplyNode(Node):
     def process(self, shared_storage, prep_result):
         shared_storage['current'] *= self.number
 
-
 class TestFlowComposition(unittest.TestCase):
-
     def test_flow_as_node(self):
         """
-        Demonstrates that a Flow can itself be chained like a Node.
-        We create a flow (f1) that starts with NumberNode(5) -> AddNode(10).
-        Then we chain f1 >> MultiplyNode(2).
-        
-        Expected result after running from f1:
-          start = 5
-          5 + 10 = 15
-          15 * 2 = 30
+        1) Create a Flow (f1) starting with NumberNode(5), then AddNode(10), then MultiplyNode(2).
+        2) Create a second Flow (f2) whose start_node is f1.
+        3) Create a wrapper Flow (f3) that contains f2 to ensure proper execution.
+        Expected final result in shared_storage['current']: (5 + 10) * 2 = 30.
         """
         shared_storage = {}
-
+        
         # Inner flow f1
         f1 = Flow(start_node=NumberNode(5))
-        f1 >> AddNode(10)
-
-        # Then chain a node after the flow
-        f1 >> MultiplyNode(2)
-
-        # Run from f1
-        f1.run(shared_storage)
-
+        f1 >> AddNode(10) >> MultiplyNode(2)
+        
+        # f2 starts with f1
+        f2 = Flow(start_node=f1)
+        
+        # Wrapper flow f3 to ensure proper execution
+        f3 = Flow(start_node=f2)
+        f3.run(shared_storage)
+        
         self.assertEqual(shared_storage['current'], 30)
 
     def test_nested_flow(self):
         """
-        Demonstrates embedding one Flow inside another Flow.
+        Demonstrates nested flows with proper wrapping:
         inner_flow: NumberNode(5) -> AddNode(3)
-        outer_flow: starts with inner_flow -> MultiplyNode(4)
-        
-        Expected result:
-          (5 + 3) * 4 = 32
+        middle_flow: starts with inner_flow -> MultiplyNode(4)
+        wrapper_flow: contains middle_flow to ensure proper execution
+        Expected final result: (5 + 3) * 4 = 32.
         """
         shared_storage = {}
         
-        # Define an inner flow
+        # Build the inner flow
         inner_flow = Flow(start_node=NumberNode(5))
         inner_flow >> AddNode(3)
-
-        # Define an outer flow, whose start node is inner_flow
-        outer_flow = Flow(start_node=inner_flow)
-        outer_flow >> MultiplyNode(4)
-
-        # Run outer_flow
-        outer_flow.run(shared_storage)
-
-        self.assertEqual(shared_storage['current'], 32)  # (5+3)*4=32
+        
+        # Build the middle flow, whose start_node is the inner flow
+        middle_flow = Flow(start_node=inner_flow)
+        middle_flow >> MultiplyNode(4)
+        
+        # Wrapper flow to ensure proper execution
+        wrapper_flow = Flow(start_node=middle_flow)
+        wrapper_flow.run(shared_storage)
+        
+        self.assertEqual(shared_storage['current'], 32)
 
     def test_flow_chaining_flows(self):
         """
-        Demonstrates chaining one flow to another flow.
-        flow1: NumberNode(10) -> AddNode(10)  # final shared_storage['current'] = 20
-        flow2: MultiplyNode(2)               # final shared_storage['current'] = 40
-        
-        flow1 >> flow2 means once flow1 finishes, flow2 starts.
-        
-        Expected result: (10 + 10) * 2 = 40
+        Demonstrates chaining two flows with proper wrapping:
+        flow1: NumberNode(10) -> AddNode(10) # final = 20
+        flow2: MultiplyNode(2) # final = 40
+        wrapper_flow: contains both flow1 and flow2 to ensure proper execution
+        Expected final result: (10 + 10) * 2 = 40.
         """
         shared_storage = {}
 
         # flow1
-        flow1 = Flow(start_node=NumberNode(10))
-        flow1 >> AddNode(10)
+        numbernode = NumberNode(10)
+        numbernode >> AddNode(10)
+        flow1 = Flow(start_node=numbernode)
 
         # flow2
         flow2 = Flow(start_node=MultiplyNode(2))
 
-        # Chain them: flow1 >> flow2
+        # Chain flow1 to flow2
         flow1 >> flow2
 
-        # Start running from flow1
-        flow1.run(shared_storage)
-
+        # Wrapper flow to ensure proper execution
+        wrapper_flow = Flow(start_node=flow1)
+        wrapper_flow.run(shared_storage)
+        
         self.assertEqual(shared_storage['current'], 40)
 
-    def test_flow_with_parameters(self):
-        """
-        Demonstrates passing parameters into a Flow (and retrieved by a Node).
-        """
-
-        class ParamNode(Node):
-            def process(self, shared_storage, prep_result):
-                # Reads 'level' from the node's (or flow's) parameters
-                shared_storage['param'] = self.parameters.get('level', 'no param')
-
-        shared_storage = {}
-
-        # Create a flow with a ParamNode
-        f = Flow(start_node=ParamNode())
-        # Set parameters on the flow
-        f.parameters = {'level': 'Level 1'}
-
-        f.run(shared_storage)
-
-        self.assertEqual(shared_storage['param'], 'Level 1')
-
-
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()