def setUp(self):
global state
state = {}
self.orig_sleep = time.sleep
time.sleep = lambda x: None
self.consumer = Consumer(test_invoker, DummyConfiguration)
self.handler = TestLogHandler()
self.consumer.logger.addHandler(self.handler)
def handle(self, *args, **options):
from huey.djhuey import HUEY
try:
consumer_options = settings.HUEY['consumer_options']
except:
consumer_options = {}
if options['workers'] is not None:
consumer_options['workers'] = options['workers']
if options['periodic'] is not None:
consumer_options['periodic'] = options['periodic']
self.autodiscover()
consumer = Consumer(HUEY, **consumer_options)
consumer.run()
def handle(self, *args, **options):
config = load_config('huey.djhuey.conf.Configuration')
if options['threads'] is not None:
config.THREADS = options['threads']
if options['periodic'] is not None:
config.PERIODIC = options['periodic']
self.autodiscover()
queue = config.QUEUE
result_store = config.RESULT_STORE
task_store = config.TASK_STORE
invoker = Invoker(queue, result_store, task_store)
consumer = Consumer(invoker, config)
consumer.run()
def handle(self, *args, **options):
config = load_config('huey.djhuey.conf.Configuration')
if options['threads'] is not None:
config.THREADS = options['threads']
if options['periodic'] is not None:
config.PERIODIC = options['periodic']
self.autodiscover()
queue = config.QUEUE
result_store = config.RESULT_STORE
task_store = config.TASK_STORE
invoker = Invoker(queue, result_store, task_store)
consumer = Consumer(invoker, config)
consumer.run()
def setUp(self):
global state
state = {}
self.orig_sleep = time.sleep
time.sleep = lambda x: None
self.consumer = Consumer(test_invoker, DummyConfiguration)
self.handler = TestLogHandler()
self.consumer.logger.addHandler(self.handler)
def handle(self, *args, **options):
from huey.djhuey import HUEY
try:
consumer_options = settings.HUEY["consumer_options"]
except:
consumer_options = {}
if options["workers"] is not None:
consumer_options["workers"] = options["workers"]
if options["periodic"] is not None:
consumer_options["periodic"] = options["periodic"]
if options["initial_delay"] is not None:
consumer_options["initial_delay"] = options["initial_delay"]
if options["max_delay"] is not None:
consumer_options["max_delay"] = options["max_delay"]
self.autodiscover()
consumer = Consumer(HUEY, **consumer_options)
consumer.run()
def setUp(self):
global state
state = {}
self.orig_pc = registry._periodic_commands
registry._periodic_commands = [every_hour.command_class()]
self.orig_sleep = time.sleep
time.sleep = lambda x: None
self.consumer = Consumer(test_invoker, DummyConfiguration)
self.consumer.invoker.queue._queue = []
self.consumer.invoker.result_store._results = {}
self.consumer.schedule._schedule = {}
self.handler = TestLogHandler()
self.consumer.logger.addHandler(self.handler)
def setUp(self):
global state
state = {}
self.orig_pc = registry._periodic_tasks
registry._periodic_commands = [every_hour.task_class()]
self.orig_sleep = time.sleep
time.sleep = lambda x: None
test_huey.queue.flush()
test_huey.result_store.flush()
test_huey.schedule.flush()
self.consumer = Consumer(test_huey, workers=2)
self.consumer.create_threads()
self.handler = TestLogHandler()
logger.addHandler(self.handler)
class SkewConsumerTestCase(unittest.TestCase):
def setUp(self):
global state
state = {}
self.orig_sleep = time.sleep
time.sleep = lambda x: None
self.consumer = Consumer(test_invoker, DummyConfiguration)
self.handler = TestLogHandler()
self.consumer.logger.addHandler(self.handler)
def tearDown(self):
self.consumer.shutdown()
self.consumer.logger.removeHandler(self.handler)
time.sleep = self.orig_sleep
def test_consumer_loader(self):
config = load_config('huey.tests.config.Config')
self.assertTrue(isinstance(config.QUEUE, DummyQueue))
self.assertEqual(config.QUEUE.name, 'test-queue')
def spawn(self, func, *args, **kwargs):
t = threading.Thread(target=func, args=args, kwargs=kwargs)
t.start()
return t
def test_iterable_queue(self):
store = []
q = IterableQueue()
def do_queue(queue, result):
for message in queue:
result.append(message)
t = self.spawn(do_queue, q, store)
q.put(1)
q.put(2)
q.put(StopIteration)
t.join()
self.assertFalse(t.is_alive())
self.assertEqual(store, [1, 2])
def test_message_processing(self):
self.consumer.start_message_receiver()
self.consumer.start_worker_pool()
self.assertFalse('k' in state)
res = modify_state('k', 'v')
res.get(blocking=True)
self.assertTrue('k' in state)
self.assertEqual(res.get(), 'v')
def test_worker(self):
res = modify_state('x', 'y')
cmd = test_invoker.dequeue()
self.assertEqual(res.get(), None)
# we will be calling release() after finishing work
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertTrue('x' in state)
self.assertEqual(res.get(), 'y')
def test_worker_exception(self):
res = blow_up()
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
])
def test_retries_and_logging(self):
# this will continually fail
res = retry_command('blampf')
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 2)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages[-2:], [
'unhandled exception in worker thread',
're-enqueueing task %s, 1 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 1)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages[-2:], [
'unhandled exception in worker thread',
're-enqueueing task %s, 0 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 0)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(len(self.handler.messages), 7)
self.assertEqual(self.handler.messages[-1:], [
'unhandled exception in worker thread',
])
self.assertEqual(test_invoker.dequeue(), None)
def test_retries_with_success(self):
# this will fail once, then succeed
res = retry_command('blampf', False)
self.assertFalse('blampf' in state)
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 2)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(state['blampf'], 'fixed')
self.assertEqual(test_invoker.dequeue(), None)
def test_pooling(self):
# simulate acquiring two worker threads
self.consumer._pool.acquire()
self.consumer._pool.acquire()
res = modify_state('x', 'y')
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
# our result is not available since all workers are blocked
self.assertEqual(res.get(), None)
self.assertFalse(self.consumer._pool.acquire(blocking=False))
# our processor is waiting
self.assertTrue(pt.is_alive())
self.assertEqual(self.consumer._queue.qsize(), 0)
# release a worker
self.consumer._pool.release()
# we can get and block now, but will set a timeout of 3 to indicate that
# something is wrong
self.assertEqual(res.get(blocking=True, timeout=3), 'y')
# this is done
pt.join()
def test_scheduling(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
pt.join()
self.assertTrue('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# it got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertTrue(r2.task_id in self.consumer.schedule._schedule)
# run through an iteration of the scheduler
self.consumer.check_schedule(dt)
# our command was not enqueued
self.assertEqual(len(self.consumer.invoker.queue), 0)
# try running the scheduler with the time the command should run
self.consumer.check_schedule(dt2)
# it was enqueued
self.assertEqual(len(self.consumer.invoker.queue), 1)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue and inspect -- it won't be executed because the scheduler will
# see that it is scheduled to run in the future and plop it back into the
# schedule
command = self.consumer.invoker.dequeue()
self.assertEqual(command.task_id, r2.task_id)
self.assertEqual(command.execute_time, dt2)
def test_retry_scheduling(self):
# this will continually fail
res = retry_command_slow('blampf')
self.assertEqual(self.consumer.schedule._schedule, {})
cur_time = datetime.datetime.utcnow()
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
self.assertEqual(self.consumer.schedule._schedule, {
cmd.task_id: cmd,
})
cmd_from_sched = self.consumer.schedule._schedule[cmd.task_id]
self.assertEqual(cmd_from_sched.retries, 2)
exec_time = cmd.execute_time
self.assertEqual((exec_time - cur_time).seconds, 10)
def test_schedule_local_utc(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2)
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
pt.join()
self.assertTrue('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# it got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertTrue(r2.task_id in self.consumer.schedule._schedule)
# run through an iteration of the scheduler
self.consumer.check_schedule(dt)
# our command was not enqueued
self.assertEqual(len(self.consumer.invoker.queue), 0)
# try running the scheduler with the time the command should run
self.consumer.check_schedule(local_to_utc(dt2))
# it was enqueued
self.assertEqual(len(self.consumer.invoker.queue), 1)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue and inspect -- it won't be executed because the scheduler will
# see that it is scheduled to run in the future and plop it back into the
# schedule
command = self.consumer.invoker.dequeue()
self.assertEqual(command.task_id, r2.task_id)
self.assertEqual(command.execute_time, local_to_utc(dt2))
def test_schedule_persistence(self):
dt = datetime.datetime(2037, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 1)
r = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)
# two messages in the queue
self.assertEqual(len(self.consumer.invoker.queue), 2)
# pull 'em down
self.consumer.check_message()
self.consumer.check_message()
self.consumer.save_schedule()
self.consumer.schedule._schedule = {}
self.consumer.load_schedule()
self.assertTrue(r.task_id in self.consumer.schedule._schedule)
self.assertTrue(r2.task_id in self.consumer.schedule._schedule)
cmd1 = self.consumer.schedule._schedule[r.task_id]
cmd2 = self.consumer.schedule._schedule[r2.task_id]
self.assertEqual(cmd1.execute_time, dt)
self.assertEqual(cmd2.execute_time, dt2)
# check w/conversion
r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt)
self.consumer.check_message()
self.consumer.save_schedule()
self.consumer.schedule._schedule = {}
self.consumer.load_schedule()
cmd3 = self.consumer.schedule._schedule[r3.task_id]
self.assertEqual(cmd3.execute_time, local_to_utc(dt))
class Configuration(BaseConfiguration):
QUEUE = queue
RESULT_STORE = result_store
LOGLEVEL = 'DEBUG'
THREAD_WORKER = False
@queue_command(invoker)
def search_yahoo(query):
url = 'http://www.yahoo.com/?q=%s' % query
resp = requests.get(url)
return resp.text
@queue_command(invoker)
def search_google(query):
url = 'http://www.google.com/search?q=%s' % query
resp = requests.get(url)
return resp.text
if __name__ == '__main__':
logging.basicConfig(level=logging.DEBUG)
try:
consumer = Consumer(invoker, Configuration)
consumer.run()
except KeyboardInterrupt:
pass
class ConsumerTestCase(unittest.TestCase):
def setUp(self):
global state
state = {}
self.orig_pc = registry._periodic_tasks
registry._periodic_commands = [every_hour.task_class()]
self.orig_sleep = time.sleep
time.sleep = lambda x: None
test_huey.queue.flush()
test_huey.result_store.flush()
test_huey.schedule.flush()
test_events._events = deque()
self.consumer = Consumer(test_huey, workers=2)
self.consumer.create_threads()
self.handler = TestLogHandler()
logger.addHandler(self.handler)
def tearDown(self):
self.consumer.shutdown()
logger.removeHandler(self.handler)
registry._periodic_tasks = self.orig_pc
time.sleep = self.orig_sleep
def assertStatusTask(self, status_task):
parsed = []
i = 0
while i < len(status_task):
event = json.loads(test_events._events[i])
status, task, extra = status_task[i]
self.assertEqual(event['status'], status)
self.assertEqual(event['id'], task.task_id)
for k, v in extra.items():
self.assertEqual(event[k], v)
i += 1
def spawn(self, func, *args, **kwargs):
t = threading.Thread(target=func, args=args, kwargs=kwargs)
t.start()
return t
def run_worker(self, task, ts=None):
worker_t = WorkerThread(
test_huey,
self.consumer.default_delay,
self.consumer.max_delay,
self.consumer.backoff,
self.consumer.utc,
self.consumer._shutdown)
ts = ts or datetime.datetime.utcnow()
worker_t.handle_task(task, ts)
def test_message_processing(self):
self.consumer.worker_threads[0].start()
self.assertFalse('k' in state)
res = modify_state('k', 'v')
res.get(blocking=True)
self.assertTrue('k' in state)
self.assertEqual(res.get(), 'v')
self.assertEqual(len(test_events._events), 2)
self.assertStatusTask([
('finished', res.task, {}),
('started', res.task, {}),
])
def test_worker(self):
modify_state('k', 'w')
task = test_huey.dequeue()
self.run_worker(task)
self.assertEqual(state, {'k': 'w'})
def test_worker_exception(self):
blow_up()
task = test_huey.dequeue()
self.run_worker(task)
self.assertTrue(
'Unhandled exception in worker thread' in self.handler.messages)
self.assertEqual(len(test_events._events), 2)
self.assertStatusTask([
('error', task, {'error': True}),
('started', task, {}),
])
def test_retries_and_logging(self):
# this will continually fail
retry_command('blampf')
for i in reversed(range(4)):
task = test_huey.dequeue()
self.assertEqual(task.retries, i)
self.run_worker(task)
if i > 0:
self.assertEqual(
self.handler.messages[-1],
'Re-enqueueing task %s, %s tries left' % (
task.task_id, i - 1))
self.assertStatusTask([
('enqueued', task, {}),
('retrying', task, {}),
('error', task,{}),
('started', task, {}),
])
last_idx = -2
else:
self.assertStatusTask([
('error', task,{}),
('started', task, {}),
])
last_idx = -1
self.assertEqual(self.handler.messages[last_idx],
'Unhandled exception in worker thread')
self.assertEqual(test_huey.dequeue(), None)
def test_retries_with_success(self):
# this will fail once, then succeed
retry_command('blampf', False)
self.assertFalse('blampf' in state)
task = test_huey.dequeue()
self.run_worker(task)
self.assertEqual(self.handler.messages, [
'Executing %s' % task,
'Unhandled exception in worker thread',
'Re-enqueueing task %s, 2 tries left' % task.task_id])
task = test_huey.dequeue()
self.assertEqual(task.retries, 2)
self.run_worker(task)
self.assertEqual(state['blampf'], 'fixed')
self.assertEqual(test_huey.dequeue(), None)
self.assertStatusTask([
('finished', task, {}),
('started', task, {}),
('enqueued', task, {'retries': 2}),
('retrying', task, {'retries': 3}),
('error', task, {'error': True}),
('started', task, {}),
])
def test_scheduling(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
ad1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
ad2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)
# dequeue the past-timestamped task and run it.
worker = self.consumer.worker_threads[0]
worker.check_message()
self.assertTrue('k' in state)
# dequeue the future-timestamped task.
worker.check_message()
# verify the task got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertStatusTask([
('scheduled', ad2.task, {}),
('finished', ad1.task, {}),
('started', ad1.task, {}),
])
# run through an iteration of the scheduler
self.consumer.scheduler_t.loop(dt)
# our command was not enqueued and no events were emitted.
self.assertEqual(len(test_queue._queue), 0)
self.assertEqual(len(test_events._events), 3)
# run through an iteration of the scheduler
self.consumer.scheduler_t.loop(dt2)
# our command was enqueued
self.assertEqual(len(test_queue._queue), 1)
self.assertEqual(len(test_events._events), 4)
self.assertStatusTask([
('enqueued', ad2.task, {}),
])
def test_retry_scheduling(self):
# this will continually fail
retry_command_slow('blampf')
cur_time = datetime.datetime.utcnow()
task = test_huey.dequeue()
self.run_worker(task, ts=cur_time)
self.assertEqual(self.handler.messages, [
'Executing %s' % task,
'Unhandled exception in worker thread',
'Re-enqueueing task %s, 2 tries left' % task.task_id,
])
in_11 = cur_time + datetime.timedelta(seconds=11)
tasks_from_sched = test_huey.read_schedule(in_11)
self.assertEqual(tasks_from_sched, [task])
task = tasks_from_sched[0]
self.assertEqual(task.retries, 2)
exec_time = task.execute_time
self.assertEqual((exec_time - cur_time).seconds, 10)
self.assertStatusTask([
('scheduled', task, {
'retries': 2,
'retry_delay': 10,
'execute_time': time.mktime(exec_time.timetuple())}),
('retrying', task, {
'retries': 3,
'retry_delay': 10,
'execute_time': None}),
('error', task, {}),
('started', task, {}),
])
def test_revoking_normal(self):
# enqueue 2 normal commands
r1 = modify_state('k', 'v')
r2 = modify_state('k2', 'v2')
# revoke the first *before it has been checked*
r1.revoke()
self.assertTrue(test_huey.is_revoked(r1.task))
self.assertFalse(test_huey.is_revoked(r2.task))
# dequeue a *single* message (r1)
task = test_huey.dequeue()
self.run_worker(task)
self.assertEqual(len(test_events._events), 1)
self.assertStatusTask([
('revoked', r1.task, {}),
])
# no changes and the task was not added to the schedule
self.assertFalse('k' in state)
# dequeue a *single* message
task = test_huey.dequeue()
self.run_worker(task)
self.assertTrue('k2' in state)
def test_revoking_schedule(self):
global state
dt = datetime.datetime(2011, 1, 1)
dt2 = datetime.datetime(2037, 1, 1)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt, convert_utc=False)
r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt2, convert_utc=False)
r4 = modify_state.schedule(args=('k4', 'v4'), eta=dt2, convert_utc=False)
# revoke r1 and r3
r1.revoke()
r3.revoke()
self.assertTrue(test_huey.is_revoked(r1.task))
self.assertFalse(test_huey.is_revoked(r2.task))
self.assertTrue(test_huey.is_revoked(r3.task))
self.assertFalse(test_huey.is_revoked(r4.task))
expected = [
#state, schedule
({}, 0),
({'k2': 'v2'}, 0),
({'k2': 'v2'}, 1),
({'k2': 'v2'}, 2),
]
for i in range(4):
estate, esc = expected[i]
# dequeue a *single* message
task = test_huey.dequeue()
self.run_worker(task)
self.assertEqual(state, estate)
self.assertEqual(len(test_huey.schedule._schedule), esc)
# lets pretend its 2037
future = dt2 + datetime.timedelta(seconds=1)
self.consumer.scheduler_t.loop(future)
self.assertEqual(len(test_huey.schedule._schedule), 0)
# There are two tasks in the queue now (r3 and r4) -- process both.
for i in range(2):
task = test_huey.dequeue()
self.run_worker(task, future)
self.assertEqual(state, {'k2': 'v2', 'k4': 'v4'})
def test_revoking_periodic(self):
global state
def loop_periodic(ts):
self.consumer.periodic_t.loop(ts)
for i in range(len(test_queue._queue)):
task = test_huey.dequeue()
self.run_worker(task, ts)
# revoke the command once
every_hour.revoke(revoke_once=True)
self.assertTrue(every_hour.is_revoked())
# it will be skipped the first go-round
dt = datetime.datetime(2011, 1, 1, 0, 0)
loop_periodic(dt)
# it has not been run
self.assertEqual(state, {})
# the next go-round it will be enqueued
loop_periodic(dt)
# our command was run
self.assertEqual(state, {'p': 'y'})
# reset state
state = {}
# revoke the command
every_hour.revoke()
self.assertTrue(every_hour.is_revoked())
# it will no longer be enqueued
loop_periodic(dt)
loop_periodic(dt)
self.assertEqual(state, {})
# restore
every_hour.restore()
self.assertFalse(every_hour.is_revoked())
# it will now be enqueued
loop_periodic(dt)
self.assertEqual(state, {'p': 'y'})
# reset
state = {}
# revoke for an hour
td = datetime.timedelta(seconds=3600)
every_hour.revoke(revoke_until=dt + td)
loop_periodic(dt)
self.assertEqual(state, {})
# after an hour it is back
loop_periodic(dt + td)
self.assertEqual(state, {'p': 'y'})
# our data store should reflect the delay
task_obj = every_hour.task_class()
self.assertEqual(len(test_huey.result_store._results), 1)
self.assertTrue(task_obj.revoke_id in test_huey.result_store._results)
class SkewConsumerTestCase(unittest.TestCase):
def setUp(self):
global state
state = {}
self.orig_sleep = time.sleep
time.sleep = lambda x: None
self.consumer = Consumer(test_invoker, DummyConfiguration)
self.handler = TestLogHandler()
self.consumer.logger.addHandler(self.handler)
def tearDown(self):
self.consumer.shutdown()
self.consumer.logger.removeHandler(self.handler)
time.sleep = self.orig_sleep
def test_consumer_loader(self):
config = load_config('huey.tests.config.Config')
self.assertTrue(isinstance(config.QUEUE, DummyQueue))
self.assertEqual(config.QUEUE.name, 'test-queue')
def spawn(self, func, *args, **kwargs):
t = threading.Thread(target=func, args=args, kwargs=kwargs)
t.start()
return t
def test_iterable_queue(self):
store = []
q = IterableQueue()
def do_queue(queue, result):
for message in queue:
result.append(message)
t = self.spawn(do_queue, q, store)
q.put(1)
q.put(2)
q.put(StopIteration)
t.join()
self.assertFalse(t.is_alive())
self.assertEqual(store, [1, 2])
def test_message_processing(self):
self.consumer.start_message_receiver()
self.consumer.start_worker_pool()
self.assertFalse('k' in state)
res = modify_state('k', 'v')
res.get(blocking=True)
self.assertTrue('k' in state)
self.assertEqual(res.get(), 'v')
def test_worker(self):
res = modify_state('x', 'y')
cmd = test_invoker.dequeue()
self.assertEqual(res.get(), None)
# we will be calling release() after finishing work
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertTrue('x' in state)
self.assertEqual(res.get(), 'y')
def test_worker_exception(self):
res = blow_up()
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
])
def test_retries_and_logging(self):
# this will continually fail
res = retry_command('blampf')
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 2)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages[-2:], [
'unhandled exception in worker thread',
're-enqueueing task %s, 1 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 1)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages[-2:], [
'unhandled exception in worker thread',
're-enqueueing task %s, 0 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 0)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(len(self.handler.messages), 7)
self.assertEqual(self.handler.messages[-1:], [
'unhandled exception in worker thread',
])
self.assertEqual(test_invoker.dequeue(), None)
def test_retries_with_success(self):
# this will fail once, then succeed
res = retry_command('blampf', False)
self.assertFalse('blampf' in state)
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 2)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(state['blampf'], 'fixed')
self.assertEqual(test_invoker.dequeue(), None)
def test_pooling(self):
# simulate acquiring two worker threads
self.consumer._pool.acquire()
self.consumer._pool.acquire()
res = modify_state('x', 'y')
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
# our result is not available since all workers are blocked
self.assertEqual(res.get(), None)
self.assertFalse(self.consumer._pool.acquire(blocking=False))
# our processor is waiting
self.assertTrue(pt.is_alive())
self.assertEqual(self.consumer._queue.qsize(), 0)
# release a worker
self.consumer._pool.release()
# we can get and block now, but will set a timeout of 3 to indicate that
# something is wrong
self.assertEqual(res.get(blocking=True, timeout=3), 'y')
# this is done
pt.join()
def test_scheduling(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'),
eta=dt2,
convert_utc=False)
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
pt.join()
self.assertTrue('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# it got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertTrue(r2.task_id in self.consumer.schedule._schedule)
# run through an iteration of the scheduler
self.consumer.check_schedule(dt)
# our command was not enqueued
self.assertEqual(len(self.consumer.invoker.queue), 0)
# try running the scheduler with the time the command should run
self.consumer.check_schedule(dt2)
# it was enqueued
self.assertEqual(len(self.consumer.invoker.queue), 1)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue and inspect -- it won't be executed because the scheduler will
# see that it is scheduled to run in the future and plop it back into the
# schedule
command = self.consumer.invoker.dequeue()
self.assertEqual(command.task_id, r2.task_id)
self.assertEqual(command.execute_time, dt2)
def test_retry_scheduling(self):
# this will continually fail
res = retry_command_slow('blampf')
self.assertEqual(self.consumer.schedule._schedule, {})
cur_time = datetime.datetime.utcnow()
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
self.assertEqual(self.consumer.schedule._schedule, {
cmd.task_id: cmd,
})
cmd_from_sched = self.consumer.schedule._schedule[cmd.task_id]
self.assertEqual(cmd_from_sched.retries, 2)
exec_time = cmd.execute_time
self.assertEqual((exec_time - cur_time).seconds, 10)
def test_schedule_local_utc(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2)
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
pt.join()
self.assertTrue('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# it got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertTrue(r2.task_id in self.consumer.schedule._schedule)
# run through an iteration of the scheduler
self.consumer.check_schedule(dt)
# our command was not enqueued
self.assertEqual(len(self.consumer.invoker.queue), 0)
# try running the scheduler with the time the command should run
self.consumer.check_schedule(local_to_utc(dt2))
# it was enqueued
self.assertEqual(len(self.consumer.invoker.queue), 1)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue and inspect -- it won't be executed because the scheduler will
# see that it is scheduled to run in the future and plop it back into the
# schedule
command = self.consumer.invoker.dequeue()
self.assertEqual(command.task_id, r2.task_id)
self.assertEqual(command.execute_time, local_to_utc(dt2))
def test_schedule_persistence(self):
dt = datetime.datetime(2037, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 1)
r = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'),
eta=dt2,
convert_utc=False)
# two messages in the queue
self.assertEqual(len(self.consumer.invoker.queue), 2)
# pull 'em down
self.consumer.check_message()
self.consumer.check_message()
self.consumer.save_schedule()
self.consumer.schedule._schedule = {}
self.consumer.load_schedule()
self.assertTrue(r.task_id in self.consumer.schedule._schedule)
self.assertTrue(r2.task_id in self.consumer.schedule._schedule)
cmd1 = self.consumer.schedule._schedule[r.task_id]
cmd2 = self.consumer.schedule._schedule[r2.task_id]
self.assertEqual(cmd1.execute_time, dt)
self.assertEqual(cmd2.execute_time, dt2)
# check w/conversion
r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt)
self.consumer.check_message()
self.consumer.save_schedule()
self.consumer.schedule._schedule = {}
self.consumer.load_schedule()
cmd3 = self.consumer.schedule._schedule[r3.task_id]
self.assertEqual(cmd3.execute_time, local_to_utc(dt))
class SkewConsumerTestCase(unittest.TestCase):
def setUp(self):
global state
state = {}
self.orig_pc = registry._periodic_commands
registry._periodic_commands = [every_hour.command_class()]
self.orig_sleep = time.sleep
time.sleep = lambda x: None
self.consumer = Consumer(test_invoker, DummyConfiguration)
self.consumer.invoker.queue._queue = []
self.consumer.invoker.result_store._results = {}
self.consumer.schedule._schedule = {}
self.handler = TestLogHandler()
self.consumer.logger.addHandler(self.handler)
def tearDown(self):
self.consumer.shutdown()
self.consumer.logger.removeHandler(self.handler)
registry._periodic_commands = self.orig_pc
time.sleep = self.orig_sleep
def test_consumer_loader(self):
config = load_config('huey.tests.config.Config')
self.assertTrue(isinstance(config.QUEUE, DummyQueue))
self.assertEqual(config.QUEUE.name, 'test-queue')
def spawn(self, func, *args, **kwargs):
t = threading.Thread(target=func, args=args, kwargs=kwargs)
t.start()
return t
def test_iterable_queue(self):
store = []
q = IterableQueue()
def do_queue(queue, result):
for message in queue:
result.append(message)
t = self.spawn(do_queue, q, store)
q.put(1)
q.put(2)
q.put(StopIteration)
t.join()
self.assertFalse(t.is_alive())
self.assertEqual(store, [1, 2])
def test_message_processing(self):
self.consumer.start_message_receiver()
self.consumer.start_worker_pool()
self.assertFalse('k' in state)
res = modify_state('k', 'v')
res.get(blocking=True)
self.assertTrue('k' in state)
self.assertEqual(res.get(), 'v')
def test_worker(self):
res = modify_state('x', 'y')
cmd = test_invoker.dequeue()
self.assertEqual(res.get(), None)
# we will be calling release() after finishing work
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertTrue('x' in state)
self.assertEqual(res.get(), 'y')
def test_worker_exception(self):
res = blow_up()
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
])
def test_retries_and_logging(self):
# this will continually fail
res = retry_command('blampf')
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 2)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages[-2:], [
'unhandled exception in worker thread',
're-enqueueing task %s, 1 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 1)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages[-2:], [
'unhandled exception in worker thread',
're-enqueueing task %s, 0 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 0)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(len(self.handler.messages), 7)
self.assertEqual(self.handler.messages[-1:], [
'unhandled exception in worker thread',
])
self.assertEqual(test_invoker.dequeue(), None)
def test_retries_with_success(self):
# this will fail once, then succeed
res = retry_command('blampf', False)
self.assertFalse('blampf' in state)
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
cmd = test_invoker.dequeue()
self.assertEqual(cmd.retries, 2)
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(state['blampf'], 'fixed')
self.assertEqual(test_invoker.dequeue(), None)
def test_pooling(self):
# simulate acquiring two worker threads
self.consumer._pool.acquire()
self.consumer._pool.acquire()
res = modify_state('x', 'y')
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
# our result is not available since all workers are blocked
self.assertEqual(res.get(), None)
self.assertFalse(self.consumer._pool.acquire(blocking=False))
# our processor is waiting
self.assertTrue(pt.is_alive())
self.assertEqual(self.consumer._queue.qsize(), 0)
# release a worker
self.consumer._pool.release()
# we can get and block now, but will set a timeout of 3 to indicate that
# something is wrong
self.assertEqual(res.get(blocking=True, timeout=3), 'y')
# this is done
pt.join()
def test_scheduling(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
pt.join()
self.assertTrue('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# it got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertTrue(r2.command.task_id in self.consumer.schedule._schedule)
# run through an iteration of the scheduler
self.consumer.check_schedule(dt)
# our command was not enqueued
self.assertEqual(len(self.consumer.invoker.queue), 0)
# try running the scheduler with the time the command should run
self.consumer.check_schedule(dt2)
# it was enqueued
self.assertEqual(len(self.consumer.invoker.queue), 1)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue and inspect -- it won't be executed because the scheduler will
# see that it is scheduled to run in the future and plop it back into the
# schedule
command = self.consumer.invoker.dequeue()
self.assertEqual(command.task_id, r2.command.task_id)
self.assertEqual(command.execute_time, dt2)
def test_retry_scheduling(self):
# this will continually fail
res = retry_command_slow('blampf')
self.assertEqual(self.consumer.schedule._schedule, {})
cur_time = datetime.datetime.utcnow()
cmd = test_invoker.dequeue()
self.consumer._pool.acquire()
self.consumer.worker(cmd)
self.assertEqual(self.handler.messages, [
'unhandled exception in worker thread',
're-enqueueing task %s, 2 tries left' % cmd.task_id,
])
self.assertEqual(self.consumer.schedule._schedule, {
cmd.task_id: cmd,
})
cmd_from_sched = self.consumer.schedule._schedule[cmd.task_id]
self.assertEqual(cmd_from_sched.retries, 2)
exec_time = cmd.execute_time
self.assertEqual((exec_time - cur_time).seconds, 10)
def test_schedule_local_utc(self):
dt = datetime.datetime(2011, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 0)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2)
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
pt.join()
self.assertTrue('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# it got stored in the schedule instead of executing
self.assertFalse('k2' in state)
self.assertTrue(r2.command.task_id in self.consumer.schedule._schedule)
# run through an iteration of the scheduler
self.consumer.check_schedule(dt)
# our command was not enqueued
self.assertEqual(len(self.consumer.invoker.queue), 0)
# try running the scheduler with the time the command should run
self.consumer.check_schedule(local_to_utc(dt2))
# it was enqueued
self.assertEqual(len(self.consumer.invoker.queue), 1)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue and inspect -- it won't be executed because the scheduler will
# see that it is scheduled to run in the future and plop it back into the
# schedule
command = self.consumer.invoker.dequeue()
self.assertEqual(command.task_id, r2.command.task_id)
self.assertEqual(command.execute_time, local_to_utc(dt2))
def test_schedule_persistence(self):
# should not error out as it will be EmptyData
self.consumer.load_schedule()
dt = datetime.datetime(2037, 1, 1, 0, 0)
dt2 = datetime.datetime(2037, 1, 1, 0, 1)
r = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)
# two messages in the queue
self.assertEqual(len(self.consumer.invoker.queue), 2)
# pull 'em down
self.consumer.check_message()
self.consumer.check_message()
self.consumer.save_schedule()
self.consumer.schedule._schedule = {}
self.consumer.load_schedule()
self.assertTrue(r.command.task_id in self.consumer.schedule._schedule)
self.assertTrue(r2.command.task_id in self.consumer.schedule._schedule)
cmd1 = self.consumer.schedule._schedule[r.command.task_id]
cmd2 = self.consumer.schedule._schedule[r2.command.task_id]
self.assertEqual(cmd1.execute_time, dt)
self.assertEqual(cmd2.execute_time, dt2)
# check w/conversion
r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt)
self.consumer.check_message()
self.consumer.save_schedule()
self.consumer.schedule._schedule = {}
self.consumer.load_schedule()
cmd3 = self.consumer.schedule._schedule[r3.command.task_id]
self.assertEqual(cmd3.execute_time, local_to_utc(dt))
def test_revoking_normal(self):
# enqueue 2 normal commands
r1 = modify_state('k', 'v')
r2 = modify_state('k2', 'v2')
# revoke the first *before it has been checked*
r1.revoke()
self.assertTrue(test_invoker.is_revoked(r1.command))
self.assertFalse(test_invoker.is_revoked(r2.command))
# dequeue a *single* message (r1)
pt = self.spawn(self.consumer.check_message)
# work on any messages generated by the processor thread
# it should throw out r1 since it is revoked and scheduled
# to run immediately
st = self.spawn(self.consumer.worker_pool)
pt.join()
# no changes and the task was not added to the schedule
self.assertFalse('k' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
self.assertTrue('k2' in state)
self.assertEqual(self.consumer.schedule._schedule, {})
def test_revoking_schedule(self):
global state
dt = datetime.datetime(2011, 1, 1)
dt2 = datetime.datetime(2037, 1, 1)
r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt, convert_utc=False)
r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt2, convert_utc=False)
r4 = modify_state.schedule(args=('k4', 'v4'), eta=dt2, convert_utc=False)
# revoke r1 and r3
r1.revoke()
r3.revoke()
self.assertTrue(test_invoker.is_revoked(r1.command))
self.assertFalse(test_invoker.is_revoked(r2.command))
self.assertTrue(test_invoker.is_revoked(r3.command))
self.assertFalse(test_invoker.is_revoked(r4.command))
expected = [
#state, schedule
({}, {}),
({'k2': 'v2'}, {}),
({'k2': 'v2'}, {r3.command.task_id: r3.command}),
({'k2': 'v2'}, {r3.command.task_id: r3.command, r4.command.task_id: r4.command}),
]
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
for i in range(4):
estate, esc = expected[i]
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
self.assertEqual(state, estate)
self.assertEqual(self.consumer.schedule._schedule, esc)
# lets pretend its 2037
self.consumer.check_schedule(dt2 + datetime.timedelta(seconds=1))
self.assertEqual(self.consumer.schedule._schedule, {})
inv = self.consumer.invoker
command = inv.dequeue()
sched = self.consumer.schedule
self.assertTrue(sched.should_run(command, dt2))
self.assertTrue(sched.can_run(command, dt2))
inv.execute(command)
self.assertEqual(state, {'k2': 'v2', 'k4': 'v4'})
def test_revoking_periodic(self):
global state
def assertQ(l):
self.assertEqual(len(self.consumer.invoker.queue._queue), l)
# grab a reference to the invoker
invoker = self.consumer.invoker
# work on any messages generated by the processor thread
st = self.spawn(self.consumer.worker_pool)
# revoke the command once
every_hour.revoke(revoke_once=True)
self.assertTrue(every_hour.is_revoked())
# it will be skipped the first go-round
dt = datetime.datetime(2011, 1, 1, 0, 0)
self.consumer.enqueue_periodic_commands(dt)
assertQ(0)
# it has not been run
self.assertEqual(state, {})
# the next go-round it will be enqueued
self.consumer.enqueue_periodic_commands(dt)
assertQ(1)
# it has still not been run
self.assertEqual(state, {})
# dequeue a *single* message
pt = self.spawn(self.consumer.check_message)
pt.join()
# our command was run
self.assertEqual(state, {'p': 'y'})
# reset state
state = {}
# revoke the command
every_hour.revoke()
self.assertTrue(every_hour.is_revoked())
# it will no longer be enqueued
self.consumer.enqueue_periodic_commands(dt)
assertQ(0)
self.consumer.enqueue_periodic_commands(dt)
assertQ(0)
# restore
every_hour.restore()
self.assertFalse(every_hour.is_revoked())
# it will now be enqueued
self.consumer.enqueue_periodic_commands(dt)
assertQ(1)
self.consumer.enqueue_periodic_commands(dt)
assertQ(2)
# reset
state = {}
invoker.queue._queue = []
# revoke for an hour
td = datetime.timedelta(seconds=3600)
every_hour.revoke(revoke_until=dt + td)
self.consumer.enqueue_periodic_commands(dt)
assertQ(0)
# after an hour it is back
self.consumer.enqueue_periodic_commands(dt+td)
assertQ(1)
self.consumer.enqueue_periodic_commands(dt+(td*2))
assertQ(2)
# our data store should reflect the delay
cmd_obj = every_hour.command_class()
self.assertEqual(len(invoker.result_store._results), 1)
self.assertTrue(cmd_obj.revoke_id in invoker.result_store._results)