update unit tests

Signed-off-by: Paul Dittamo <[email protected]>

Author: pvditt

tests/flytekit/unit/core/test_array_node.py

@@ -378,3 +378,47 @@ def test_map_task_wrapper():
 
     mapped_lp = map_task(lp)(a=[1, 3, 5], b=[2, 4, 6], c=[7, 8, 9])
     assert mapped_lp == [14, 96, 270]
+
+
+def test_run_all_sub_nodes_default():
+    node = array_node(lp, concurrency=10, min_success_ratio=0.9)
+    assert node.run_all_sub_nodes is False
+
+
+def test_run_all_sub_nodes_set():
+    node = array_node(lp, concurrency=10, min_success_ratio=0.9, run_all_sub_nodes=True)
+    assert node.run_all_sub_nodes is True
+
+
+def test_run_all_sub_nodes_serialization(serialization_settings):
+    @workflow
+    def wf_run_all() -> typing.List[int]:
+        return map_task(lp, concurrency=10, min_success_ratio=0.9, run_all_sub_nodes=True)(
+            a=[1, 3, 5], b=["two", 4, "six"], c=[7, 8, 9]
+        )
+
+    od = OrderedDict()
+    wf_spec = get_serializable(od, serialization_settings, wf_run_all)
+
+    parent_node = wf_spec.template.nodes[0]
+    assert parent_node.array_node._run_all_sub_nodes is True
+
+    pb = parent_node.array_node.to_flyte_idl()
+    assert pb.run_all_sub_nodes is True
+
+
+def test_run_all_sub_nodes_serialization_default(serialization_settings):
+    @workflow
+    def wf_no_run_all() -> typing.List[int]:
+        return map_task(lp, concurrency=10, min_success_ratio=0.9)(
+            a=[1, 3, 5], b=["two", 4, "six"], c=[7, 8, 9]
+        )
+
+    od = OrderedDict()
+    wf_spec = get_serializable(od, serialization_settings, wf_no_run_all)
+
+    parent_node = wf_spec.template.nodes[0]
+    assert parent_node.array_node._run_all_sub_nodes is False
+
+    pb = parent_node.array_node.to_flyte_idl()
+    assert pb.run_all_sub_nodes is False

tests/flytekit/unit/core/test_array_node_map_task.py

@@ -578,6 +578,7 @@ def wf1(x: typing.List[int]):
     assert array_node.array_node._parallelism == 10
     assert not array_node.array_node._is_original_sub_node_interface
     assert array_node.array_node._execution_mode == _core_workflow.ArrayNode.MINIMAL_STATE
+    assert not array_node.array_node._run_all_sub_nodes
     task_spec = od[arraynode_maptask]
     assert task_spec.template.metadata.retries.retries == 2
     assert task_spec.template.metadata.interruptible
@@ -588,6 +589,67 @@ def wf1(x: typing.List[int]):
     assert array_node.array_node._execution_mode == _core_workflow.ArrayNode.FULL_STATE
 
 
+def test_run_all_sub_nodes_default():
+    @task
+    def t1(a: int) -> int:
+        return a + 1
+
+    mt = map_task(t1)
+    assert mt.run_all_sub_nodes is False
+
+
+def test_run_all_sub_nodes_set():
+    @task
+    def t1(a: int) -> int:
+        return a + 1
+
+    mt = map_task(t1, run_all_sub_nodes=True)
+    assert mt.run_all_sub_nodes is True
+
+
+def test_run_all_sub_nodes_serialization(serialization_settings):
+    @task
+    def t1(a: int) -> int:
+        return a + 1
+
+    arraynode_maptask = map_task(t1, run_all_sub_nodes=True)
+
+    @workflow
+    def wf(x: typing.List[int]):
+        return arraynode_maptask(a=x)
+
+    od = OrderedDict()
+    wf_spec = get_serializable(od, serialization_settings, wf)
+
+    array_node = wf_spec.template.nodes[0]
+    assert array_node.array_node._run_all_sub_nodes is True
+
+    # Verify it serializes to the protobuf correctly
+    pb = array_node.array_node.to_flyte_idl()
+    assert pb.run_all_sub_nodes is True
+
+
+def test_run_all_sub_nodes_serialization_default(serialization_settings):
+    @task
+    def t1(a: int) -> int:
+        return a + 1
+
+    arraynode_maptask = map_task(t1)
+
+    @workflow
+    def wf(x: typing.List[int]):
+        return arraynode_maptask(a=x)
+
+    od = OrderedDict()
+    wf_spec = get_serializable(od, serialization_settings, wf)
+
+    array_node = wf_spec.template.nodes[0]
+    assert array_node.array_node._run_all_sub_nodes is False
+
+    pb = array_node.array_node.to_flyte_idl()
+    assert pb.run_all_sub_nodes is False
+
+
 def test_serialization_extended_resources(serialization_settings):
     @task(
         accelerator=GPUAccelerator("test_gpu"),