add acl token support to agent.check.register and agent.service.register

cablehead · cablehead · commit 1a039801d676 · 2015-10-21T13:34:13.000-07:00
diff --git a/consul/base.py b/consul/base.py
@@ -602,6 +602,7 @@ def register(
                     port=None,
                     tags=None,
                     check=None,
+                    token=None,
                     # *deprecated* use check parameter
                     script=None,
                     interval=None,
@@ -625,6 +626,10 @@ def register(
                 An optional health *check* can be created for this service is
                 one of `Check.script`_, `Check.http`_, or `Check.ttl`_.
 
+                *token* is an optional `ACL token`_ to apply to this request.
+                Note this call will return successful even if the token doesn't
+                have permissions to register this service.
+
                 *script*, *interval*, *ttl*, *http*, and *timeout* arguments
                 are deprecated. use *check* instead.
                 """
@@ -649,9 +654,15 @@ def register(
                         http=http,
                         timeout=timeout))
 
+                params = {}
+                token = token or self.agent.token
+                if token:
+                    params['token'] = token
+
                 return self.agent.http.put(
                     lambda x: x.code == 200,
                     '/v1/agent/service/register',
+                    params=params,
                     data=json.dumps(payload))
 
             def deregister(self, service_id):
@@ -701,6 +712,7 @@ def register(
                     check_id=None,
                     notes=None,
                     service_id=None,
+                    token=None,
                     # *deprecated* use check parameter
                     script=None,
                     interval=None,
@@ -726,6 +738,10 @@ def register(
                 Optionally, a *service_id* can be specified to associate a
                 registered check with an existing service.
 
+                *token* is an optional `ACL token`_ to apply to this request.
+                Note this call will return successful even if the token doesn't
+                have permissions to register this check.
+
                 *script*, *interval*, *ttl*, *http*, and *timeout* arguments
                 are deprecated. use *check* instead.
 
@@ -754,9 +770,15 @@ def register(
                 if service_id:
                     payload['serviceid'] = service_id
 
+                params = {}
+                token = token or self.agent.token
+                if token:
+                    params['token'] = token
+
                 return self.agent.http.put(
                     lambda x: x.code == 200,
                     '/v1/agent/check/register',
+                    params=params,
                     data=json.dumps(payload))
 
             def deregister(self, check_id):
diff --git a/tests/consul.osx b/tests/consul.osx
diff --git a/tests/test_std.py b/tests/test_std.py
@@ -680,6 +680,12 @@ def test_acl_explict_token_use(self, acl_consul):
             key "private/" {
                 policy = "deny"
             }
+            service "foo-" {
+                policy = "write"
+            }
+            service "bar-" {
+                policy = "read"
+            }
         """
 
         token = c.acl.create(rules=rules, token=master_token)
@@ -712,7 +718,17 @@ def test_acl_explict_token_use(self, acl_consul):
             consul.ACLPermissionDenied,
             c.kv.delete, 'private/foo', token=token)
 
+        # test token pass through for service registration
+        c.agent.service.register("bar-1", token=token)
+        c.agent.service.register("foo-1", token=token)
+        index, data = c.health.service('foo-1')
+        assert data
+        index, data = c.health.service('bar-1')
+        assert not data
+
         # clean up
+        assert c.agent.service.deregister('foo-1') is True
+        assert c.agent.service.deregister('bar-1') is True
         c.acl.destroy(token, token=master_token)
         acls = c.acl.list(token=master_token)
         assert set([x['ID'] for x in acls]) == \
