Add option for get_worker_tasks to also return ASSIGNED tasks

kobo/hub/xmlrpc/worker.py

@@ -56,18 +56,25 @@ def get_worker_id(request):
 
 
 @validate_worker
-def get_worker_tasks(request):
+def get_worker_tasks(request, also_assigned=False):
     """
-    Get list of OPEN tasks executed on a worker.
+    Get list of OPEN (and maybe ASSIGNED) tasks executed on a worker.
 
+    @param also_assigned: Whether to return ASSIGNED tasks in addition to OPEN tasks.
+    @type  also_assigned: bool
     @rtype: list
     """
     task_list = []
 
     # Worker.running_tasks returns both OPEN and ASSIGNED tasks but all calls
     # by the worker assume that only OPEN tasks will be returned.
     # See: https://github.com/release-engineering/kobo/pull/251#issue-2183712995
-    for task in request.worker.running_tasks().filter(state=TASK_STATES['OPEN']).order_by("-exclusive", "-awaited", "id"):
+    # Old behavior can be toggled with also_assigned parameter
+    if also_assigned:
+        tasks = request.worker.running_tasks().order_by("-exclusive", "-awaited", "id")
+    else:
+        tasks = request.worker.running_tasks().filter(state=TASK_STATES['OPEN']).order_by("-exclusive", "-awaited", "id")
+    for task in tasks:
         task_info = task.export()
 
         # set wakeup alert

tests/test_xmlrpc_worker.py

@@ -91,6 +91,25 @@ def test_get_worker_tasks(self):
 
         self.assertEqual(len(tasks), 1)
         self.assertEqual(tasks[0]['state'], TASK_STATES['OPEN'])
+
+    def test_get_worker_tasks_assigned(self):
+        for state in TASK_STATES:
+            Task.objects.create(
+                worker=self._worker,
+                arch=self._arch,
+                channel=self._channel,
+                owner=self._user,
+                state=TASK_STATES[state],
+            )
+
+        req = _make_request(self._worker)
+        tasks = worker.get_worker_tasks(req, True)
+
+        self.assertEqual(len(tasks), 2)
+        self.assertTrue(tasks[0]['id'] < tasks[1]['id'])
+
+        for task in tasks:
+            self.assertTrue(task['state'] in [TASK_STATES['ASSIGNED'], TASK_STATES['OPEN']])
 
     def test_get_worker_tasks_check_wait(self):
         t_parent = Task.objects.create(