More refactoring

django_q/puller.py

@@ -16,10 +16,45 @@ from django_q.exceptions import TimeoutException, TimeoutHandler
 from django_q.process_manager import ProcessManager
 
 
+
 class Puller(ProcessManager):
     """The Puller is responsible for pulling the tasks from the broker, then return them to be picked up by the
     guard"""
 
+    @staticmethod
+    def get_tasks_from_broker(broker=None):
+        queued_tasks = []
+        logger.debug("pulling new tasks")
+        if broker is None:
+            broker = get_broker()
+        try:
+            task_set = broker.dequeue()
+        except Exception:
+            # broker probably crashed. Let the sentinel handle it.
+            raise ValueError("Failed to pull task from broker")
+        if task_set:
+            logger.info(
+                _("Found %(amount_tasks)s tasks") % {"amount_tasks": len(task_set)}
+            )
+            for task in task_set:
+                print(task)
+                logger.info("ONE TASK")
+                ack_id = task[0]
+                # unpack the task
+                try:
+                    queue_task = SignedPackage.loads(task[1])
+                except (TypeError, BadSignature):
+                    logger.exception("Failed to pull task from broker - bad task")
+                    broker.fail(ack_id)
+                    continue
+                queue_task.ack_id = ack_id
+                # send back to main process
+                queued_tasks.append(queue_task)
+            logger.debug(
+                _("queueing from %(list_key)s") % {"list_key": broker.list_key}
+            )
+        return queued_tasks
+
     def get_target(self):
         return self.run_puller
 
@@ -40,29 +75,11 @@ class Puller(ProcessManager):
                 logger.info("Stopping Puller")
                 break
             try:
-                task_set = broker.dequeue()
+                queued_tasks = Puller.get_tasks_from_broker(broker=broker)
             except Exception:
-                logger.exception("Failed to pull task from broker")
-                # broker probably crashed. Let the sentinel handle it.
+                logger.exception("Couldn't get items from broker")
                 sleep(10)
                 break
-            if task_set:
-                logger.info(
-                    _("Found %(amount_tasks)s tasks") % {"amount_tasks": len(task_set)}
-                )
-                for task in task_set:
-                    ack_id = task[0]
-                    # unpack the task
-                    try:
-                        queue_task = SignedPackage.loads(task[1])
-                    except (TypeError, BadSignature):
-                        logger.exception("Failed to pull task from broker - bad task")
-                        broker.fail(ack_id)
-                        continue
-                    queue_task.ack_id = ack_id
-                    # send back to main process
-                    pipe.send(queue_task)
-                logger.debug(
-                    _("queueing from %(list_key)s") % {"list_key": broker.list_key}
-                )
+            for queue_task in queued_tasks:
+                pipe.send(queue_task)
         logger.info(_("%(name)s stopped pushing tasks") % {"name": current_process().name})