def mock_get_quota(self, allocated, consumed, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, message='test') resp.result = Result( getQuotaResult=GetQuotaResult(quota=deepcopy(allocated), consumed=deepcopy(consumed))) self._scheduler.getQuota.return_value = resp
def expect_get_tasks(self, tasks, ignore_ids=None, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code response = Response(responseCode=response_code, message='test') scheduled = [] for index, task in enumerate(tasks): if not ignore_ids or index not in ignore_ids: scheduled.append(ScheduledTask(assignedTask=AssignedTask(task=task, instanceId=index))) response.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=scheduled)) query = TaskQuery( owner=Identity(role=self._job_key.role), environment=self._job_key.environment, jobName=self._job_key.name, statuses=ACTIVE_STATES) self._scheduler.getTasksStatus(query).AndReturn(response)
def expect_start(self, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code response = Response(responseCode=response_code, message='test') response.result = Result(acquireLockResult=AcquireLockResult(lock=self._lock)) self._scheduler.acquireLock(LockKey(job=self._job_key), self._session_key).AndReturn(response)
def expect_populate(self, job_config, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, message='test') result = set([deepcopy(job_config.taskConfig)]) resp.result = Result(populateJobResult=PopulateJobResult(populated=result)) self._scheduler.populateJobConfig(job_config, JobConfigValidation.RUN_FILTERS).AndReturn(resp)
def mock_get_tasks(self, tasks, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, message='test') resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks)) self._scheduler.getTasksStatus.return_value = resp