Adds tests for task chains

adds several small improvements based on the problems that surfaced during writing the tests

django_q/cluster.py

@@ -328,7 +328,7 @@ def monitor(result_queue, broker=None):
         if task.get('cached', False):
             save_cached(task, broker)
         else:
-            save_task(task)
+            save_task(task, broker)
         # log the result
         if task['success']:
             logger.info(_("Processed [{}]").format(task['name']))
@@ -386,7 +386,7 @@ def worker(task_queue, result_queue, timer, timeout=Conf.TIMEOUT):
     logger.info(_('{} stopped doing work').format(name))
 
 
-def save_task(task):
+def save_task(task, broker):
     """
     Saves the task package to Django or the cache
     """
@@ -395,7 +395,7 @@ def save_task(task):
         return
     # async next in a chain
     if task.get('chain', None):
-        tasks.async_chain(task['chain'], group=task['group'], cached=task['cached'])
+        tasks.async_chain(task['chain'], group=task['group'], cached=task['cached'], sync=task['sync'], broker=broker)
     # SAVE LIMIT > 0: Prune database, SAVE_LIMIT 0: No pruning
     db.close_old_connections()
     try:
@@ -430,7 +430,7 @@ def save_cached(task, broker):
             group_key = '{}:{}:keys'.format(broker.list_key, group)
             group_list = broker.cache.get(group_key) or []
             # if it's an iter group, check if we are ready
-            if iter_count and len(group_list) == iter_count-1:
+            if iter_count and len(group_list) == iter_count - 1:
                 group_args = '{}:{}:args'.format(broker.list_key, group)
                 # collate the results into a Task result
                 results = [signing.SignedPackage.loads(broker.cache.get(k))['result'] for k in group_list]
@@ -444,7 +444,7 @@ def save_cached(task, broker):
                     task['cached'] = task.pop('iter_cached', None)
                     save_cached(task, broker=broker)
                 else:
-                    save_task(task)
+                    save_task(task, broker)
                 broker.cache.delete_many(group_list)
                 broker.cache.delete_many([group_key, group_args])
                 return
@@ -453,7 +453,7 @@ def save_cached(task, broker):
             broker.cache.set(group_key, group_list)
             # async next in a chain
             if task.get('chain', None):
