More tests

Signed-off-by: Kevin Su <[email protected]>

tests/flytekit/unit/core/test_composition.py

@@ -1,4 +1,4 @@
-import typing
+from typing import Dict, List, NamedTuple, Optional, Union
 
 import pytest
 
@@ -10,7 +10,7 @@
 
 def test_wf1_with_subwf():
     @task
-    def t1(a: int) -> typing.NamedTuple("OutputsBC", t1_int_output=int, c=str):
+    def t1(a: int) -> NamedTuple("OutputsBC", t1_int_output=int, c=str):
         a = a + 2
         return a, "world-" + str(a)
 
@@ -132,7 +132,7 @@ def my_wf2(a: int, b: int = 42) -> (str, str, int, int):
 
 def test_wf1_with_lp_node():
     @task
-    def t1(a: int) -> typing.NamedTuple("OutputsBC", t1_int_output=int, c=str):
+    def t1(a: int) -> NamedTuple("OutputsBC", t1_int_output=int, c=str):
         a = a + 2
         return a, "world-" + str(a)
 
@@ -175,26 +175,26 @@ def my_wf3(a: int = 42) -> (int, str, str, str):
 
 def test_optional_input():
     @task()
-    def t1(b: typing.Optional[int] = None) -> typing.Optional[int]:
+    def t1(a: Optional[int] = None, b: Optional[List[int]] = None, c: Optional[Dict[str, int]] = None) -> Optional[int]:
         ...
 
     @task()
-    def t2(c: typing.Optional[int] = None) -> typing.Optional[int]:
+    def t2(a: Union[int, Optional[List[int]], None] = None) -> Union[int, Optional[List[int]], None]:
         ...
 
     @workflow
-    def wf(a: typing.Optional[int] = 1) -> typing.Optional[int]:
+    def wf(a: Optional[int] = 1) -> Optional[int]:
         t1()
-        return t2(c=a)
+        return t2(a=a)
 
     assert wf() is None
 
-    @task()
-    def t3(c: typing.Optional[int] = 3) -> typing.Optional[int]:
-        ...
-
     with pytest.raises(ValueError, match="The default value for the optional type must be None, but got 3"):
 
+        @task()
+        def t3(c: Optional[int] = 3) -> Optional[int]:
+            ...
+
         @workflow
         def wf():
             return t3()

tests/flytekit/unit/core/test_interface.py

@@ -200,7 +200,9 @@ def z(a: Annotated[int, "some annotation"]) -> Annotated[int, "some annotation"]
     assert our_interface.inputs == {"a": Annotated[int, "some annotation"]}
     assert our_interface.outputs == {"o0": Annotated[int, "some annotation"]}
 
-    def z(a: typing.Optional[int] = None, b: typing.Optional[str] = None) -> typing.Tuple[int, str]:
+    def z(
+        a: typing.Optional[int] = None, b: typing.Optional[str] = None, c: typing.Union[typing.List[int], None] = None
+    ) -> typing.Tuple[int, str]:
         ...
 
     our_interface = transform_function_to_interface(z)
@@ -209,6 +211,8 @@ def z(a: typing.Optional[int] = None, b: typing.Optional[str] = None) -> typing.
     assert params.parameters["a"].default.scalar.none_type == Void()
     assert not params.parameters["b"].required
     assert params.parameters["b"].default.scalar.none_type == Void()
+    assert not params.parameters["c"].required
+    assert params.parameters["c"].default.scalar.none_type == Void()
 
 
 def test_parameters_with_docstring():