From bbb347457b281b6ba1ad0f40d0f9dda515c8b865 Mon Sep 17 00:00:00 2001
From: Lubomir Gallovic <lgallovi@redhat.com>
Date: Thu, 27 Jun 2024 12:19:59 +0200
Subject: [PATCH] Add option for get_worker_tasks to also return ASSIGNED tasks

This behavior was recently changed in [1], where this function now only
returns OPEN tasks. This caused a regression in our service - which uses
a custom task manager. Add an parameter to restore the old behavior. The
parameter will be unused by default, so the default behavior remains
unchanged.

[1]: https://github.com/release-engineering/kobo/pull/251
---
 kobo/hub/xmlrpc/worker.py   | 13 ++++++++++---
 tests/test_xmlrpc_worker.py | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/kobo/hub/xmlrpc/worker.py b/kobo/hub/xmlrpc/worker.py
index 2205c53..abc20a0 100644
--- a/kobo/hub/xmlrpc/worker.py
+++ b/kobo/hub/xmlrpc/worker.py
@@ -56,10 +56,12 @@ 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 = []
@@ -67,7 +69,12 @@ def get_worker_tasks(request):
     # 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
diff --git a/tests/test_xmlrpc_worker.py b/tests/test_xmlrpc_worker.py
index b8df4fa..a26c994 100644
--- a/tests/test_xmlrpc_worker.py
+++ b/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(