-                tasks.async_chain(task['chain'], group=group, cached=task['cached'])
+                tasks.async_chain(task['chain'], group=group, cached=task['cached'], sync=task['sync'], broker=broker)
         # save the task
         broker.cache.set(task_key,
                          signing.SignedPackage.dumps(task),

django_q/tasks.py

@@ -19,11 +19,11 @@ def async(func, *args, **kwargs):
     # get options from q_options dict or direct from kwargs
     options = kwargs.pop('q_options', kwargs)
     broker = options.pop('broker', get_broker())
-    sync = options.pop('sync', False)
     # pop optionals
     opts = {'hook': None,
             'group': None,
             'save': None,
+            'sync': None,
             'cached': Conf.CACHED,
             'iter_count': None,
             'iter_cached': None,
@@ -44,7 +44,7 @@ def async(func, *args, **kwargs):
             task[key] = opts[key]
     # sign it
     pack = signing.SignedPackage.dumps(task)
-    if sync or Conf.SYNC:
+    if task.get('sync', False) or Conf.SYNC:
         return _sync(pack)
     # push it
     broker.enqueue(pack)
@@ -402,7 +402,7 @@ def async_iter(func, args_iter, **kwargs):
     return iter_group
 
 
-def async_chain(chain, group=None, cached=Conf.CACHED, sync=Conf.SYNC):
+def async_chain(chain, group=None, cached=Conf.CACHED, sync=Conf.SYNC, broker=None):
     """
     async a chain of tasks
     the chain must be in the format [(func,(args),{kwargs}),(func,(args),{kwargs})]
@@ -422,6 +422,7 @@ def async_chain(chain, group=None, cached=Conf.CACHED, sync=Conf.SYNC):
     kwargs['group'] = group
     kwargs['cached'] = cached
     kwargs['sync'] = sync
+    kwargs['broker'] = broker or get_broker()
     async(task[0], *args, **kwargs)
     return group
 
@@ -430,26 +431,34 @@ class Chain(object):
     """
     A sequential chain of tasks
     """
+
     def __init__(self, chain=None, group=None, cached=Conf.CACHED, sync=Conf.SYNC):
         self.chain = chain or []
         self.group = group or ''
+        self.broker = get_broker()
         self.cached = cached
         self.sync = sync
+        self.started = False
 
     def append(self, func, *args, **kwargs):
         """
         add a task to the chain
         takes the same parameters as async()
         """
-        task = (func, args, kwargs)
+        self.chain.append((func, args, kwargs))
-        self.chain.append(task)
+        # remove existing results
+        if self.started:
+            delete_group(self.group)
+            self.started = False
 
     def run(self):
         """
         Start queueing the chain to the worker cluster
         :return: the chain's group id
         """
-        self.group = async_chain(self.chain, group=self.group, cached=self.cached, sync=self.sync)
+        self.group = async_chain(chain=self.chain.copy(), group=self.group, cached=self.cached, sync=self.sync,
+                                 broker=self.broker)
+        self.started = True
         return self.group
 
     def result(self, wait=0):
@@ -458,7 +467,8 @@ class Chain(object):
         :param int wait: how many milliseconds to wait for a result
         :return: an unsorted list of results
         """
-        return result_group(self.group, wait=wait, count=len(self.chain), cached=self.cached)
+        if self.started:
+            return result_group(self.group, wait=wait, count=self.length(), cached=self.cached)
 
     def fetch(self, failures=True, wait=0):
         """
@@ -467,15 +477,25 @@ class Chain(object):
         :param int wait: how many milliseconds to wait for a result
         :return: an unsorted list of task objects
         """
-        return fetch_group(self.group, failures=failures, wait=wait, count=len(self.chain), cached=self.cached)
+        if self.started:
+            return fetch_group(self.group, failures=failures, wait=wait, count=self.length(), cached=self.cached)
 
     def current(self):
         """
         get the index of the currently executing chain element
         :return int: current chain index
         """
+        if not self.started:
+            return None
         return count_group(self.group, cached=self.cached)
 
+    def length(self):
+        """
+        get the length of the chain
+        :return int: length of the chain
+        """
+        return len(self.chain)
+
 
 def _sync(pack):
     """Simulate a package travelling through the cluster."""

django_q/tests/tasks.py

@@ -38,5 +38,9 @@ def get_user_id(user):
     return user.id
 
 
+def hello():
+    return 'hello'
+
+
 def result(obj):
     print('RESULT HOOK {} : {}'.format(obj.name, obj.result))

django_q/tests/test_cached.py

@@ -1,11 +1,11 @@
 from multiprocessing import Event, Queue, Value
 
 import pytest
-from django_q.cluster import pusher, worker, monitor
 
+from django_q.cluster import pusher, worker, monitor
 from django_q.conf import Conf
 from django_q.tasks import async, result, fetch, count_group, result_group, fetch_group, delete_group, delete_cached, \
-    async_iter
+    async_iter, Chain, async_chain
 from django_q.brokers import get_broker
 
 
@@ -96,9 +96,38 @@ def test_iter(broker):
     result_t = result(t)
     assert result_t is not None
     task_t = fetch(t)
-    assert task_t. __unicode__ is not None
     assert task_t.result == result_t
     assert result(t2) is not None
     assert result(t3) is not None
     assert result(t4)[0] == 1
     # test cached iter result
+
+
+@pytest.mark.django_db
+def test_chain(broker):
+    broker.purge_queue()
+    broker.cache.clear()
+    task_chain = Chain(sync=True)
+    task_chain.append('math.floor', 1)
+    task_chain.append('math.copysign', 1, -1)
+    task_chain.append('math.floor', 2)
+    assert task_chain.length() == 3
+    assert task_chain.current() is None
+    task_chain.run()
+    r = task_chain.result(wait=1000)
+    assert task_chain.current() == task_chain.length()
+    assert len(r) == task_chain.length()
+    t = task_chain.fetch()
+    assert len(t) == task_chain.length()
+    task_chain.cached = True
+    task_chain.append('math.floor', 3)
+    assert task_chain.length() == 4
+    task_chain.run()
+    r = task_chain.result(wait=1000)
+    assert task_chain.current() == task_chain.length()
+    assert len(r) == task_chain.length()
+    t = task_chain.fetch()
+    assert len(t) == task_chain.length()
+    # test single
+    rid = async_chain(['django_q.tests.tasks.hello', 'django_q.tests.tasks.hello'], sync=True, cached=True)
+    assert result_group(rid, cached=True) == ['hello', 'hello']