From f0b243a68f4579ca18a36232d13d203ae685f777 Mon Sep 17 00:00:00 2001
From: byhsu <byhsu@linkedin.com>
Date: Thu, 26 Jan 2023 01:55:48 -0800
Subject: [PATCH] fix test

Signed-off-by: byhsu <byhsu@linkedin.com>
---
 flytekit/core/pod_template.py                       | 5 ++++-
 flytekit/models/task.py                             | 1 +
 tests/flytekit/unit/models/test_tasks.py            | 6 +++++-
 tests/flytekit/unit/models/test_workflow_closure.py | 1 +
 4 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/flytekit/core/pod_template.py b/flytekit/core/pod_template.py
index e9484f4c83..5798e92590 100644
--- a/flytekit/core/pod_template.py
+++ b/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
diff --git a/flytekit/models/task.py b/flytekit/models/task.py
index 6d3cf14861..61f31b626a 100644
--- a/flytekit/models/task.py
+++ b/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,
         )
 
 
diff --git a/tests/flytekit/unit/models/test_tasks.py b/tests/flytekit/unit/models/test_tasks.py
index fed32b63aa..a979a39b66 100644
--- a/tests/flytekit/unit/models/test_tasks.py
+++ b/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
diff --git a/tests/flytekit/unit/models/test_workflow_closure.py b/tests/flytekit/unit/models/test_workflow_closure.py
index d229d0d5c9..2b5b06696b 100644
--- a/tests/flytekit/unit/models/test_workflow_closure.py
+++ b/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")