Changed how verify_token brings in env var, fixed verify_token test

gdm · Apr 24, 2017 · f22f9ca · f22f9ca
1 parent f5e1fc3
commit f22f9ca
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/warrant/__init__.py b/warrant/__init__.py
@@ -95,8 +95,8 @@ def get_keys(self):
             return self.pool_jwk
         except AttributeError:
             #Check for the dictionary in environment variables.
-            pool_jwk_env = env('COGNITO_JWKS', var_type='dict')
-            if isinstance(pool_jwk_env, dict):
+            pool_jwk_env = env('COGNITO_JWKS', {},var_type='dict')
+            if len(pool_jwk_env.keys()) > 0:
                 self.pool_jwk = pool_jwk_env
                 return self.pool_jwk
             #If it is not there use the requests library to get it
@@ -251,9 +251,9 @@ def authenticate(self, password):
         aws = AWSSRP(username=self.username, password=password, pool_id=self.user_pool_id,
                      client_id=self.client_id, client=self.client)
         tokens = aws.authenticate_user()
-        self.verify_token(tokens['AuthenticationResult']['IdToken'],'id_token')
+        self.verify_token(tokens['AuthenticationResult']['IdToken'],'id_token','id')
         self.refresh_token = tokens['AuthenticationResult']['RefreshToken']
-        self.verify_token(tokens['AuthenticationResult']['AccessToken'], 'access_token')
+        self.verify_token(tokens['AuthenticationResult']['AccessToken'], 'access_token','access')
         self.token_type = tokens['AuthenticationResult']['TokenType']
 
     def logout(self):

diff --git a/warrant/tests/tests.py b/warrant/tests/tests.py
@@ -4,7 +4,7 @@
 from envs import env
 from botocore.exceptions import ClientError
 
-from warrant import Cognito,UserObj
+from warrant import Cognito, UserObj, TokenVerificationException
 from warrant.aws_srp import AWSSRP
 
 
@@ -53,8 +53,9 @@ def test_authenticate(self):
     def test_verify_token(self):
         self.user.authenticate(self.password)
         bad_access_token = '{}wrong'.format(self.user.access_token)
-        verify = self.user.verify_token(bad_access_token,'access_token')
-        self.assertFalse(verify)
+
+        with self.assertRaises(TokenVerificationException) as vm:
+            self.user.verify_token(bad_access_token, 'access_token', 'access')
 
     def test_logout(self):
         self.user.authenticate(self.password)