Callback should only be called on drained hosts.
authorDavid Robinson <drobinson@twopensource.com>
Mon, 9 Nov 2015 23:37:53 +0000 (15:37 -0800)
committerMaxim Khutornenko <maxim@apache.org>
Mon, 9 Nov 2015 23:37:53 +0000 (15:37 -0800)
Bugs closed: AURORA-1537

Reviewed at https://reviews.apache.org/r/40104/

src/main/python/apache/aurora/admin/host_maintenance.py
src/test/python/apache/aurora/admin/test_host_maintenance.py

index cfb13de..677f870 100644 (file)
@@ -161,11 +161,11 @@ class HostMaintenance(object):
     """Perform a given operation on a list of hosts that are ready for maintenance.
 
     :param drained_hosts: Hosts that have been drained (via _drain_hosts)
-    :type drained_hosts: gen.apache.aurora.ttypes.Hosts
+    :type drained_hosts: list of strings
     :param callback: Function to call one hostname at a time
     :type callback: function
     """
-    for hostname in drained_hosts.hostNames:
+    for hostname in drained_hosts:
       callback(hostname)
 
   def perform_maintenance(self, hostnames, grouping_function=DEFAULT_GROUPING,
@@ -215,7 +215,7 @@ class HostMaintenance(object):
       not_drained_hostnames |= self._drain_hosts(hosts)
 
       if callback:
-        self._operate_on_hosts(hosts, callback)
+        self._operate_on_hosts(hosts.hostNames - not_drained_hostnames, callback)
 
     if not_drained_hostnames:
       output = '\n'.join(list(not_drained_hostnames))
index 6264d0c..868dd18 100644 (file)
@@ -175,9 +175,8 @@ class TestHostMaintenance(unittest.TestCase):
 
   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)
+    maintenance._operate_on_hosts(TEST_HOSTNAMES, mock_callback)
     assert mock_callback.call_count == 3
 
   @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance._drain_hosts",
@@ -193,6 +192,11 @@ class TestHostMaintenance(unittest.TestCase):
     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)
@@ -201,8 +205,8 @@ class TestHostMaintenance(unittest.TestCase):
     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(Hosts(set([hostname])), mock_callback) for hostname in TEST_HOSTNAMES]
+    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:]]
 
   @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance._drain_hosts",
               spec=HostMaintenance._drain_hosts)