diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 91f29e3576..ca5d20f48f 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -241,16 +241,16 @@ def test_union_type1(input): @pytest.mark.parametrize( "extra_cli_args, task_name, expected_output", [ - (("--a",), "test_task_boolean", True), - (("--no_a",), "test_task_boolean", False), + (("--a",), "test_boolean", True), + (("--no_a",), "test_boolean", False), - (tuple(), "test_task_boolean_default_true", True), - (("--a",), "test_task_boolean_default_true", True), - (("--no_a",), "test_task_boolean_default_true", False), + (tuple(), "test_boolean_default_true", True), + (("--a",), "test_boolean_default_true", True), + (("--no_a",), "test_boolean_default_true", False), - (tuple(), "test_task_boolean_default_false", False), - (("--a",), "test_task_boolean_default_false", True), - (("--no_a",), "test_task_boolean_default_false", False), + (tuple(), "test_boolean_default_false", False), + (("--a",), "test_boolean_default_false", True), + (("--no_a",), "test_boolean_default_false", False), ], ) def test_boolean_type(extra_cli_args, task_name, expected_output): @@ -422,6 +422,9 @@ def test_get_entities_in_file(workflow_file): "task_with_env_vars", "task_with_list", "task_with_optional", + "test_boolean", + "test_boolean_default_false", + "test_boolean_default_true", "test_union1", "test_union2", ] @@ -436,6 +439,9 @@ def test_get_entities_in_file(workflow_file): "task_with_env_vars", "task_with_list", "task_with_optional", + "test_boolean", + "test_boolean_default_false", + "test_boolean_default_true", "test_union1", "test_union2", ] diff --git a/tests/flytekit/unit/cli/pyflyte/workflow.py b/tests/flytekit/unit/cli/pyflyte/workflow.py index 496f5c28b9..7aca5d3bf4 100644 --- a/tests/flytekit/unit/cli/pyflyte/workflow.py +++ b/tests/flytekit/unit/cli/pyflyte/workflow.py @@ -81,15 +81,15 @@ def test_union2(a: typing.Union[float, typing.List[int], MyDataclass]): print(a) @task -def test_task_boolean(a: bool): +def test_boolean(a: bool): print(a) @task -def test_task_boolean_default_true(a: bool = True): +def test_boolean_default_true(a: bool = True): print(a) @task -def test_task_boolean_default_false(a: bool = False): +def test_boolean_default_false(a: bool = False): print(a)