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): 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_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)