diff --git a/consulate/api/agent.py b/consulate/api/agent.py index 67270c3..868cd2e 100644 --- a/consulate/api/agent.py +++ b/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 diff --git a/tests/agent_tests.py b/tests/agent_tests.py index a9e94a3..e1c3fb9 100644 --- a/tests/agent_tests.py +++ b/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'))