Python host_maintenance.HostMaintenance类代码示例

def test_rack_grouping():
    old_grouping_functions = HostMaintenance.GROUPING_FUNCTIONS.copy()
    HostMaintenance.GROUPING_FUNCTIONS["by_rack"] = rack_grouping

    example_host_list = [
        "west-aaa-001.example.com",
        "west-aaa-002.example.com",
        "west-xyz-002.example.com",
        "east-xyz-003.example.com",
        "east-xyz-004.example.com",
    ]

    try:
        batches = list(HostMaintenance.iter_batches(example_host_list, 1, "by_rack"))
        assert batches[0] == Hosts(set(["west-aaa-001.example.com", "west-aaa-002.example.com"]))
        assert batches[1] == Hosts(
            set(["west-xyz-002.example.com", "east-xyz-003.example.com", "east-xyz-004.example.com"])
        )

        batches = list(HostMaintenance.iter_batches(example_host_list, 2, "by_rack"))
        assert batches[0] == Hosts(set(example_host_list))

        batches = list(HostMaintenance.iter_batches(example_host_list, 3, "by_rack"))
        assert batches[0] == Hosts(set(example_host_list))

        with pytest.raises(ValueError):
            list(HostMaintenance.iter_batches(example_host_list, 0))

    finally:
        HostMaintenance.GROUPING_FUNCTIONS = old_grouping_functions

  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_perform_maintenance_partial_sla_failure(self, mock_check_sla, mock_start_maintenance,
                               mock_drain_hosts, mock_operate_on_hosts, mock_complete_maintenance):
    mock_callback = mock.Mock()
    failed_host = 'us-west-001.example.com'
    mock_check_sla.return_value = set([failed_host])
    drained_hosts = set(TEST_HOSTNAMES) - set([failed_host])
    maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')

    with temporary_file() as fp:
      with group_by_rack():
        maintenance.perform_maintenance(
            TEST_HOSTNAMES,
            callback=mock_callback,
            grouping_function='by_rack',
            output_file=fp.name)

        with open(fp.name, 'r') as fpr:
          content = fpr.read()
          assert failed_host in content

        mock_start_maintenance.assert_called_once_with(TEST_HOSTNAMES)
        assert mock_check_sla.call_count == 1
        assert mock_drain_hosts.call_count == 1
        assert mock_drain_hosts.call_args_list == [mock.call(Hosts(drained_hosts))]
        assert mock_operate_on_hosts.call_count == 1
        assert mock_operate_on_hosts.call_args_list == [
            mock.call(Hosts(drained_hosts), mock_callback)]
        assert mock_complete_maintenance.call_count == 2
        assert mock_complete_maintenance.call_args_list == [
            mock.call(Hosts(set([failed_host]))), mock.call(Hosts(drained_hosts))]

 def test_perform_maintenance(self, mock_check_sla, mock_start_maintenance, mock_drain_hosts):
   mock_check_sla.return_value = set()
   mock_start_maintenance.return_value = TEST_HOSTNAMES
   maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')
   maintenance.perform_maintenance(TEST_HOSTNAMES)
   mock_start_maintenance.assert_called_once_with(TEST_HOSTNAMES)
   assert mock_check_sla.call_count == 3
   assert mock_drain_hosts.call_count == 3
   assert mock_drain_hosts.call_args_list == [
       mock.call(Hosts(set([hostname]))) for hostname in TEST_HOSTNAMES]

