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 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 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 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_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 setup_mock_quota_call_with_consumption(cls, mock_context):
api = mock_context.get_api('west')
response = cls.create_simple_success_response()
response.result = Result(getQuotaResult=GetQuotaResult(
quota=ResourceAggregate(resources=frozenset([
Resource(numCpus=5),
Resource(ramMb=20480),
Resource(diskMb=40960)
])),
prodSharedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=1),
Resource(ramMb=512),
Resource(diskMb=1024)
])),
prodDedicatedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=2),
Resource(ramMb=1024),
Resource(diskMb=2048)
])),
nonProdSharedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=3),
Resource(ramMb=2048),
Resource(diskMb=4096)
])),
nonProdDedicatedConsumption=ResourceAggregate(resources=frozenset([
Resource(numCpus=4),
Resource(ramMb=4096),
Resource(diskMb=8192)
])),
))
api.get_quota.return_value = 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 test_start_update_and_wait_success(self):
mock_config = self.create_mock_config()
self._fake_context.get_job_config = Mock(return_value=mock_config)
self._mock_options.wait = True
resp = self.create_simple_success_response()
resp.result = Result(startJobUpdateResult=StartJobUpdateResult(
key=JobUpdateKey(job=JobKey(
role="role", environment="env", name="name"),
id="id")))
self._mock_api.start_job_update.return_value = resp
self._mock_api.query_job_updates.side_effect = [
get_status_query_response(status=JobUpdateStatus.ROLLED_FORWARD)
]
assert self._command.execute(self._fake_context) == EXIT_OK
assert self._mock_api.start_job_update.mock_calls == [
call(ANY, None, None)
]
assert self._mock_api.query_job_updates.mock_calls == [
call(update_key=resp.result.startJobUpdateResult.key)
]
assert self._fake_context.get_out() == [
StartUpdate.UPDATE_MSG_TEMPLATE %
('http://something_or_other/scheduler/role/env/name/update/id'),
'Current state ROLLED_FORWARD'
]
assert self._fake_context.get_err() == []
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 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 test_kill_job_with_empty_instances_batched(self):
"""Test kill client-side API logic."""
mock_context = FakeAuroraCommandContext()
with contextlib.nested(
patch('apache.aurora.client.cli.jobs.Job.create_context',
return_value=mock_context),
patch('apache.aurora.client.factory.CLUSTERS',
new=self.TEST_CLUSTERS)):
api = mock_context.get_api('west')
# set up an empty instance list in the getTasksWithoutConfigs response
status_response = self.create_simple_success_response()
status_response.result = Result(
scheduleStatusResult=ScheduleStatusResult(tasks=[]))
mock_context.add_expected_query_result(status_response)
api.kill_job.return_value = self.create_simple_success_response()
with temporary_file() as fp:
fp.write(self.get_valid_config())
fp.flush()
cmd = AuroraCommandLine()
cmd.execute([
'job', 'kill',
'--config=%s' % fp.name,
self.get_instance_spec('0,2,4-13')
])
assert api.kill_job.call_count == 0
def setup_get_tasks_status_calls(cls, scheduler):
status_response = cls.create_simple_success_response()
scheduler.getTasksStatus.return_value = status_response
scheduler.getTasksWithoutConfigs.return_value = status_response
task_config = TaskConfig(numCpus=1.0,
ramMb=10,
diskMb=1,
job=JobKey(role='bozo',
environment='test',
name='hello'))
# This should be a list of ScheduledTask's.
tasks = []
for i in range(20):
task_status = create_autospec(spec=ScheduledTask, instance=True)
task_status.assignedTask = create_autospec(spec=AssignedTask,
instance=True)
task_status.assignedTask.instanceId = i
task_status.assignedTask.taskId = "Task%s" % i
task_status.assignedTask.slaveId = "Slave%s" % i
task_status.slaveHost = "Slave%s" % i
task_status.assignedTask.task = task_config
tasks.append(task_status)
status_response.result = Result(
scheduleStatusResult=ScheduleStatusResult(tasks=tasks))
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 test_start_update_command_line_succeeds(self):
resp = self.create_simple_success_response()
resp.result = Result(startJobUpdateResult=StartJobUpdateResult(
key=JobUpdateKey(job=JobKey(
role="role", environment="env", name="name"),
id="id")))
self._mock_api.start_job_update.return_value = resp
mock_config = self.create_mock_config()
self._fake_context.get_job_config = Mock(return_value=mock_config)
self._mock_options.instance_spec = TaskInstanceKey(self._job_key, None)
self._mock_options.message = 'hello'
with patch(
'apache.aurora.client.cli.update.DiffFormatter') as formatter:
formatter.return_value = self._formatter
assert self._command.execute(self._fake_context) == EXIT_OK
assert self._formatter.show_job_update_diff.mock_calls == [
call(self._mock_options.instance_spec.instance)
]
assert self._mock_api.start_job_update.mock_calls == [
call(ANY, 'hello', None, ANY)
]
assert self._fake_context.get_out() == [
StartUpdate.UPDATE_MSG_TEMPLATE %
('http://something_or_other/scheduler/role/env/name/update/id'),
]
assert self._fake_context.get_err() == []
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 create_mock_status_query_result(cls, scheduleStatus):
query_result = cls.create_simple_success_response()
query_result.result = Result(scheduleStatusResult=ScheduleStatusResult(tasks=[
cls.create_scheduled_task(0, initial_time=1000, status=scheduleStatus),
cls.create_scheduled_task(1, initial_time=1004, status=scheduleStatus)
]))
return query_result
def create_acquire_lock_response(cls, code, msg, token, rolling):
"""Set up the response to a startUpdate API call."""
start_update_response = cls.create_blank_response(code, msg)
start_update_response.result = Result(
acquireLockResult=AcquireLockResult(
lock=Lock(key='foo', token='token')))
return start_update_response
def expect_populate(self, job_config, response_code=ResponseCode.OK):
resp = make_response(response_code)
config = deepcopy(job_config.taskConfig)
resp.result = Result(populateJobResult=PopulateJobResult(
taskConfig=config))
self._scheduler.populateJobConfig(job_config).AndReturn(resp)
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 test_cron_status_multiple_jobs(self):
_, mock_scheduler_proxy = self.create_mock_api()
with contextlib.nested(
patch('time.sleep'),
patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS),
patch('apache.aurora.client.cli.context.AuroraCommandContext.print_out')) as (
_, _, _, mock_print):
response = self.create_simple_success_response()
response.result = Result(getJobsResult=GetJobsResult(configs=[
JobConfiguration(
key=JobKey(role='bozo', environment='test', name='hello'),
cronSchedule='* * * * *'),
JobConfiguration(
key=JobKey(role='bozo', environment='test', name='hello2'),
cronSchedule='* * * * *')
]))
mock_scheduler_proxy.getJobs.return_value = response
cmd = AuroraCommandLine()
result = cmd.execute(['cron', 'show', 'west/bozo/test/hello'])
assert result == EXIT_OK
mock_scheduler_proxy.getJobs.assert_called_once_with("bozo")
mock_print.assert_called_with("west/bozo/test/hello\t * * * * *")
def test_start_update_command_line_succeeds(self):
mock_context = FakeAuroraCommandContext()
resp = self.create_simple_success_response()
resp.result = Result(startJobUpdateResult=StartJobUpdateResult(updateId="id"))
with contextlib.nested(
patch('apache.aurora.client.cli.update.Update.create_context', return_value=mock_context),
patch('apache.aurora.client.factory.CLUSTERS', new=self.TEST_CLUSTERS)):
mock_api = mock_context.get_api(self.TEST_CLUSTER)
mock_api.start_job_update.return_value = resp
with temporary_file() as fp:
fp.write(self.get_valid_config())
fp.flush()
cmd = AuroraCommandLine()
result = cmd.execute(['beta-update', 'start', self.TEST_JOBSPEC, fp.name])
assert result == EXIT_OK
update_url_msg = StartUpdate.UPDATE_MSG_TEMPLATE % (
mock_context.get_update_page(mock_api, AuroraJobKey.from_path(self.TEST_JOBSPEC), "id"))
assert mock_api.start_job_update.call_count == 1
args, kwargs = mock_api.start_job_update.call_args
assert isinstance(args[0], AuroraConfig)
assert args[1] is None
assert mock_context.get_out() == [update_url_msg]
assert mock_context.get_err() == []
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 _create_getjobs_response(cls):
response = cls.create_simple_success_response()
response.result = Result(getJobsResult=GetJobsResult(configs=[
JobConfiguration(
cronSchedule='* * * * *',
key=JobKey(role='bozo', environment='test', name='hello'))
]))
return response
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 setup_populate_job_config(cls, api):
populate = cls.create_simple_success_response()
api.populateJobConfig.return_value = populate
tasks = set(task.assignedTask.task
for task in cls.create_scheduled_tasks())
populate.result = Result(populateJobResult=PopulateJobResult(
populatedDEPRECATED=tasks, taskConfig=list(tasks)[0]))
return populate
def get_mock_tier_configurations(cls):
response = cls.create_simple_success_response()
response.result = Result(getTierConfigResult=GetTierConfigResult(
defaultTierName=cls.PREEMPTIBLE_TIER.name,
tiers=frozenset([
cls.PREFERRED_TIER, cls.PREEMPTIBLE_TIER, cls.REVOCABLE_TIER
])))
return response
def setup_mock_quota_call_no_consumption(cls, mock_context):
api = mock_context.get_api('west')
response = cls.create_simple_success_response()
response.result = Result(getQuotaResult=GetQuotaResult(
quota=ResourceAggregate(numCpus=5, ramMb=20480, diskMb=40960),
prodConsumption=None,
nonProdConsumption=None))
api.get_quota.return_value = 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,
messageDEPRECATED='test')
response.result = Result(acquireLockResult=AcquireLockResult(
lock=self._lock))
self._scheduler.acquireLock(LockKey(job=self._job_key),
self._session_key).AndReturn(response)
def setup_populate_job_config(cls, api):
populate = cls.create_simple_success_response()
populate.result = Result(populateJobResult=PopulateJobResult(
populatedDEPRECATED={TaskConfig()},
taskConfig=TaskConfig()
))
api.populateJobConfig.return_value = populate
return populate
def create_getjobs_response(cls):
mock_job_one = JobConfiguration(
key=JobKey(role='RoleA', environment='test', name='hithere'))
mock_job_two = JobConfiguration(
key=JobKey(role='bozo', environment='test', name='hello'))
result = cls.create_simple_success_response()
result.result = Result(getJobsResult=GetJobsResult(
configs=[mock_job_one, mock_job_two]))
return result
