def setUp(self): self.tearDown() self.scheduler = scheduler.TaskScheduler() self.scheduler.task_manager = TaskManagerProxy() self.manager = ModuleManagerProxy() self.scheduler._register(self.manager) # intercept _schedule so that its calls can be recorded CallProxy.patch(self.scheduler, '_schedule') # hook the threads modules in the scheduling module so we can intercept # any calls to deferToThread() self.threads_ = ThreadsProxy(self) scheduler.threads = self.threads_
class TaskScheduler_Base(django.TestCase, ModuleTestCaseMixIn): """ Base Test class for TaskScheduler - the class responsible for tracking and decision making for the task queue. This class contains some base setup and utility methods. """ def setUp(self): self.tearDown() self.scheduler = scheduler.TaskScheduler() self.scheduler.task_manager = TaskManagerProxy() self.manager = ModuleManagerProxy() self.scheduler._register(self.manager) # intercept _schedule so that its calls can be recorded CallProxy.patch(self.scheduler, '_schedule') # hook the threads modules in the scheduling module so we can intercept # any calls to deferToThread() self.threads_ = ThreadsProxy(self) scheduler.threads = self.threads_ def add_worker(self, connect=False): """ Helper function for adding a worker to the scheduler """ worker = RemoteProxy('localhost:%d'%len(self.scheduler.workers)) self.scheduler.workers[worker.name] = worker if connect: # connect the proxy worker fully so that it can be scheduled self.scheduler.worker_status_returned([WORKER_STATUS_IDLE], worker, worker.name) return worker def queue_and_run_task(self, success=None): """ Helper for setting up a running task """ s = self.scheduler worker = self.add_worker(True) task = s._queue_task('foo.bar') s._schedule.enable() response = s._schedule() # complete start sequence for task, or fail it. if no flag is given # task will be left waiting for response from remote worker. if success == True: s.run_task_successful(None, worker.name) elif success == False: s.run_task_failed(None, worker.name) return response, worker, task def queue_and_run_subtask(self, worker, success=None, workunit_key=1): """ Helper for setting up a running subtask """ s = self.scheduler s._schedule.disable() subtask = s.request_worker(worker.name, 'test.foo.bar', 'args', workunit_key) self.assert_(subtask, "subtask was not created") s._schedule.enable() response = s._schedule() # complete start sequence for subtask, or fail it. if no flag is given # subtask will be left waiting for response from remote worker. if success == True: s.run_task_successful(None, worker.name, subtask.subtask_key) elif success == False: s.run_task_failed(None, worker.name, subtask.subtask_key) return response, subtask def assertWorkerStatus(self, worker, status, scheduler, main=True): """ Assertion function for checking the status of a worker. This encapsulates checking the presence of the worker in various pools which denote its state. This ensures that if the code for tracking worker state changes, then only this assertion needs to change. """ if status==WORKER_IDLE: in_ = '_idle_workers' not_in = ['_active_workers','_waiting_workers'] self.assertFalse(worker.name in scheduler._main_workers, "Worker (%s) shouldn't be in main workers" % worker.name) elif status==WORKER_WAITING: in_ = '_waiting_workers' not_in = ['_active_workers','_idle_workers'] elif status==WORKER_ACTIVE: in_ = '_active_workers' not_in = ['_waiting_workers','_idle_workers'] if main: self.assert_(worker.name in scheduler._main_workers, "Worker (%s) isn't in main workers" % worker.name) for pool in not_in: self.assertFalse(worker.name in getattr(scheduler, pool), "Worker (%s) shouldn't be in %s" % (worker.name,pool)) self.assert_(worker.name in getattr(scheduler, in_), "Worker (%s) isn't in %s" % (worker.name,in_)) def assertSchedulerAdvanced(self): """ asserts that the scheduler was attempted to be advanced using either deferToThread or directly called """ self.assert_(self.scheduler._schedule.calls or self.threads_.was_deferred(self.scheduler._schedule), "No Attempt was made to advance the scheduler") # clear calls self.scheduler._schedule.reset() self.threads_.calls = [] def tearDown(self): self.scheduler = None Batch.objects.all().delete() WorkUnit.objects.all().delete() TaskInstance.objects.all().delete() clean_reactor() def validate_queue_format(self): """ helper function for validating format of objects in the queue. This function should be used by anything that adds or modifies the queue. This will help ensure that any changes to the queue structure will be detected in all places that it needs to be changed. """ s = self.scheduler # check formatting of _queue for task in s._queue: self.assert_(isinstance(task, (list,)), type(task)) score, taskinstance = task self.assert_(isinstance(score, (list,tuple)), type(score)) self.assert_(isinstance(taskinstance, (TaskInstance,)), type(taskinstance)) # check formatting of _active_tasks for key, value in s._active_tasks.items(): self.assert_(isinstance(key, (long, int)), type(key)) self.assert_(isinstance(value, (TaskInstance,)), type(value))