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 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 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