def test_work_once(self): class TaskType1(testutil.MockTask): work = mock.MagicMock() class TaskType2(testutil.MockTask): work = mock.MagicMock() cache_task = TaskType1( TaskPriority.CACHE, 'a', 'b', ) cleanup_task = TaskType2( TaskPriority.CLEANUP, 'a', 'b', ) queue = TaskQueue() queue.add_task(cache_task) queue.add_task(cleanup_task) log_list = [] attempted = queue.work_once(log_list=log_list) TaskType1.work.assert_called_with([cache_task]) self.assertEqual( TaskType2.work.call_count, 0, ) self.assertEqual( attempted, 1, ) self.assert_work_log(log_list, [ ('CACHE', [ ('TaskType1', 'a', 1) ]), ]) attempted = queue.work_once(log_list=log_list) self.assertEqual( attempted, 1, ) TaskType2.work.assert_called_with([cleanup_task]) self.assert_work_log(log_list, [ ('CACHE', [ ('TaskType1', 'a', 1) ]), ('CLEANUP', [ ('TaskType2', 'a', 1) ]), ])
def test_cache_lookup_promise(self): fake_cache = {} class TryCache(Task): priority = TaskPriority.CACHE def __init__(self, key): self.key = key super(TryCache, self).__init__() def __repr__(self): return "<TryCache %s>" % self.key @classmethod def work(cls, tasks): logging.debug("in TryCache.work with %r" % tasks) for task in tasks: task.resolve( fake_cache.get(task.key, CACHE_MISS) ) class LoadData(Task): def __init__(self, data): self.data = data super(LoadData, self).__init__() def __repr__(self): return "<LoadData %r>" % self.data @classmethod def work(cls, tasks): logging.debug("in LoadData.work with %r" % tasks) for task in tasks: task.resolve(task.data) class CachePopulate(Task): priority = TaskPriority.CLEANUP def __init__(self, key, value): self.key = key self.value = value super(CachePopulate, self).__init__() @classmethod def work(cls, tasks): logging.debug("in CachePopulate.work with %r" % tasks) for task in tasks: fake_cache[task.key] = task.value cache_key = "baz" data_value = 5 overall_promise = cache_lookup_promise( TryCache(cache_key), LoadData(data_value), cache_update_task_builder=lambda x: CachePopulate(cache_key, x) ) self.assertEqual( type(overall_promise), Promise, ) self.assertEqual( type(overall_promise.task), TryCache, ) callback = mock.MagicMock() overall_promise.then(callback) # Now let's actually do this thing. task_queue = TaskQueue() task_queue.add_task(overall_promise.task) log_list = [] task_queue.work(log_list=log_list) callback.assert_called_with(5) self.assert_work_log(log_list, [ ('CACHE', [ ('TryCache', (), 1), ]), ('SYNC_LOOKUP', [ ('LoadData', (), 1), ]), ('CLEANUP', [ ('CachePopulate', (), 1), ]), ])
def test_add_task(self): class TaskType1(testutil.MockTask): pass class TaskType2(testutil.MockTask): pass queue = TaskQueue() for task_type in (TaskType1, TaskType2): for priority in (TaskPriority.CACHE, TaskPriority.CLEANUP): for batch_key in ('a', 'b'): for coal_key in ('Q', 'Q', 'R'): mock_task = task_type( priority, batch_key, coal_key, ) queue.add_task(mock_task) self.assertEqual( set(queue.subqueues.keys()), set([1, 2, 3, 4]), ) self.assertEqual( len(queue.subqueues[1]), 4, ) self.assertEqual( len(queue.subqueues[2]), 0, ) self.assertEqual( len(queue.subqueues[3]), 0, ) self.assertEqual( len(queue.subqueues[4]), 4, ) self.assertEqual( set(queue.subqueues[1].keys()), set( [ (TaskType1, 'a'), (TaskType1, 'b'), (TaskType2, 'a'), (TaskType2, 'b'), ] ), ) self.assertEqual( set(queue.subqueues[4].keys()), set( [ (TaskType1, 'a'), (TaskType1, 'b'), (TaskType2, 'a'), (TaskType2, 'b'), ] ), ) for compound_key in ((TaskType1, 'a'), (TaskType2, 'a')): self.assertEqual( set(queue.subqueues[1][compound_key].keys()), set(['Q', 'R']), )