def test_default_grouping():
    example_host_list = ["xyz321.example.com", "bar337.example.com", "foo001.example.com"]

    batches = list(HostMaintenance.iter_batches(example_host_list, 1))
    assert batches[0] == Hosts(set(["bar337.example.com"]))
    assert batches[1] == Hosts(set(["foo001.example.com"]))
    assert batches[2] == Hosts(set(["xyz321.example.com"]))

    batches = list(HostMaintenance.iter_batches(example_host_list, 2))
    assert batches[0] == Hosts(set(["bar337.example.com", "foo001.example.com"]))
    assert batches[1] == Hosts(set(["xyz321.example.com"]))

  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),
                HostStatus(host=TEST_HOSTNAMES[1], mode=MaintenanceMode.DRAINED),
                HostStatus(host=TEST_HOSTNAMES[2], 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)))]

 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_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_none_grouping():
  example_host_list = [
    'xyz321.example.com',
    'bar337.example.com',
    'foo001.example.com',
  ]

  batches = list(HostMaintenance.iter_batches(example_host_list, 'none'))
  assert batches[0] == Hosts(set(example_host_list))
  assert len(batches) == 1

  def test_drain_hosts(self, mock_drain_hosts, mock_maintenance_status):
    fake_maintenance_status_response = [
        Response(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(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(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(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)

    clock = mock.Mock(time)
    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')
    maintenance._drain_hosts(test_hosts, clock)
    mock_drain_hosts.assert_called_once_with(test_hosts)
    assert clock.sleep.call_count == 4
    assert clock.sleep.call_args == mock.call(
        HostMaintenance.START_MAINTENANCE_DELAY.as_(Time.SECONDS))
    assert mock_maintenance_status.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_default_grouping():
  example_host_list = [
    'xyz321.example.com',
    'bar337.example.com',
    'foo001.example.com',
  ]

  batches = list(HostMaintenance.iter_batches(example_host_list))
  assert batches[0] == Hosts(set(['bar337.example.com']))
  assert batches[1] == Hosts(set(['foo001.example.com']))
  assert batches[2] == Hosts(set(['xyz321.example.com']))

 def test_perform_maintenance(self, mock_operate_on_hosts, mock_check_sla, mock_start_maintenance,
                              mock_drain_hosts):
   mock_callback = mock.Mock()
   mock_check_sla.return_value = set()
   mock_start_maintenance.return_value = TEST_HOSTNAMES
   drain_hosts_results = [set() for _ in TEST_HOSTNAMES]
   drain_hosts_results[0] = set([TEST_HOSTNAMES[0]])
   def drain_hosts_result(*args):
     return drain_hosts_results.pop(0)
   mock_drain_hosts.side_effect = drain_hosts_result
   maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')
   maintenance.perform_maintenance(TEST_HOSTNAMES, callback=mock_callback)
   mock_start_maintenance.assert_called_once_with(TEST_HOSTNAMES)
   assert mock_check_sla.call_count == 3
   assert mock_drain_hosts.call_count == 3
   assert mock_drain_hosts.call_args_list == [
       mock.call(Hosts(set([hostname]))) for hostname in TEST_HOSTNAMES]
   assert mock_operate_on_hosts.call_count == 3
   assert mock_operate_on_hosts.call_args_list == [mock.call(set(), mock_callback)] + [
       mock.call(set([hostname]), mock_callback) for hostname in TEST_HOSTNAMES[1:]]

  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_rack_grouping():
  example_host_list = [
    'west-aaa-001.example.com',
    'west-aaa-002.example.com',
    'west-xyz-002.example.com',
    'east-xyz-003.example.com',
    'east-xyz-004.example.com',
  ]

  with group_by_rack():
    batches = list(HostMaintenance.iter_batches(example_host_list, 'by_rack'))
    assert batches[0] == Hosts(set([
        'west-aaa-001.example.com',
        'west-aaa-002.example.com'
    ]))
    assert batches[1] == Hosts(set([
        'west-xyz-002.example.com',
        'east-xyz-003.example.com',
        'east-xyz-004.example.com',
    ]))

 def test_start_maintenance(self, mock_api):
   mock_api.return_value = Response(responseCode=ResponseCode.OK)
   maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')
   maintenance.start_maintenance(TEST_HOSTNAMES)
   mock_api.assert_called_once_with(Hosts(set(TEST_HOSTNAMES)))