Rename test tasks and fix test_get_entities_in_file

flyteorg · Sep 23, 2024 · b844fa3 · b844fa3
1 parent 928bee3
commit b844fa3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git 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
@@ -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)