Update tests/flytekit/unit/extras/templates/test_pod_templates.py

Co-authored-by: Flyte Bot <[email protected]>
flyteorg · Dec 30, 2024 · 067c902 · 067c902
1 parent 44e2ea3
commit 067c902
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tests/flytekit/unit/extras/templates/test_pod_templates.py b/tests/flytekit/unit/extras/templates/test_pod_templates.py
@@ -2,9 +2,13 @@
 from flytekit.core.task import task
 
 def test_attach_shm():
-
+ 
     shm = attach_shm("SHM", "5Gi")
-
-    @task(pod_template=shm)
+    assert shm.name == "SHM"
+    assert shm.size == "5Gi"
+
     def my_task():
         pass
+
+    # Verify pod template is attached to task
+    assert my_task.pod_template == shm