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) logger.setLevel(logging.INFO) 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 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) logger.setLevel(logging.INFO) 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)