Fix taskmetadata

Signed-off-by: byhsu <[email protected]>
flyteorg · Jan 26, 2023 · 32ea24b · 32ea24b
1 parent f0b243a
commit 32ea24b
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py
@@ -11,7 +11,7 @@
 from kubernetes.client.models import V1EnvVar, V1ResourceRequirements
 
 from flytekit.configuration import ImageConfig, SerializationSettings
-from flytekit.core.base_task import PythonTask, TaskResolverMixin
+from flytekit.core.base_task import PythonTask, TaskMetadata, TaskResolverMixin
 from flytekit.core.context_manager import FlyteContextManager
 from flytekit.core.pod_template import PodTemplate
 from flytekit.core.resources import Resources, ResourceSpec
@@ -86,8 +86,8 @@ def __init__(
             sec_ctx = SecurityContext(secrets=secret_requests)
 
         # pod_template_name overwrites the metedata.pod_template_name
-        if "metadata" in kwargs:
-            kwargs["metadata"].pod_template_name = pod_template_name
+        kwargs["metadata"] = kwargs["metadata"] if "metadata" in kwargs else TaskMetadata()
+        kwargs["metadata"].pod_template_name = pod_template_name
 
         super().__init__(
             task_type=task_type,

diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py
@@ -233,6 +233,7 @@ def test_pod_template(default_serialization_settings):
             ]
         ),
     ),
+    pod_template_name="A",
 )
 
 
@@ -277,10 +278,16 @@ def test_minimum_pod_template(default_serialization_settings):
         "task_with_minimum_pod_template",
     ]
 
+    #################
+    # Test pod_teamplte_name
+    #################
+    assert task_with_minimum_pod_template.metadata.pod_template_name == "A"
+
     #################
     # Test Serialization
     #################
     ts = get_serializable_task(default_serialization_settings, task_with_minimum_pod_template)
     assert ts.template.container is None
     # k8s_pod content is already verified above, so only check the existence here
     assert ts.template.k8s_pod is not None
+    assert ts.template.metadata.pod_template_name == "A"
diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py
@@ -142,7 +142,7 @@ def test_pod_template():
                 ]
             ),
         ),
-        pod_template_name="podTemplateA",
+        pod_template_name="A",
     )
     def func_with_pod_template(i: str):
         print(i + 3)
@@ -192,10 +192,16 @@ def func_with_pod_template(i: str):
         "func_with_pod_template",
     ]
 
+    #################
+    # Test pod_teamplte_name
+    #################
+    assert func_with_pod_template.metadata.pod_template_name == "A"
+
     #################
     # Test Serialization
     #################
     ts = get_serializable_task(default_serialization_settings, func_with_pod_template)
     assert ts.template.container is None
     # k8s_pod content is already verified above, so only check the existence here
     assert ts.template.k8s_pod is not None
+    assert ts.template.metadata.pod_template_name == "A"