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 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_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_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_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 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 create_drained_status_result(self, hosts):
host_statuses = [
HostStatus(host=hostname, mode=MaintenanceMode.DRAINED)
for hostname in hosts.hostNames
]
response = self.create_simple_success_response()
response.result.maintenanceStatusResult = MaintenanceStatusResult(
statuses=set(host_statuses))
return response
def create_host_statuses(self, maintenance_mode, skip_hosts=None):
return [
HostStatus(host=hostname, mode=maintenance_mode)
for hostname in self.HOSTNAMES
if not skip_hosts or hostname not in skip_hosts
]