Add Service Maintenance Functionality

consulate/api/agent.py

@@ -274,6 +274,21 @@ def deregister(self, service_id):
             """
             return self._put_no_response_body(['deregister', service_id])
 
+        def maintenance(self, service_id, enable=True, reason=None):
+            """Place given service into "maintenance mode".
+
+            :param str service_id: The id for the service
+            :param bool enable: Enable maintenance mode
+            :param str reason: Reason for putting node in maintenance
+            :rtype: bool
+
+            """
+            query_params = {'enable': enable}
+            if reason:
+                query_params['reason'] = reason
+            return self._put_no_response_body(['maintenance', service_id],
+                                              query_params)
+
     def checks(self):
         """Return the all the checks that are registered with the local agent.
         These checks were either provided through configuration files, or

tests/agent_tests.py

@@ -88,6 +88,22 @@ def test_service_registration(self):
         self.assertIn('test-service', self.consul.agent.services())
         self.consul.agent.service.deregister('test-service')
 
+    def test_service_maintenance(self):
+        self.consul.agent.service.register(
+            'test-service', address='10.0.0.1', port=5672, tags=['foo', 'bar'])
+        self.assertIn('test-service', self.consul.agent.services())
+        reason = 'Down for Acceptance'
+        self.consul.agent.service.maintenance('test-service', reason=reason)
+        node_in_maintenance = self.consul.catalog.nodes()[0]['Node']
+        health_check = self.consul.health.node(node_in_maintenance)
+        self.assertEqual(len(health_check), 2)
+        self.assertIn(reason, [check['Notes'] for check in health_check])
+        self.consul.agent.service.maintenance('test-service', enable=False)
+        health_check = self.consul.health.node(node_in_maintenance)
+        self.assertEqual(len(health_check), 1)
+        self.assertNotEqual(reason, health_check[0]['Notes'])
+        self.consul.agent.service.deregister('test-service')
+
     def test_token(self):
         self.assertTrue(
             self.consul.agent.token('acl_replication_token', 'foo'))