aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject git commit: Moving post_drain script execution into host_maintenance.py
Date Fri, 17 Oct 2014 20:03:32 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master b5d66b3e2 -> 465ff3419


Moving post_drain script execution into host_maintenance.py

Bugs closed: AURORA-806

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


Project: http://git-wip-us.apache.org/repos/asf/incubator-aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-aurora/commit/465ff341
Tree: http://git-wip-us.apache.org/repos/asf/incubator-aurora/tree/465ff341
Diff: http://git-wip-us.apache.org/repos/asf/incubator-aurora/diff/465ff341

Branch: refs/heads/master
Commit: 465ff3419b384fd035c7ccb8c6136475a80f2d13
Parents: b5d66b3
Author: Maxim Khutornenko <maxim@apache.org>
Authored: Fri Oct 17 13:03:08 2014 -0700
Committer: Maxim Khutornenko <maxim@apache.org>
Committed: Fri Oct 17 13:03:08 2014 -0700

----------------------------------------------------------------------
 .../apache/aurora/admin/host_maintenance.py       | 18 +++++++++++++++++-
 .../apache/aurora/client/commands/maintenance.py  |  9 +++------
 .../apache/aurora/admin/test_host_maintenance.py  | 18 ++++++++++++++++--
 3 files changed, 36 insertions(+), 9 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/465ff341/src/main/python/apache/aurora/admin/host_maintenance.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/admin/host_maintenance.py b/src/main/python/apache/aurora/admin/host_maintenance.py
index 697de4b..bff8afd 100644
--- a/src/main/python/apache/aurora/admin/host_maintenance.py
+++ b/src/main/python/apache/aurora/admin/host_maintenance.py
@@ -165,8 +165,19 @@ class HostMaintenance(object):
 
     return result
 
+  def _operate_on_hosts(self, drained_hosts, callback):
+    """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
+    :param callback: Function to call one hostname at a time
+    :type callback: function
+    """
+    for hostname in drained_hosts.hostNames:
+      callback(hostname)
+
   def perform_maintenance(self, hostnames, grouping_function=DEFAULT_GROUPING,
-                          percentage=None, duration=None, output_file=None):
+                          percentage=None, duration=None, output_file=None, callback=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
@@ -183,6 +194,8 @@ class HostMaintenance(object):
     :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
+    :param callback: Function to call once hosts are drained
+    :type callback: function
     :rtype: set of host names that were successfully drained
     """
     hostnames = self.start_maintenance(hostnames)
@@ -209,6 +222,9 @@ class HostMaintenance(object):
 
       not_drained_hostnames |= self._drain_hosts(hosts)
 
+      if callback:
+        self._operate_on_hosts(hosts, callback)
+
     if not_drained_hostnames:
       output = '\n'.join(list(not_drained_hostnames))
       log.info('The following hosts WERE NOT DRAINED due to failed SLA check or external
failures:')

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/465ff341/src/main/python/apache/aurora/client/commands/maintenance.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/commands/maintenance.py b/src/main/python/apache/aurora/client/commands/maintenance.py
index e465d97..5321eb6 100644
--- a/src/main/python/apache/aurora/client/commands/maintenance.py
+++ b/src/main/python/apache/aurora/client/commands/maintenance.py
@@ -111,16 +111,13 @@ def host_drain(cluster):
 
   post_drain_callback = parse_script(options.post_drain_script)
 
-  drained_hostnames = HostMaintenance(CLUSTERS[cluster], options.verbosity).perform_maintenance(
+  HostMaintenance(CLUSTERS[cluster], options.verbosity).perform_maintenance(
       drainable_hosts,
       grouping_function=options.grouping,
       percentage=override_percentage,
       duration=override_duration,
-      output_file=options.unsafe_hosts_filename)
-
-  if post_drain_callback:
-    for hostname in drained_hostnames:
-      post_drain_callback(hostname)
+      output_file=options.unsafe_hosts_filename,
+      callback=post_drain_callback)
 
 
 @app.command

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/465ff341/src/test/python/apache/aurora/admin/test_host_maintenance.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/admin/test_host_maintenance.py b/src/test/python/apache/aurora/admin/test_host_maintenance.py
index 246081a..4b8072c 100644
--- a/src/test/python/apache/aurora/admin/test_host_maintenance.py
+++ b/src/test/python/apache/aurora/admin/test_host_maintenance.py
@@ -168,22 +168,36 @@ class TestHostMaintenance(unittest.TestCase):
     maintenance.start_maintenance(TEST_HOSTNAMES)
     mock_api.assert_called_once_with(Hosts(set(TEST_HOSTNAMES)))
 
+  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
+
   @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance._drain_hosts",
     spec=HostMaintenance._drain_hosts)
   @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance.start_maintenance",
     spec=HostMaintenance.start_maintenance)
   @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance._check_sla",
     spec=HostMaintenance._check_sla)
-  def test_perform_maintenance(self, mock_check_sla, mock_start_maintenance, mock_drain_hosts):
+  @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance._operate_on_hosts",
+    spec=HostMaintenance._operate_on_hosts)
+  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
     maintenance = HostMaintenance(DEFAULT_CLUSTER, 'quiet')
-    maintenance.perform_maintenance(TEST_HOSTNAMES)
+    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(Hosts(set([hostname])), mock_callback) for hostname in TEST_HOSTNAMES]
 
   @mock.patch("apache.aurora.admin.host_maintenance.HostMaintenance._drain_hosts",
               spec=HostMaintenance._drain_hosts)


Mime
View raw message