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']),
)
