Fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · Feb 10, 2022 · 7a2320b · 7a2320b
1 parent 7629730
commit 7a2320b
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py b/plugins/flytekit-aws-batch/flytekitplugins/awsbatch/task.py
@@ -43,7 +43,6 @@ def __init__(self, task_config: AWSBatchConfig, task_function: Callable, **kwarg
         super(AWSBatchFunctionTask, self).__init__(
             task_config=task_config, task_type=self._AWS_BATCH_TASK_TYPE, task_function=task_function, **kwargs
         )
-        self._run_task = PythonFunctionTask(task_config=None, task_function=task_function)
         self._task_config = task_config
 
     def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
@@ -66,9 +65,9 @@ def get_command(self, settings: SerializationSettings) -> List[str]:
             "{{.rawOutputDataPrefix}}",
             "--is-aws-batch-single-job",
             "--resolver",
-            self._run_task.task_resolver.location,
+            self.task_resolver.location,
             "--",
-            *self._run_task.task_resolver.loader_args(settings, self._run_task),
+            *self.task_resolver.loader_args(settings, self),
         ]
 
         return container_args