Add TaskRouter Grant

wanjunsli · wanjunsli · commit ff1c2b31f947 · 2017-03-02T14:03:39.000-08:00
diff --git a/tests/unit/jwt/test_access_token.py b/tests/unit/jwt/test_access_token.py
@@ -10,7 +10,8 @@
     IpMessagingGrant,
     SyncGrant,
     VoiceGrant,
-    VideoGrant
+    VideoGrant,
+    TaskRouterGrant
 )
 
 ACCOUNT_SID = 'AC123'
@@ -174,6 +175,27 @@ def test_programmable_voice_grant(self):
             }
         }, decoded_token.payload['grants']['voice'])
 
+    def test_task_router_grant(self):
+        grant = TaskRouterGrant(
+            workspace_sid='WS123',
+            worker_sid='WK123',
+            role='worker'
+        )
+
+        scat = AccessToken(ACCOUNT_SID, SIGNING_KEY_SID, 'secret')
+        scat.add_grant(grant)
+
+        token = scat.to_jwt()
+        assert_is_not_none(token)
+        decoded_token = AccessToken.from_jwt(token, 'secret')
+        self._validate_claims(decoded_token.payload)
+        assert_equal(1, len(decoded_token.payload['grants']))
+        assert_equal({
+            'workspace_sid': 'WS123',
+            'worker_sid': 'WK123',
+            'role': 'worker'
+        }, decoded_token.payload['grants']['task_router'])
+
     def test_pass_grants_in_constructor(self):
         grants = [
             ConversationsGrant(),
diff --git a/twilio/jwt/access_token/grants.py b/twilio/jwt/access_token/grants.py
@@ -118,3 +118,26 @@ def to_payload(self):
             grant['configuration_profile_sid'] = self.configuration_profile_sid
 
         return grant
+
+
+class TaskRouterGrant(AccessTokenGrant):
+    """Grant to access Twilio TaskRouter"""
+    def __init__(self, workspace_sid=None, worker_sid=None, role=None):
+        self.workspace_sid = workspace_sid
+        self.worker_sid = worker_sid
+        self.role = role
+
+    @property
+    def key(self):
+        return "task_router"
+
+    def to_payload(self):
+        grant = {}
+        if self.workspace_sid:
+            grant['workspace_sid'] = self.workspace_sid
+        if self.worker_sid:
+            grant['worker_sid'] = self.worker_sid
+        if self.role:
+            grant['role'] = self.role
+
+        return grant
