Add tests for the new Pydantic I/O syntax

Signed-off-by: Alice Purcell <[email protected]>

tests/test_pydantic_io_syntax.py

@@ -0,0 +1,62 @@
+import pytest
+
+from hera.shared._global_config import _SCRIPT_PYDANTIC_IO_FLAG
+from hera.workflows import Input, Output, Steps, Workflow, script
+
+
+class IntInput(Input):
+    field: int
+
+
+class IntOutput(Output):
+    field: int
+
+
+@pytest.fixture(autouse=True)
+def enable_pydantic_io(global_config_fixture):
+    global_config_fixture.experimental_features[_SCRIPT_PYDANTIC_IO_FLAG] = True
+
+
+def test_output_field_contains_argo_template(global_config_fixture):
+    @script()
+    def double(input: IntInput) -> IntOutput:
+        return IntOutput(field=input.field * 2)
+
+    with Workflow(name="foo"):
+        with Steps(name="bar"):
+            result = double(IntInput(field=5)).field
+
+    assert result == "{{steps.double.outputs.parameters.field}}"
+
+
+def test_script_can_return_none():
+    @script()
+    def print_field(input: IntInput) -> None:
+        print(input.field)
+
+    with Workflow(name="foo"):
+        with Steps(name="bar"):
+            result = print_field(IntInput(field=5))
+
+    assert result is None
+
+
+def test_invalid_pydantic_io_outside_of_context():
+    @script()
+    def double(input: IntInput) -> IntOutput:
+        return IntOutput(field=input.field * 2)
+
+    with Workflow(name="foo"):
+        with pytest.raises(SyntaxError, match="Cannot use Pydantic I/O outside of a .* context"):
+            double(IntInput(field=5))
+
+
+def test_invalid_non_pydantic_return_type():
+    @script()
+    def double(input: IntInput) -> int:
+        return input.field * 2
+
+    with Workflow(name="foo"):
+        with Steps(name="bar"):
+            with pytest.raises(SyntaxError, match="Cannot use Pydantic input type without a Pydantic output type"):
+                double(IntInput(field=5))