439 lines
13 KiB
Python
439 lines
13 KiB
Python
import datetime
|
|
import unittest
|
|
|
|
from huey import crontab
|
|
from huey import exceptions as huey_exceptions
|
|
from huey import Huey
|
|
from huey.api import QueueTask
|
|
from huey.backends.dummy import DummyDataStore
|
|
from huey.backends.dummy import DummyQueue
|
|
from huey.backends.dummy import DummySchedule
|
|
from huey.registry import registry
|
|
from huey.utils import EmptyData
|
|
from huey.utils import local_to_utc
|
|
|
|
|
|
queue_name = 'test-queue'
|
|
queue = DummyQueue(queue_name)
|
|
schedule = DummySchedule(queue_name)
|
|
huey = Huey(queue, schedule=schedule)
|
|
|
|
res_queue_name = 'test-queue-2'
|
|
res_queue = DummyQueue(res_queue_name)
|
|
res_store = DummyDataStore(res_queue_name)
|
|
|
|
res_huey = Huey(res_queue, res_store, schedule)
|
|
res_huey_nones = Huey(res_queue, res_store, store_none=True)
|
|
|
|
# store some global state
|
|
state = {}
|
|
last_executed_task_class = []
|
|
|
|
# create a decorated queue command
|
|
@huey.task()
|
|
def add(key, value):
|
|
state[key] = value
|
|
|
|
@huey.task(include_task=True)
|
|
def self_aware(key, value, task=None):
|
|
last_executed_task_class.append(task.__class__.__name__)
|
|
|
|
# create a periodic queue command
|
|
@huey.periodic_task(crontab(minute='0'))
|
|
def add_on_the_hour():
|
|
state['periodic'] = 'x'
|
|
|
|
# define a command using the class
|
|
class AddTask(QueueTask):
|
|
def execute(self):
|
|
k, v = self.data
|
|
state[k] = v
|
|
|
|
# create a command that raises an exception
|
|
class BampfException(Exception):
|
|
pass
|
|
|
|
@huey.task()
|
|
def throw_error():
|
|
raise BampfException('bampf')
|
|
|
|
@res_huey.task()
|
|
def add2(a, b):
|
|
return a + b
|
|
|
|
@res_huey.periodic_task(crontab(minute='0'))
|
|
def add_on_the_hour2():
|
|
state['periodic'] = 'x'
|
|
|
|
@res_huey.task()
|
|
def returns_none():
|
|
return None
|
|
|
|
@res_huey_nones.task()
|
|
def returns_none2():
|
|
return None
|
|
|
|
|
|
class HueyTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
global state
|
|
global last_executed_task_class
|
|
queue.flush()
|
|
res_queue.flush()
|
|
schedule.flush()
|
|
state = {}
|
|
last_executed_task_class = []
|
|
|
|
def test_registration(self):
|
|
self.assertTrue('queuecmd_add' in registry)
|
|
self.assertTrue('queuecmd_add_on_the_hour' in registry)
|
|
self.assertTrue('AddTask' in registry)
|
|
|
|
def test_enqueue(self):
|
|
# sanity check
|
|
self.assertEqual(len(queue), 0)
|
|
|
|
# initializing the command does not enqueue it
|
|
ac = AddTask(('k', 'v'))
|
|
self.assertEqual(len(queue), 0)
|
|
|
|
# ok, enqueue it, then check that it was enqueued
|
|
huey.enqueue(ac)
|
|
self.assertEqual(len(queue), 1)
|
|
|
|
# it can be enqueued multiple times
|
|
huey.enqueue(ac)
|
|
self.assertEqual(len(queue), 2)
|
|
|
|
# no changes to state
|
|
self.assertFalse('k' in state)
|
|
|
|
def test_enqueue_decorator(self):
|
|
# sanity check
|
|
self.assertEqual(len(queue), 0)
|
|
|
|
add('k', 'v')
|
|
self.assertEqual(len(queue), 1)
|
|
|
|
add('k', 'v')
|
|
self.assertEqual(len(queue), 2)
|
|
|
|
# no changes to state
|
|
self.assertFalse('k' in state)
|
|
|
|
def test_schedule(self):
|
|
dt = datetime.datetime(2011, 1, 1, 0, 1)
|
|
add('k', 'v')
|
|
self.assertEqual(len(queue), 1)
|
|
|
|
task = huey.dequeue()
|
|
self.assertEqual(task.execute_time, None)
|
|
|
|
add.schedule(args=('k2', 'v2'), eta=dt)
|
|
self.assertEqual(len(queue), 1)
|
|
task = huey.dequeue()
|
|
self.assertEqual(task.execute_time, local_to_utc(dt))
|
|
|
|
add.schedule(args=('k3', 'v3'), eta=dt, convert_utc=False)
|
|
self.assertEqual(len(queue), 1)
|
|
task = huey.dequeue()
|
|
self.assertEqual(task.execute_time, dt)
|
|
|
|
def test_error_raised(self):
|
|
throw_error()
|
|
|
|
# no error
|
|
task = huey.dequeue()
|
|
|
|
# error
|
|
self.assertRaises(BampfException, huey.execute, task)
|
|
|
|
def test_internal_error(self):
|
|
"""
|
|
Verify that exceptions are wrapped with the special "huey"
|
|
exception classes.
|
|
"""
|
|
class SpecialException(Exception):
|
|
pass
|
|
|
|
class BrokenQueue(DummyQueue):
|
|
def read(self):
|
|
raise SpecialException('read error')
|
|
|
|
def write(self, data):
|
|
raise SpecialException('write error')
|
|
|
|
class BrokenDataStore(DummyDataStore):
|
|
def get(self, key):
|
|
raise SpecialException('get error')
|
|
|
|
def put(self, key, value):
|
|
raise SpecialException('put error')
|
|
|
|
class BrokenSchedule(DummySchedule):
|
|
def add(self, data, ts):
|
|
raise SpecialException('add error')
|
|
|
|
def read(self, ts):
|
|
raise SpecialException('read error')
|
|
|
|
task = AddTask()
|
|
huey = Huey(
|
|
BrokenQueue('q'),
|
|
BrokenDataStore('q'),
|
|
BrokenSchedule('q'))
|
|
|
|
self.assertRaises(
|
|
huey_exceptions.QueueWriteException,
|
|
huey.enqueue,
|
|
AddTask())
|
|
self.assertRaises(
|
|
huey_exceptions.QueueReadException,
|
|
huey.dequeue)
|
|
self.assertRaises(
|
|
huey_exceptions.DataStorePutException,
|
|
huey.revoke,
|
|
task)
|
|
self.assertRaises(
|
|
huey_exceptions.DataStoreGetException,
|
|
huey.restore,
|
|
task)
|
|
self.assertRaises(
|
|
huey_exceptions.ScheduleAddException,
|
|
huey.add_schedule,
|
|
task)
|
|
self.assertRaises(
|
|
huey_exceptions.ScheduleReadException,
|
|
huey.read_schedule,
|
|
1)
|
|
|
|
def test_dequeueing(self):
|
|
res = huey.dequeue() # no error raised if queue is empty
|
|
self.assertEqual(res, None)
|
|
|
|
add('k', 'v')
|
|
task = huey.dequeue()
|
|
|
|
self.assertTrue(isinstance(task, QueueTask))
|
|
self.assertEqual(task.get_data(), (('k', 'v'), {}))
|
|
|
|
def test_execution(self):
|
|
self.assertFalse('k' in state)
|
|
add('k', 'v')
|
|
|
|
task = huey.dequeue()
|
|
self.assertFalse('k' in state)
|
|
|
|
huey.execute(task)
|
|
self.assertEqual(state['k'], 'v')
|
|
|
|
add('k', 'X')
|
|
self.assertEqual(state['k'], 'v')
|
|
|
|
huey.execute(huey.dequeue())
|
|
self.assertEqual(state['k'], 'X')
|
|
|
|
self.assertRaises(TypeError, huey.execute, huey.dequeue())
|
|
|
|
def test_self_awareness(self):
|
|
self_aware('k', 'v')
|
|
task = huey.dequeue()
|
|
huey.execute(task)
|
|
self.assertEqual(last_executed_task_class.pop(), "queuecmd_self_aware")
|
|
|
|
self_aware('k', 'v')
|
|
huey.execute(huey.dequeue())
|
|
self.assertEqual(last_executed_task_class.pop(), "queuecmd_self_aware")
|
|
|
|
add('k', 'x')
|
|
huey.execute(huey.dequeue())
|
|
self.assertEqual(len(last_executed_task_class), 0)
|
|
|
|
def test_call_local(self):
|
|
self.assertEqual(len(queue), 0)
|
|
self.assertEqual(state, {})
|
|
add.call_local('nugget', 'green')
|
|
|
|
self.assertEqual(len(queue), 0)
|
|
self.assertEqual(state['nugget'], 'green')
|
|
|
|
def test_revoke(self):
|
|
ac = AddTask(('k', 'v'))
|
|
ac2 = AddTask(('k2', 'v2'))
|
|
ac3 = AddTask(('k3', 'v3'))
|
|
|
|
res_huey.enqueue(ac)
|
|
res_huey.enqueue(ac2)
|
|
res_huey.enqueue(ac3)
|
|
res_huey.enqueue(ac2)
|
|
res_huey.enqueue(ac)
|
|
|
|
self.assertEqual(len(res_queue), 5)
|
|
res_huey.revoke(ac2)
|
|
|
|
while res_queue:
|
|
task = res_huey.dequeue()
|
|
if not res_huey.is_revoked(task):
|
|
res_huey.execute(task)
|
|
|
|
self.assertEqual(state, {'k': 'v', 'k3': 'v3'})
|
|
|
|
def test_revoke_periodic(self):
|
|
add_on_the_hour2.revoke()
|
|
self.assertTrue(add_on_the_hour2.is_revoked())
|
|
|
|
# it is still revoked
|
|
self.assertTrue(add_on_the_hour2.is_revoked())
|
|
|
|
add_on_the_hour2.restore()
|
|
self.assertFalse(add_on_the_hour2.is_revoked())
|
|
|
|
add_on_the_hour2.revoke(revoke_once=True)
|
|
self.assertTrue(add_on_the_hour2.is_revoked()) # it is revoked once, but we are preserving that state
|
|
self.assertTrue(add_on_the_hour2.is_revoked(peek=False)) # is revoked once, but clear state
|
|
self.assertFalse(add_on_the_hour2.is_revoked()) # no longer revoked
|
|
|
|
d = datetime.datetime
|
|
add_on_the_hour2.revoke(revoke_until=d(2011, 1, 1, 11, 0))
|
|
self.assertTrue(add_on_the_hour2.is_revoked(dt=d(2011, 1, 1, 10, 0)))
|
|
self.assertTrue(add_on_the_hour2.is_revoked(dt=d(2011, 1, 1, 10, 59)))
|
|
self.assertFalse(add_on_the_hour2.is_revoked(dt=d(2011, 1, 1, 11, 0)))
|
|
|
|
add_on_the_hour2.restore()
|
|
self.assertFalse(add_on_the_hour2.is_revoked())
|
|
|
|
def test_result_store(self):
|
|
res = add2(1, 2)
|
|
res2 = add2(4, 5)
|
|
res3 = add2(0, 0)
|
|
|
|
# none have been executed as yet
|
|
self.assertEqual(res.get(), None)
|
|
self.assertEqual(res2.get(), None)
|
|
self.assertEqual(res3.get(), None)
|
|
|
|
# execute the first task
|
|
res_huey.execute(res_huey.dequeue())
|
|
self.assertEqual(res.get(), 3)
|
|
self.assertEqual(res2.get(), None)
|
|
self.assertEqual(res3.get(), None)
|
|
|
|
# execute the second task
|
|
res_huey.execute(res_huey.dequeue())
|
|
self.assertEqual(res.get(), 3)
|
|
self.assertEqual(res2.get(), 9)
|
|
self.assertEqual(res3.get(), None)
|
|
|
|
# execute the 3rd, which returns a zero value
|
|
res_huey.execute(res_huey.dequeue())
|
|
self.assertEqual(res.get(), 3)
|
|
self.assertEqual(res2.get(), 9)
|
|
self.assertEqual(res3.get(), 0)
|
|
|
|
# check that it returns None when nothing is present
|
|
res = returns_none()
|
|
self.assertEqual(res.get(), None)
|
|
|
|
# execute, it will still return None, but underneath it is an EmptyResult
|
|
# indicating its actual result was not persisted
|
|
res_huey.execute(res_huey.dequeue())
|
|
self.assertEqual(res.get(), None)
|
|
self.assertEqual(res._result, EmptyData)
|
|
|
|
# execute again, this time note that we're pointing at the invoker
|
|
# that *does* accept None as a store-able result
|
|
res = returns_none2()
|
|
self.assertEqual(res.get(), None)
|
|
|
|
# it stores None
|
|
res_huey_nones.execute(res_huey_nones.dequeue())
|
|
self.assertEqual(res.get(), None)
|
|
self.assertEqual(res._result, None)
|
|
|
|
def test_task_store(self):
|
|
dt1 = datetime.datetime(2011, 1, 1, 0, 0)
|
|
dt2 = datetime.datetime(2035, 1, 1, 0, 0)
|
|
|
|
add2.schedule(args=('k', 'v'), eta=dt1, convert_utc=False)
|
|
task1 = res_huey.dequeue()
|
|
|
|
add2.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)
|
|
task2 = res_huey.dequeue()
|
|
|
|
add2('k3', 'v3')
|
|
task3 = res_huey.dequeue()
|
|
|
|
# add the command to the schedule
|
|
res_huey.add_schedule(task1)
|
|
self.assertEqual(len(res_huey.schedule._schedule), 1)
|
|
|
|
# add a future-dated command
|
|
res_huey.add_schedule(task2)
|
|
self.assertEqual(len(res_huey.schedule._schedule), 2)
|
|
|
|
res_huey.add_schedule(task3)
|
|
|
|
tasks = res_huey.read_schedule(dt1)
|
|
self.assertEqual(tasks, [task3, task1])
|
|
|
|
tasks = res_huey.read_schedule(dt1)
|
|
self.assertEqual(tasks, [])
|
|
|
|
tasks = res_huey.read_schedule(dt2)
|
|
self.assertEqual(tasks, [task2])
|
|
|
|
def test_ready_to_run_method(self):
|
|
dt1 = datetime.datetime(2011, 1, 1, 0, 0)
|
|
dt2 = datetime.datetime(2035, 1, 1, 0, 0)
|
|
|
|
add2.schedule(args=('k', 'v'), eta=dt1)
|
|
task1 = res_huey.dequeue()
|
|
|
|
add2.schedule(args=('k2', 'v2'), eta=dt2)
|
|
task2 = res_huey.dequeue()
|
|
|
|
add2('k3', 'v3')
|
|
task3 = res_huey.dequeue()
|
|
|
|
add2.schedule(args=('k4', 'v4'), task_id='test_task_id')
|
|
task4 = res_huey.dequeue()
|
|
|
|
# sanity check what should be run
|
|
self.assertTrue(res_huey.ready_to_run(task1))
|
|
self.assertFalse(res_huey.ready_to_run(task2))
|
|
self.assertTrue(res_huey.ready_to_run(task3))
|
|
self.assertTrue(res_huey.ready_to_run(task4))
|
|
self.assertEqual('test_task_id', task4.task_id)
|
|
|
|
def test_task_delay(self):
|
|
curr = datetime.datetime.utcnow()
|
|
curr50 = curr + datetime.timedelta(seconds=50)
|
|
curr70 = curr + datetime.timedelta(seconds=70)
|
|
|
|
add2.schedule(args=('k', 'v'), delay=60)
|
|
task1 = res_huey.dequeue()
|
|
|
|
add2.schedule(args=('k2', 'v2'), delay=600)
|
|
task2 = res_huey.dequeue()
|
|
|
|
add2('k3', 'v3')
|
|
task3 = res_huey.dequeue()
|
|
|
|
# add the command to the schedule
|
|
res_huey.add_schedule(task1)
|
|
res_huey.add_schedule(task2)
|
|
res_huey.add_schedule(task3)
|
|
|
|
# sanity check what should be run
|
|
self.assertFalse(res_huey.ready_to_run(task1))
|
|
self.assertFalse(res_huey.ready_to_run(task2))
|
|
self.assertTrue(res_huey.ready_to_run(task3))
|
|
|
|
self.assertFalse(res_huey.ready_to_run(task1, curr50))
|
|
self.assertFalse(res_huey.ready_to_run(task2, curr50))
|
|
self.assertTrue(res_huey.ready_to_run(task3, curr50))
|
|
|
|
self.assertTrue(res_huey.ready_to_run(task1, curr70))
|
|
self.assertFalse(res_huey.ready_to_run(task2, curr70))
|
|
self.assertTrue(res_huey.ready_to_run(task3, curr70))
|