Adding tests. Fixing up docs a bit.

cruatta · cruatta · commit 78110a6ec6eb · 2015-06-02T23:38:49.000-07:00
diff --git a/consul/base.py b/consul/base.py
@@ -475,14 +475,13 @@ def members(self, wan=False):
                 '/v1/agent/members',
                 params=params)
 
-
         def maintenance(self, enable, reason=None):
             """
             The node maintenance endpoint can place the agent into
             "maintenance mode".
 
-            *enable* is either True or False. True enables maintenance mode,
-            False disables maintenance mode.
+            *enable* is either 'true' or 'false'. True enables maintenance
+            mode, 'false' disables maintenance mode.
 
             *reason* is an optional string. This is simply to aid human
             operators.
@@ -499,7 +498,6 @@ def maintenance(self, enable, reason=None):
                 '/v1/agent/maintenance',
                 params=params)
 
-
         class Service(object):
             def __init__(self, agent):
                 self.agent = agent
@@ -567,11 +565,11 @@ def maintenance(self, service_id, enable, reason=None):
                 The service maintenance endpoint allows placing a given service
                 into "maintenance mode".
 
-                *service_id* is the id of the service that is to be targeted for
-                maintenance. 
+                *service_id* is the id of the service that is to be targeted
+                for maintenance.
 
-                *enable* is either True or False. True enables maintenance mode,
-                False disables maintenance mode.
+                *enable* is either True or False. 'true' enables maintenance
+                mode, 'false' disables maintenance mode.
 
                 *reason* is an optional string. This is simply to aid human
                 operators.
diff --git a/tests/test_std.py b/tests/test_std.py
@@ -289,6 +289,51 @@ def test_agent_register_check_no_service_id(self, consul_port):
 
         time.sleep(40/1000.0)
 
+    def test_agent_service_maintenance(self, consul_port):
+        c = consul.Consul(port=consul_port)
+
+        c.agent.service.register('foo', ttl='100ms')
+
+        time.sleep(40/1000.0)
+
+        c.agent.service.maintenance('foo', 'true', "test")
+
+        time.sleep(40/1000.0)
+
+        checks_pre = c.agent.checks()
+        assert '_service_maintenance:foo' in checks_pre.keys()
+        assert 'test' == checks_pre['_service_maintenance:foo']['Notes']
+
+        c.agent.service.maintenance('foo', 'false')
+
+        time.sleep(40/1000.0)
+
+        checks_post = c.agent.checks()
+        assert '_service_maintenance:foo' not in checks_post.keys()
+
+        # Cleanup
+        c.agent.service.deregister('foo')
+
+        time.sleep(40/1000.0)
+
+    def test_agent_node_maintenance(self, consul_port):
+        c = consul.Consul(port=consul_port)
+
+        c.agent.maintenance('true', "test")
+
+        time.sleep(40/1000.0)
+
+        checks_pre = c.agent.checks()
+        assert '_node_maintenance' in checks_pre.keys()
+        assert 'test' == checks_pre['_node_maintenance']['Notes']
+
+        c.agent.maintenance('false')
+
+        time.sleep(40/1000.0)
+
+        checks_post = c.agent.checks()
+        assert '_node_maintenance' not in checks_post.keys()
+
     def test_agent_members(self, consul_port):
         c = consul.Consul(port=consul_port)
         members = c.agent.members()
