From 0e16de11e720d0882546224db13d1ebd4a44247d Mon Sep 17 00:00:00 2001
From: Enrico Losavio <enrico.losavio@trixter.de>
Date: Fri, 5 May 2023 17:09:47 +0200
Subject: [PATCH] Fixed issue #46

---
 author/task.py            | 13 +++++++++++--
 tests/author/test_task.py | 20 ++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/author/task.py b/author/task.py
index 34a98e4..68e69a8 100644
--- a/author/task.py
+++ b/author/task.py
@@ -201,8 +201,17 @@ def __init__(self, arguments, is_handle_task=True, wait_for_task=None, *args, **
             _LOG.debug("Handletask {}. Adding simple dependency...".format(self))
             self._add_command_tasks(*args, **kwargs)
         elif self.elements_id and self.per_element and self._is_expected_iterable(getattr(self.elements, "processed", None)):
-            for element in self.elements.processed:
-                element = ArgumentValue(self.elements.initial, element)
+            elements = []
+            if isinstance(self.elements.processed, dict):
+                for key_value in self.elements.processed.items():
+                    element = ArgumentValue(self.elements.initial, key_value)
+                    elements.append(element)
+            else:
+                for element in self.elements.processed:
+                    element = ArgumentValue(self.elements.initial, element)
+                    elements.append(element)
+
+            for element in elements:
                 self._add_required_tasks(self.required_tasks, self, elements=element, *args, **kwargs)
         else:
             self._add_required_tasks(self.required_tasks, self, elements=self.elements, *args, **kwargs)
diff --git a/tests/author/test_task.py b/tests/author/test_task.py
index 93ac0fb..69c269d 100644
--- a/tests/author/test_task.py
+++ b/tests/author/test_task.py
@@ -31,6 +31,7 @@
 )
 
 from jobtronaut.author import (
+    Arguments,
     ArgumentValue,
     Task,
     TaskWithOverrides
@@ -296,6 +297,25 @@ def test_add_command_tasks(self, mock_per_element):
         _task = TaskFixture(TASK_FIXTURE_ARGUMENTS)
         self.assertEqual(len(_task.arguments[_task.elements_id].initial), len(_task.subtasks))
 
+        # lets check if the PER_ELEMENT resolved contain the proper element/ArgumentValue
+        index = 0
+        for handle_task in _task.subtasks:
+            for command_task in handle_task.subtasks:
+                self.assertEqual(command_task.elements.processed, _task.arguments[_task.elements_id].initial[index])
+            index += 1
+
+        TASK_FIXTURE_ARGUMENTS["test_argument"] = {"a": 1, "b": 2}
+        _task = TaskFixture(TASK_FIXTURE_ARGUMENTS)
+        all_items_sorted = sorted(list(TASK_FIXTURE_ARGUMENTS["test_argument"].items()), key=lambda item: item[0])
+        all_command_tasks = []
+        for handle_task in _task.subtasks:
+            for command_task in handle_task.subtasks:
+                all_command_tasks.append(command_task)
+        all_elements_sorted = sorted([_.elements.processed for _ in all_command_tasks], key=lambda item: item[0])
+
+        self.assertEqual(all_elements_sorted, all_items_sorted)
+
+
     @patch.object(TaskFixture, "view", create=True, new=lambda x: x.cmd()[-1])
     @patch.object(TaskFixture, "cmd", create=True, new=lambda x: ["/some/executable", "test"])
     @patch("jobtronaut.author.task.EXECUTABLE_RESOLVER", new=lambda x: "/bin/echo")