def test_endMaintenance(self): self.mock_thrift_client.endMaintenance(IsA(Hosts), SESSION).AndReturn(DEFAULT_RESPONSE) self.mox.ReplayAll() self.make_scheduler_proxy().endMaintenance(Hosts())
def test_endMaintenance(self):
self.mock_thrift_client.endMaintenance(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().endMaintenance(Hosts())
def test_drainHosts(self): self.mock_thrift_client.drainHosts(IsA(Hosts), SESSION).AndReturn(DEFAULT_RESPONSE) self.mox.ReplayAll() self.make_scheduler_proxy().drainHosts(Hosts())
def test_drainHosts(self):
self.mock_thrift_client.drainHosts(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().drainHosts(Hosts())
def test_maintenanceStatus(self):
self.mock_thrift_client.maintenanceStatus(IsA(Hosts), IsA(SessionKey))
self.mox.ReplayAll()
self.make_scheduler_proxy().maintenanceStatus(Hosts())
def iter_batches(cls, hostnames, grouping_function=DEFAULT_GROUPING):
groups = group_hosts(hostnames, grouping_function)
groups = sorted(groups.items(), key=lambda v: v[0])
for group in groups:
yield Hosts(group[1])
def test_maintenanceStatus(self):
self.mock_thrift_client.maintenanceStatus(
IsA(Hosts)).AndReturn(DEFAULT_RESPONSE)
self.mox.ReplayAll()
self.make_scheduler_proxy().maintenanceStatus(Hosts())
def perform_maintenance(self,
hostnames,
grouping_function=DEFAULT_GROUPING,
percentage=None,
duration=None,
output_file=None):
"""Put hosts into maintenance mode and drain them.
Walk through the process of putting hosts into maintenance and draining them of tasks. The hosts
will remain in maintenance mode upon completion.
:param hostnames: A list of hostnames to operate upon
:type hostnames: list of strings
:param grouping_function: How to split up the hostname into groups
:type grouping_function: function
:param percentage: SLA percentage to use
:type percentage: float
:param duration: SLA duration to use
:type duration: twitter.common.quantity.Time
:param output_file: file to write hosts that were not drained due to failed SLA check
:type output_file: string
:rtype: set of host names that were successfully drained
"""
hostnames = self.start_maintenance(hostnames)
not_drained_hostnames = set()
for hosts in self.iter_batches(hostnames, grouping_function):
log.info('Beginning SLA check for %s' % hosts.hostNames)
unsafe_hostnames = self._check_sla(list(hosts.hostNames),
grouping_function, percentage,
duration)
if unsafe_hostnames:
log.warning(
'Some hosts did not pass SLA check and will not be drained! '
'Skipping hosts: %s' % unsafe_hostnames)
not_drained_hostnames |= unsafe_hostnames
drainable_hostnames = hosts.hostNames - unsafe_hostnames
if not drainable_hostnames:
continue
hosts = Hosts(drainable_hostnames)
else:
log.info('All hosts passed SLA check.')
self._drain_hosts(hosts)
if not_drained_hostnames:
output = '\n'.join(list(not_drained_hostnames))
log.info(
'The following hosts did not pass SLA check and were not drained:'
)
print(output)
if output_file:
try:
with open(output_file, 'w') as fp:
fp.write(output)
fp.write('\n')
log.info('Written unsafe host names into: %s' %
output_file)
except IOError as e:
log.error('Failed to write into the output file: %s' % e)
return set(hostnames) - not_drained_hostnames
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_operate_on_hosts(self):
mock_callback = mock.Mock()
test_hosts = Hosts(TEST_HOSTNAMES)
maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')
maintenance._operate_on_hosts(test_hosts, mock_callback)
assert mock_callback.call_count == 3
def test_end_maintenance(self, mock_complete_maintenance):
maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')
maintenance.end_maintenance(TEST_HOSTNAMES)
mock_complete_maintenance.assert_called_once_with(
Hosts(set(TEST_HOSTNAMES)))
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)))