#26 add weight support

consul/base.py

@@ -284,6 +284,20 @@ def cb(response):
         return cb
 
 
+#
+# Convenience to define weight
+
+class Weight(object):
+    """
+    There object for set weights parameters like this
+    {'passing': 100, 'warning': 100}
+    """
+
+    @classmethod
+    def weights(cls, passing, warning):
+        return {'passing': passing, 'warning': warning}
+
+
 class HTTPClient(six.with_metaclass(abc.ABCMeta, object)):
     def __init__(self, host='127.0.0.1', port=8500, scheme='http',
                  verify=True, cert=None, timeout=None):
@@ -1364,6 +1378,7 @@ def register(
                     check=None,
                     token=None,
                     meta=None,
+                    weights=None,
                     # *deprecated* use check parameter
                     script=None,
                     interval=None,
@@ -1396,6 +1411,10 @@ def register(
                 *meta* specifies arbitrary KV metadata linked to the service
                 formatted as {k1:v1, k2:v2}.
 
+                *weights* specifies weights for the service.
+                If this field is not provided weights
+                will default to {"Passing": 1, "Warning": 1}.
+
                 *script*, *interval*, *ttl*, *http*, and *timeout* arguments
                 are deprecated. use *check* instead.
 
@@ -1425,6 +1444,8 @@ def register(
                     payload['meta'] = meta
                 if check:
                     payload['check'] = check
+                if weights:
+                    payload['weights'] = weights
 
                 else:
                     payload.update(Check._compat(

tests/test_std.py

@@ -8,6 +8,7 @@
 
 import consul
 import consul.std
+from consul.base import Weight
 
 Check = consul.Check
 
@@ -348,6 +349,28 @@ def test_agent_register_enable_tag_override(self, consul_port):
         # Cleanup tasks
         c.agent.check.deregister('foo')
 
+    def test_agent_register_enable_weights(self, consul_port):
+        c = consul.Consul(port=consul_port)
+        index, nodes = c.health.service("foo1")
+        assert nodes == []
+
+        c.agent.service.register('foo', weights=Weight.weights(10, 10))
+        weight = {"Passing": 10, "Warning": 10}
+        assert c.agent.services()['foo']['Weights'] == weight
+        # Cleanup tasks
+        c.agent.check.deregister('foo')
+
+    def test_agent_register_disable_weights(self, consul_port):
+        c = consul.Consul(port=consul_port)
+        index, nodes = c.health.service("foo1")
+        assert nodes == []
+
+        c.agent.service.register('foo')
+        weight = {"Passing": 1, "Warning": 1}
+        assert c.agent.services()['foo']['Weights'] == weight
+        # Cleanup tasks
+        c.agent.check.deregister('foo')
+
     def test_agent_service_maintenance(self, consul_port):
         c = consul.Consul(port=consul_port)