Adapt tests

Signed-off-by: Fabio Grätz <[email protected]>
flyteorg · Feb 29, 2024 · e68a4e8 · e68a4e8
1 parent 5aa854d
commit e68a4e8
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/tests/flytekit/unit/core/test_array_node_map_task.py b/tests/flytekit/unit/core/test_array_node_map_task.py
@@ -305,4 +305,4 @@ def my_mappable_task(a: int) -> typing.Optional[str]:
     def wf(x: typing.List[int]):
         array_node_map_task(my_mappable_task)(a=x).with_overrides(container_image="random:image")
 
-    assert wf.nodes[0].run_entity.container_image == "random:image"
+    assert wf.nodes[0]._container_image == "random:image"
diff --git a/tests/flytekit/unit/core/test_map_task.py b/tests/flytekit/unit/core/test_map_task.py
@@ -352,7 +352,7 @@ def my_mappable_task(a: int) -> typing.Optional[str]:
     def wf(x: typing.List[int]):
         map_task(my_mappable_task)(a=x).with_overrides(container_image="random:image")
 
-    assert wf.nodes[0].flyte_entity.run_task.container_image == "random:image"
+    assert wf.nodes[0]._container_image == "random:image"
 
 
 def test_bounded_inputs_vars_order(serialization_settings):

diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py
@@ -465,7 +465,7 @@ def wf() -> str:
         bar().with_overrides(container_image="hello/world")
         return "hi"
 
-    assert wf.nodes[0].flyte_entity.container_image == "hello/world"
+    assert wf.nodes[0]._container_image == "hello/world"
 
 
 def test_override_accelerator():