Merge pull request #230 from lzaoral/hub-create-preassigned-subtasks

hub: support creation of subtasks with an inherited worker from parent task
release-engineering · Oct 19, 2023 · 4e81c06 · 4e81c06
2 parents 7ac1d18 + db85e67
commit 4e81c06
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 4 deletions.
diff --git a/kobo/hub/xmlrpc/worker.py b/kobo/hub/xmlrpc/worker.py
@@ -234,7 +234,8 @@ def get_awaited_tasks(request, awaited_task_list):
 
 
 @validate_worker
-def create_subtask(request, label, method, args, parent_id, subtask_priority = None):
+def create_subtask(request, label, method, args, parent_id, subtask_priority = None,
+                   inherit_worker=False):
     parent_task = Task.objects.get_and_verify(task_id=parent_id, worker=request.worker)
 
     return Task.create_task(
@@ -243,6 +244,7 @@ def create_subtask(request, label, method, args, parent_id, subtask_priority = N
         method,
         args=args,
         parent_id=parent_id,
+        worker_name=(request.worker.name if inherit_worker else None),
         arch_name=parent_task.arch.name,
         channel_name=parent_task.channel.name,
         priority=subtask_priority or parent_task.priority

diff --git a/kobo/worker/task.py b/kobo/worker/task.py
@@ -90,12 +90,13 @@ def cleanup(cls, hub, conf, task_info):
     def notification(cls, hub, conf, task_info):
         pass
 
-    def spawn_subtask(self, method, args, label="", priority = None):
+    def spawn_subtask(self, method, args, label="", priority = None,
+                      inherit_worker=False):
         """Spawn a new subtask."""
         if self.foreground:
             raise RuntimeError("Foreground tasks can't spawn subtasks.")
 
-        subtask_id = self.hub.worker.create_subtask(label, method, args, self.task_id, priority)
+        subtask_id = self.hub.worker.create_subtask(label, method, args, self.task_id, priority, inherit_worker)
         self._running_subtask_list.append(subtask_id)
         return subtask_id
 

diff --git a/tests/test_taskbase.py b/tests/test_taskbase.py
@@ -144,7 +144,8 @@ def test_spawn_subtask(self):
         ret_id = t.spawn_subtask('method', [], 'label')
         self.assertEqual(ret_id, subtask_id)
 
-        hub.worker.create_subtask.assert_called_once_with('label', 'method', [], task_id, None)
+        hub.worker.create_subtask.assert_called_once_with(
+                'label', 'method', [], task_id, None, False)
 
     def test_spawn_subtask_foreground_task(self):
         task_info = {'id': 100}

diff --git a/tests/test_xmlrpc_worker.py b/tests/test_xmlrpc_worker.py
@@ -897,6 +897,29 @@ def test_create_subtask(self):
         self.assertEqual(t_child.parent.id, t_parent.id)
         self.assertEqual(t_child.label, 'Label')
         self.assertEqual(t_child.method, 'Method')
+        self.assertEqual(t_child.worker, None)
+        self.assertEqual(t_child.state, TASK_STATES['FREE'])
+
+    def test_create_subtask_with_worker(self):
+        t_parent = Task.objects.create(
+            worker=self._worker,
+            arch=self._arch,
+            channel=self._channel,
+            owner=self._user,
+            state=TASK_STATES['FREE'],
+        )
+
+        req = _make_request(self._worker)
+        task_id = worker.create_subtask(req, 'Label', 'Method', None, t_parent.id,
+                                        inherit_worker=True)
+        self.assertTrue(task_id > 0)
+
+        t_child = Task.objects.get(id=task_id)
+        self.assertEqual(t_child.parent.id, t_parent.id)
+        self.assertEqual(t_child.label, 'Label')
+        self.assertEqual(t_child.method, 'Method')
+        self.assertEqual(t_child.worker, self._worker)
+        self.assertEqual(t_child.state, TASK_STATES['ASSIGNED'])
 
     def test_create_subtask_if_another_worker_task(self):
         w = Worker.objects.create(