fix test

flytekit/core/pod_template.py

@@ -1,10 +1,13 @@
 from dataclasses import dataclass
-from kubernetes.client.models import V1PodSpec
 from typing import Dict, Optional
+
+from kubernetes.client.models import V1PodSpec
+
 from flytekit.exceptions import user as _user_exceptions
 
 PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
 
+
 @dataclass
 class PodTemplate(object):
     pod_spec: V1PodSpec

flytekit/models/task.py

@@ -319,6 +319,7 @@ def from_flyte_idl(cls, pb2_object):
             discovery_version=pb2_object.discovery_version,
             deprecated_error_message=pb2_object.deprecated_error_message,
             cache_serializable=pb2_object.cache_serializable,
+            pod_template_name=pb2_object.pod_template_name,
         )
 
 

tests/flytekit/unit/models/test_tasks.py

@@ -71,6 +71,7 @@ def test_task_metadata():
         "0.1.1b0",
         "This is deprecated!",
         True,
+        "A",
     )
 
     assert obj.discoverable is True
@@ -82,6 +83,7 @@ def test_task_metadata():
     assert obj.runtime.version == "1.0.0"
     assert obj.deprecated_error_message == "This is deprecated!"
     assert obj.discovery_version == "0.1.1b0"
+    assert obj.pod_template_name == "A"
     assert obj == task.TaskMetadata.from_flyte_idl(obj.to_flyte_idl())
 
 
@@ -134,6 +136,7 @@ def test_task_spec():
         "0.1.1b0",
         "This is deprecated!",
         True,
+        "A",
     )
 
     int_type = types.LiteralType(types.SimpleType.INTEGER)
@@ -178,7 +181,7 @@ def test_task_spec():
     assert obj.template == template
 
 
-def test_task_template__k8s_pod_target():
+def test_task_template_k8s_pod_target():
     int_type = types.LiteralType(types.SimpleType.INTEGER)
     obj = task.TaskTemplate(
         identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version"),
@@ -192,6 +195,7 @@ def test_task_template__k8s_pod_target():
             "1.0",
             "deprecated",
             False,
+            "A",
         ),
         interface_models.TypedInterface(
             # inputs

tests/flytekit/unit/models/test_workflow_closure.py

@@ -38,6 +38,7 @@ def test_workflow_closure():
         "0.1.1b0",
         "This is deprecated!",
         True,
+        "A",
     )
 
     cpu_resource = _task.Resources.ResourceEntry(_task.Resources.ResourceName.CPU, "1")