diff --git a/tests/flytekit/unit/core/test_array_node_map_task.py b/tests/flytekit/unit/core/test_array_node_map_task.py
index c4329041fd..f2d5b927c0 100644
--- a/tests/flytekit/unit/core/test_array_node_map_task.py
+++ b/tests/flytekit/unit/core/test_array_node_map_task.py
@@ -9,7 +9,7 @@
 import pytest
 from flyteidl.core import workflow_pb2 as _core_workflow
 
-from flytekit import dynamic, map_task, task, workflow, eager, PythonFunctionTask, Resources
+from flytekit import dynamic, map, task, workflow, eager, PythonFunctionTask, Resources
 from flytekit.configuration import FastSerializationSettings, Image, ImageConfig, SerializationSettings
 from flytekit.core import context_manager
 from flytekit.core.array_node_map_task import ArrayNodeMapTask, ArrayNodeMapTaskResolver