From 2df29b908c07e3bb809b75ac8cb53a3fdafeb26f Mon Sep 17 00:00:00 2001 From: Ilan Steemers Date: Sat, 17 Oct 2015 18:49:45 +0200 Subject: [PATCH] Adds tests for task chains adds several small improvements based on the problems that surfaced during writing the tests --- django_q/cluster.py | 12 ++++++------ django_q/tasks.py | 36 +++++++++++++++++++++++++++-------- django_q/tests/tasks.py | 4 ++++ django_q/tests/test_cached.py | 35 +++++++++++++++++++++++++++++++--- 4 files changed, 70 insertions(+), 17 deletions(-) diff --git a/django_q/cluster.py b/django_q/cluster.py index d69a34a..b5e8d0c 100644 --- a/django_q/cluster.py +++ b/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), diff --git a/django_q/tasks.py b/django_q/tasks.py index 6c04f80..e572c7c 100644 --- a/django_q/tasks.py +++ b/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(task) + self.chain.append((func, args, kwargs)) + # 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.""" diff --git a/django_q/tests/tasks.py b/django_q/tests/tasks.py index 62e5674..27ad6bb 100644 --- a/django_q/tests/tasks.py +++ b/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)) diff --git a/django_q/tests/test_cached.py b/django_q/tests/test_cached.py index 2885eae..58bf11b 100644 --- a/django_q/tests/test_cached.py +++ b/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']