def get_status_query_response(cls): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() summaries = GetJobUpdateSummariesResult() query_response.result.getJobUpdateSummariesResult = summaries summaries.updateSummaries = [ JobUpdateSummary( updateId="hello", jobKey=AuroraJobKey('west', 'mcc', 'test', 'hello'), user="******", state=JobUpdateState(status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1411404927, lastModifiedTimestampMs=14114056030)), JobUpdateSummary( updateId="goodbye", jobKey=AuroraJobKey('west', 'mch', 'prod', 'goodbye'), user="******", state=JobUpdateState(status=JobUpdateStatus.ROLLING_BACK, createdTimestampMs=1411300632, lastModifiedTimestampMs=14114092632)), JobUpdateSummary(updateId="gasp", jobKey=AuroraJobKey('west', 'mcq', 'devel', 'gasp'), user="******", state=JobUpdateState( status=JobUpdateStatus.ROLL_FORWARD_PAUSED, createdTimestampMs=1411600891, lastModifiedTimestampMs=1411800891)) ] return query_response
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, details=[ResponseDetail(message='test')]) resp.result = Result(scheduleStatusResult=ScheduleStatusResult( tasks=tasks)) self._scheduler.getTasksWithoutConfigs.return_value = resp
def mock_scheduler(cls, response_code=None): scheduler = create_autospec(spec=SchedulerThriftApiSpec, instance=True) response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, details=[ResponseDetail(message='test')]) resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=cls.create_tasks())) scheduler.getTasksWithoutConfigs.return_value = resp return scheduler
def test_drain_hosts_timed_out_wait(self, _, mock_drain_hosts, mock_maintenance_status, mock_log): fake_maintenance_status_response = Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult(set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.SCHEDULED), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.SCHEDULED), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.SCHEDULED) ])))) mock_drain_hosts.return_value = Response(responseCode=ResponseCode.OK) mock_maintenance_status.return_value = fake_maintenance_status_response test_hosts = Hosts(set(TEST_HOSTNAMES)) maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet') maintenance.MAX_STATUS_WAIT = Amount(1, Time.MILLISECONDS) not_drained_hostnames = maintenance._drain_hosts(test_hosts) assert TEST_HOSTNAMES == sorted(not_drained_hostnames) assert mock_maintenance_status.call_count == 1 mock_drain_hosts.assert_called_once_with(test_hosts) mock_maintenance_status.assert_called_once_with((Hosts(set(TEST_HOSTNAMES)))) assert mock_log.mock_calls == [mock.call(textwrap.dedent("""\ Failed to move all hosts into DRAINED within 1 ms: \tHost:us-west-001.example.com\tStatus:SCHEDULED \tHost:us-west-002.example.com\tStatus:SCHEDULED \tHost:us-west-003.example.com\tStatus:SCHEDULED"""))]
def get_update_details_response(cls): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() details = JobUpdateDetails( update=JobUpdate( summary=JobUpdateSummary( key=UPDATE_KEY, user="******", state=JobUpdateState( status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1000, lastModifiedTimestampMs=2000 ), ) ), updateEvents=[ JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=3000), JobUpdateEvent( status=JobUpdateStatus.ROLL_FORWARD_PAUSED, message="Investigating issues", timestampMs=4000 ), JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=5000), ], instanceEvents=[ JobInstanceUpdateEvent(instanceId=1, timestampMs=6000, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent(instanceId=2, timestampMs=7000, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent(instanceId=1, timestampMs=8000, action=JobUpdateAction.INSTANCE_UPDATED), JobInstanceUpdateEvent(instanceId=2, timestampMs=9000, action=JobUpdateAction.INSTANCE_UPDATED), ], ) query_response.result.getJobUpdateDetailsResult = GetJobUpdateDetailsResult(details=details) return query_response
def create_response(cls, quota, prod, non_prod, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, details=[ResponseDetail(message='test')]) resp.result = Result(getQuotaResult=GetQuotaResult( quota=quota, prodConsumption=prod, nonProdConsumption=non_prod)) return resp
def test_transient_error(_, client): mock_scheduler_client = mock.create_autospec( spec=scheduler_client.SchedulerClient, spec_set=False, instance=True) mock_thrift_client = mock.create_autospec(spec=AuroraAdmin.Client, instance=True) mock_thrift_client.killTasks.side_effect = [ Response(responseCode=ResponseCode.ERROR_TRANSIENT, details=[ ResponseDetail(message="message1"), ResponseDetail(message="message2") ], serverInfo=DEFAULT_RESPONSE.serverInfo), Response(responseCode=ResponseCode.ERROR_TRANSIENT, serverInfo=DEFAULT_RESPONSE.serverInfo), Response(responseCode=ResponseCode.OK, serverInfo=DEFAULT_RESPONSE.serverInfo) ] mock_scheduler_client.get_thrift_client.return_value = mock_thrift_client client.get.return_value = mock_scheduler_client proxy = TestSchedulerProxy(Cluster(name='local')) proxy.killTasks(TaskQuery(), None) assert mock_thrift_client.killTasks.call_count == 3
def expect_start(self, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code response = Response(responseCode=response_code, messageDEPRECATED='test', serverInfo=SERVER_INFO) response.result = Result(acquireLockResult=AcquireLockResult(lock=self._lock)) self._scheduler.acquireLock(LockKey(job=self._job_key), self._session_key).AndReturn(response)
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, messageDEPRECATED='test') resp.result = Result(getQuotaResult=GetQuotaResult( quota=deepcopy(allocated), prodConsumption=deepcopy(consumed))) self._scheduler.getQuota.return_value = resp
def get_update_details_response(cls): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() details = JobUpdateDetails() query_response.result.getJobUpdateDetailsResult = GetJobUpdateDetailsResult(details=details) details.update = JobUpdate() details.update.summary = JobUpdateSummary( jobKey=AuroraJobKey("west", "mcc", "test", "hello"), updateId="fake-update-identifier", user="******", state=JobUpdateState( status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1411404927, lastModifiedTimestampMs=14114056030, ), ) details.updateEvents = [ JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=1411404927), JobUpdateEvent(status=JobUpdateStatus.ROLL_FORWARD_PAUSED, timestampMs=1411405000), JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=1411405100), ] details.instanceEvents = [ JobInstanceUpdateEvent(instanceId=1, timestampMs=1411404930, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent(instanceId=2, timestampMs=1411404940, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent(instanceId=1, timestampMs=1411404950, action=JobUpdateAction.INSTANCE_UPDATED), JobInstanceUpdateEvent(instanceId=2, timestampMs=1411404960, action=JobUpdateAction.INSTANCE_UPDATED), ] return query_response
def mock_scheduler(cls, response_code=None): scheduler = Mock() response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, messageDEPRECATED='test') resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=cls.create_tasks())) scheduler.getTasksWithoutConfigs.return_value = resp return scheduler
def validate_quota_from_requested(self, job_key, production, released, acquired): """Validates requested change will not exceed the available quota. Arguments: job_key -- job key. production -- production flag. released -- production CapacityRequest to be released (in case of job update). acquired -- production CapacityRequest to be acquired. Returns: ResponseCode.OK if check is successful. """ resp_ok = Response(responseCode=ResponseCode.OK, messageDEPRECATED='Quota check successful.') if not production: return resp_ok resp = self._scheduler.getQuota(job_key.role) if resp.responseCode != ResponseCode.OK: log.error('Failed to get quota from scheduler: %s' % resp.messageDEPRECATED) return resp allocated = CapacityRequest(resp.result.getQuotaResult.quota) consumed = CapacityRequest(resp.result.getQuotaResult.prodConsumption) requested = acquired - released effective = allocated - consumed - requested if not effective.valid(): log.info('Not enough quota to create/update job.') print_quota(allocated.quota(), 'Total allocated quota', job_key.role) print_quota(consumed.quota(), 'Consumed quota', job_key.role) print_quota(requested.quota(), 'Requested', job_key.name) return Response( responseCode=ResponseCode.INVALID_REQUEST, messageDEPRECATED='Failed quota check.') return resp_ok
def test_drain_hosts_timed_out_wait(self, _, mock_drain_hosts, mock_maintenance_status): fake_maintenance_status_response = Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult( set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.SCHEDULED), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.SCHEDULED), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.SCHEDULED) ])))) mock_drain_hosts.return_value = Response(responseCode=ResponseCode.OK) mock_maintenance_status.return_value = fake_maintenance_status_response test_hosts = Hosts(set(TEST_HOSTNAMES)) maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet') maintenance.MAX_STATUS_WAIT = Amount(1, Time.MILLISECONDS) not_drained_hostnames = maintenance._drain_hosts(test_hosts) assert TEST_HOSTNAMES == sorted(not_drained_hostnames) assert mock_maintenance_status.call_count == 1 mock_drain_hosts.assert_called_once_with(test_hosts) mock_maintenance_status.assert_called_once_with( (Hosts(set(TEST_HOSTNAMES))))
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, messageDEPRECATED='test') result = set([deepcopy(job_config.taskConfig)]) resp.result = Result(populateJobResult=PopulateJobResult( populatedDEPRECATED=result)) self._scheduler.populateJobConfig(job_config).AndReturn(resp)
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, details=[ResponseDetail(message='test')]) resp.result = Result( getQuotaResult=GetQuotaResult( quota=deepcopy(allocated), prodSharedConsumption=deepcopy(consumed))) self._scheduler.getQuota.return_value = resp
def expect_start(self, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code response = Response(responseCode=response_code, messageDEPRECATED='test') response.result = Result(acquireLockResult=AcquireLockResult( lock=self._lock)) self._scheduler.acquireLock(LockKey(job=self._job_key), self._session_key).AndReturn(response)
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, messageDEPRECATED="test") resp.result = Result( getQuotaResult=GetQuotaResult(quota=deepcopy(allocated), prodConsumption=deepcopy(consumed)) ) self._scheduler.getQuota.return_value = resp
def mock_restart_fails(self): response = Response(responseCode=ResponseCode.ERROR, messageDEPRECATED='test error') response.serverInfo = ServerInfo(thriftAPIVersion=THRIFT_API_VERSION) self.mock_scheduler.restartShards( JOB.to_thrift(), IgnoreArg(), self.lock, SESSION_KEY).AndReturn(response)
def mock_scheduler(cls, response_code=None): scheduler = create_autospec(spec=SchedulerThriftApiSpec, instance=True) response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, details=[ResponseDetail(message='test')]) resp.result = Result(scheduleStatusResult=ScheduleStatusResult( tasks=cls.create_tasks())) scheduler.getTasksWithoutConfigs.return_value = resp return scheduler
def expect_get_statuses(self, instance_ids=WATCH_INSTANCES, num_calls=EXPECTED_CYCLES): tasks = [self.create_task(instance_id) for instance_id in instance_ids] response = Response(responseCode=ResponseCode.OK, messageDEPRECATED="test") response.result = Result() response.result.scheduleStatusResult = ScheduleStatusResult(tasks=tasks) query = self.get_tasks_status_query(instance_ids) for _ in range(int(num_calls)): self._scheduler.getTasksWithoutConfigs(query).AndReturn(response)
def mock_restart_instances(self, instances, lock=None): response = Response(responseCode=ResponseCode.OK, messageDEPRECATED='test') response.serverInfo = ServerInfo(thriftAPIVersion=THRIFT_API_VERSION) self.mock_scheduler.restartShards( JOB.to_thrift(), instances, lock, SESSION_KEY).AndReturn(response) self.mock_instance_watcher.watch(instances).AndReturn([])
def test_combine_messages(self): resp = Response(responseCode=ResponseCode.ERROR) assert base.combine_messages(resp) == '' resp = Response(responseCode=ResponseCode.ERROR, details=[]) assert base.combine_messages(resp) == '' resp = Response(responseCode=ResponseCode.ERROR, details=[ResponseDetail(message='Error')]) assert base.combine_messages(resp) == 'Error' resp = Response( responseCode=ResponseCode.ERROR, details=[ResponseDetail(message='Error1'), ResponseDetail(message='Error2')]) assert base.combine_messages(resp) == 'Error1, Error2'
def mock_status_active_tasks(self, instance_ids): tasks = [] for i in instance_ids: tasks.append(ScheduledTask( status=ScheduleStatus.RUNNING, assignedTask=AssignedTask(task=TaskConfig(), instanceId=i) )) response = Response(responseCode=ResponseCode.OK, messageDEPRECATED='test') response.result = Result() response.result.scheduleStatusResult = ScheduleStatusResult(tasks=tasks) self.mock_scheduler.getTasksWithoutConfigs(IgnoreArg()).AndReturn(response)
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, messageDEPRECATED='test', serverInfo=SERVER_INFO) 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(jobKeys=[self._job_key], statuses=ACTIVE_STATES) self._scheduler.getTasksStatus(query).AndReturn(response)
def expect_get_statuses(self, instance_ids=WATCH_INSTANCES, num_calls=EXPECTED_CYCLES): tasks = [self.create_task(instance_id) for instance_id in instance_ids] response = Response(responseCode=ResponseCode.OK, messageDEPRECATED='test') response.result = Result() response.result.scheduleStatusResult = ScheduleStatusResult( tasks=tasks) query = self.get_tasks_status_query(instance_ids) for x in range(int(num_calls)): self._scheduler.getTasksWithoutConfigs(query).AndReturn(response)
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, messageDEPRECATED='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 test_drain_hosts(self, mock_event_wait, mock_drain_hosts, mock_maintenance_status): fake_maintenance_status_response = [ Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult(set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.SCHEDULED), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.SCHEDULED), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.SCHEDULED) ])))), Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult(set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.DRAINING), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.DRAINING), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.DRAINING) ])))), Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult(set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.DRAINING), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.DRAINED), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.DRAINED) ])))), Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult(set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.DRAINED) ]))))] fake_maintenance_status_call_args = [] def fake_maintenance_status_side_effect(hosts): fake_maintenance_status_call_args.append(copy.deepcopy(hosts)) return fake_maintenance_status_response.pop(0) mock_drain_hosts.return_value = Response(responseCode=ResponseCode.OK) mock_maintenance_status.side_effect = fake_maintenance_status_side_effect test_hosts = Hosts(set(TEST_HOSTNAMES)) maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet') not_drained_hostnames = maintenance._drain_hosts(test_hosts) assert len(not_drained_hostnames) == 0 mock_drain_hosts.assert_called_once_with(test_hosts) assert mock_maintenance_status.call_count == 4 assert mock_event_wait.call_count == 4 assert fake_maintenance_status_call_args == [ (Hosts(set(TEST_HOSTNAMES))), (Hosts(set(TEST_HOSTNAMES))), (Hosts(set(TEST_HOSTNAMES))), (Hosts(set([TEST_HOSTNAMES[0]])))]
def get_status_query_response(self, count=3): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() summaries = GetJobUpdateSummariesResult() query_response.result.getJobUpdateSummariesResult = summaries summaries.updateSummaries = [JobUpdateSummary( updateId="%s" % i, jobKey=self.TEST_JOBKEY.to_thrift(), user="******", state=JobUpdateState( status=JobUpdateStatus.ROLLED_FORWARD, createdTimestampMs=1411404927, lastModifiedTimestampMs=14114056030)) for i in range(count)] return query_response
def test_complete_maintenance(self, mock_end_maintenance, mock_maintenance_status, mock_warning): mock_maintenance_status.return_value = Response(result=Result( maintenanceStatusResult=MaintenanceStatusResult(set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.NONE), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.NONE), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.DRAINED) ])) )) mock_end_maintenance.return_value = Response(responseCode=ResponseCode.OK) test_hosts = Hosts(set(TEST_HOSTNAMES)) maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet') maintenance._complete_maintenance(test_hosts) mock_end_maintenance.assert_called_once_with(test_hosts) mock_maintenance_status.assert_called_once_with(test_hosts) mock_warning.assert_called_once_with('%s is DRAINING or in DRAINED' % TEST_HOSTNAMES[2])
def mock_status_active_tasks(self, instance_ids): tasks = [] for i in instance_ids: tasks.append( ScheduledTask(status=ScheduleStatus.RUNNING, assignedTask=AssignedTask(task=TaskConfig(), instanceId=i))) response = Response(responseCode=ResponseCode.OK, messageDEPRECATED='test') response.result = Result() response.result.scheduleStatusResult = ScheduleStatusResult( tasks=tasks) self.mock_scheduler.getTasksWithoutConfigs( IgnoreArg()).AndReturn(response)
def mock_restart_instances(self, instances, lock=None): response = Response(responseCode=ResponseCode.OK, messageDEPRECATED='test') self.mock_scheduler.restartShards(JOB.to_thrift(), instances, lock, SESSION_KEY).AndReturn(response) self.mock_instance_watcher.watch(instances).AndReturn([])
def mock_restart_fails(self): response = Response(responseCode=ResponseCode.ERROR, messageDEPRECATED='test error') self.mock_scheduler.restartShards(JOB.to_thrift(), IgnoreArg(), self.lock, SESSION_KEY).AndReturn(response)
def create_response(cls, quota, prod, non_prod, response_code=ResponseCode.OK): return Response( responseCode=response_code, details=[ResponseDetail(message='test')], result=Result(getQuotaResult=GetQuotaResult( quota=quota, prodSharedConsumption=prod, nonProdSharedConsumption=non_prod)) )
def test_check_and_log_response(self, mock_sys_exit, mock_log): resp = Response(responseCode=ResponseCode.LOCK_ERROR) out = base.check_and_log_response(resp) self.assertIsNone(out) mock_sys_exit.assert_called_once_with(1) mock_log.assert_any_call( 'Response from scheduler: LOCK_ERROR (message: )')
def test_check_status(self, mock_maintenance_status): mock_maintenance_status.return_value = Response( responseCode=ResponseCode.OK, result=Result(maintenanceStatusResult=MaintenanceStatusResult( set([ HostStatus(host=TEST_HOSTNAMES[0], mode=MaintenanceMode.DRAINING), HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.DRAINED), HostStatus(host=TEST_HOSTNAMES[2], mode=MaintenanceMode.NONE) ])))) maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet') result = maintenance.check_status(TEST_HOSTNAMES) mock_maintenance_status.assert_called_once_with( Hosts(set(TEST_HOSTNAMES))) assert len(result) == 3 assert (TEST_HOSTNAMES[0], MaintenanceMode._VALUES_TO_NAMES[MaintenanceMode.DRAINING] ) in result assert (TEST_HOSTNAMES[1], MaintenanceMode._VALUES_TO_NAMES[MaintenanceMode.DRAINED] ) in result assert ( TEST_HOSTNAMES[2], MaintenanceMode._VALUES_TO_NAMES[MaintenanceMode.NONE]) in result
def test_api_version_mismatch(self): resp = Response(serverInfo=ServerInfo( thriftAPIVersion=THRIFT_API_VERSION + 1)) self.mock_thrift_client.getQuota(IgnoreArg()).AndReturn(resp) self.mox.ReplayAll() with pytest.raises( scheduler_client.SchedulerProxy.ThriftInternalError): self.make_scheduler_proxy().getQuota(ROLE)
def test_start_maintenance(self, mock_api): mock_api.return_value = Response( responseCode=ResponseCode.OK, result=Result(startMaintenanceResult=StartMaintenanceResult( statuses=set([HostStatus()])))) maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet') maintenance.start_maintenance(TEST_HOSTNAMES) mock_api.assert_called_once_with(Hosts(set(TEST_HOSTNAMES)))
def _expect_get_jobs(self, *envs): self._api.get_jobs(self.ROLE).AndReturn( Response(responseCode=ResponseCode.OK, messageDEPRECATED='Mock OK', result=Result(getJobsResult=GetJobsResult(configs=set( JobConfiguration(key=JobKey( role=self.ROLE, environment=env, name=self.NAME)) for env in envs)))))
def setUp(self): self.mox = Mox() self.mox.StubOutClassWithMocks(AuroraAdmin, 'Client') self.mox.StubOutClassWithMocks(scheduler_client, 'SchedulerClient') self.mock_scheduler_client = self.mox.CreateMock(scheduler_client.SchedulerClient) self.mock_thrift_client = self.mox.CreateMock(AuroraAdmin.Client) scheduler_client.SchedulerClient.get(IgnoreArg(), verbose=IgnoreArg()).AndReturn( self.mock_scheduler_client) self.mock_scheduler_client.get_thrift_client().AndReturn(self.mock_thrift_client) version_resp = Response(responseCode=ResponseCode.OK) version_resp.result = Result(getVersionResult=CURRENT_API_VERSION) self.mock_thrift_client.getVersion().AndReturn(version_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, messageDEPRECATED='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 validate_quota_from_requested(self, job_key, production, released, acquired): """Validates requested change will not exceed the available quota. Arguments: job_key -- job key. production -- production flag. released -- production CapacityRequest to be released (in case of job update). acquired -- production CapacityRequest to be acquired. Returns: ResponseCode.OK if check is successful. """ # TODO(wfarner): Avoid synthesizing scheduler responses. resp_ok = Response( responseCode=ResponseCode.OK, details=[ResponseDetail(message='Quota check successful.')]) if not production: return resp_ok resp = self._scheduler.getQuota(job_key.role) if resp.responseCode != ResponseCode.OK: log.error('Failed to get quota from scheduler: %s' % combine_messages(resp)) return resp allocated = CapacityRequest(resp.result.getQuotaResult.quota) consumed = CapacityRequest( resp.result.getQuotaResult.prodSharedConsumption) requested = acquired - released effective = allocated - consumed - requested if not effective.valid(): log.info('Not enough quota to create/update job.') print_quota(allocated.quota(), 'Total allocated quota', job_key.role) print_quota(consumed.quota(), 'Consumed quota', job_key.role) print_quota(requested.quota(), 'Requested', job_key.name) print_quota(effective.invert_or_reset().quota(), 'Additional quota required', job_key.role) # TODO(wfarner): Avoid synthesizing scheduler responses. return Response( responseCode=ResponseCode.INVALID_REQUEST, details=[ResponseDetail(message='Failed quota check.')]) return resp_ok
def expect_restart(self, instance_ids, response_code=None): for i in instance_ids: response_code = ResponseCode.OK if response_code is None else response_code response = Response(responseCode=response_code, messageDEPRECATED='test') self._scheduler.restartShards( self._job_key, [i], self._lock, self._session_key).AndReturn(response)
def get_update_details_response(cls): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() details = JobUpdateDetails( update=JobUpdate(summary=JobUpdateSummary( key=UPDATE_KEY, user="******", state=JobUpdateState(status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1000, lastModifiedTimestampMs=2000), metadata={ Metadata("issue", "test"), Metadata("country", "America"), Metadata("country", "Canada") })), updateEvents=[ JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=3000), JobUpdateEvent(status=JobUpdateStatus.ROLL_FORWARD_PAUSED, message="Investigating issues", timestampMs=4000), JobUpdateEvent(status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=5000) ], instanceEvents=[ JobInstanceUpdateEvent( instanceId=1, timestampMs=6000, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent( instanceId=2, timestampMs=7000, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent( instanceId=1, timestampMs=8000, action=JobUpdateAction.INSTANCE_UPDATED), JobInstanceUpdateEvent(instanceId=2, timestampMs=9000, action=JobUpdateAction.INSTANCE_UPDATED) ]) query_response.result.getJobUpdateDetailsResult = GetJobUpdateDetailsResult( detailsList=[details]) return query_response
def get_update_details_response(self): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() details = JobUpdateDetails( update=JobUpdate( summary=JobUpdateSummary( jobKey=self.TEST_JOBKEY.to_thrift(), updateId="0", user="******", state=JobUpdateState( status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1411404927, lastModifiedTimestampMs=14114056030))), updateEvents=[ JobUpdateEvent( status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=1411404927), JobUpdateEvent( status=JobUpdateStatus.ROLL_FORWARD_PAUSED, timestampMs=1411405000), JobUpdateEvent( status=JobUpdateStatus.ROLLING_FORWARD, timestampMs=1411405100)], instanceEvents=[ JobInstanceUpdateEvent( instanceId=1, timestampMs=1411404930, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent( instanceId=2, timestampMs=1411404940, action=JobUpdateAction.INSTANCE_UPDATING), JobInstanceUpdateEvent( instanceId=1, timestampMs=1411404950, action=JobUpdateAction.INSTANCE_UPDATED), JobInstanceUpdateEvent( instanceId=2, timestampMs=1411404960, action=JobUpdateAction.INSTANCE_UPDATED)]) query_response.result.getJobUpdateDetailsResult = GetJobUpdateDetailsResult(details=details) return query_response
def get_status_query_response(cls): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() summaries = GetJobUpdateSummariesResult() query_response.result.getJobUpdateSummariesResult = summaries summaries.updateSummaries = [ JobUpdateSummary( updateId="hello", jobKey=AuroraJobKey("west", "mcc", "test", "hello"), user="******", state=JobUpdateState( status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1411404927, lastModifiedTimestampMs=14114056030, ), ), JobUpdateSummary( updateId="goodbye", jobKey=AuroraJobKey("west", "mch", "prod", "goodbye"), user="******", state=JobUpdateState( status=JobUpdateStatus.ROLLING_BACK, createdTimestampMs=1411300632, lastModifiedTimestampMs=14114092632, ), ), JobUpdateSummary( updateId="gasp", jobKey=AuroraJobKey("west", "mcq", "devel", "gasp"), user="******", state=JobUpdateState( status=JobUpdateStatus.ROLL_FORWARD_PAUSED, createdTimestampMs=1411600891, lastModifiedTimestampMs=1411800891, ), ), ] return query_response
def get_status_query_response(cls): query_response = Response() query_response.responseCode = ResponseCode.OK query_response.result = Result() summaries = GetJobUpdateSummariesResult() query_response.result.getJobUpdateSummariesResult = summaries summaries.updateSummaries = [ JobUpdateSummary( updateId="hello", jobKey=AuroraJobKey('west', 'mcc', 'test', 'hello'), user="******", state=JobUpdateState(status=JobUpdateStatus.ROLLING_FORWARD, createdTimestampMs=1411404927, lastModifiedTimestampMs=14114056030)), JobUpdateSummary( updateId="goodbye", jobKey=AuroraJobKey('west', 'mch', 'prod', 'goodbye'), user="******", state=JobUpdateState(status=JobUpdateStatus.ROLLING_BACK, createdTimestampMs=1411300632, lastModifiedTimestampMs=14114092632)), JobUpdateSummary( updateId="gasp", jobKey=AuroraJobKey('west', 'mcq', 'devel', 'gasp'), user="******", state=JobUpdateState(status=JobUpdateStatus.ROLL_FORWARD_PAUSED, createdTimestampMs=1411600891, lastModifiedTimestampMs=1411800891))] return query_response
def create_response(cls, locks, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, details=[ResponseDetail(message='test')]) resp.result = Result(getLocksResult=GetLocksResult(locks=locks)) return 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, details=[ResponseDetail(message="test")]) resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks)) self._scheduler.getTasksWithoutConfigs.return_value = resp
def mock_status_no_active_task(self): response = Response(responseCode=ResponseCode.INVALID_REQUEST, messageDEPRECATED='test') response.serverInfo = ServerInfo(thriftAPIVersion=THRIFT_API_VERSION) self.mock_scheduler.getTasksWithoutConfigs(IgnoreArg()).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, messageDEPRECATED='test') result = set([deepcopy(job_config.taskConfig)]) resp.result = Result(populateJobResult=PopulateJobResult(populated=result)) self._scheduler.populateJobConfig(job_config).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, messageDEPRECATED='test') resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks)) self._scheduler.getTasksWithoutConfigs.return_value = resp
def create_response(cls, locks, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, messageDEPRECATED='test') resp.result = Result(getLocksResult=GetLocksResult(locks=locks)) return resp
def create_response(cls, tasks, response_code=None): response_code = ResponseCode.OK if response_code is None else response_code resp = Response(responseCode=response_code, messageDEPRECATED='test') resp.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=tasks)) return resp