feat(auth) mock id_token_from_cognito in tests

ZXYmania · ZXYmania · commit 50798b400fdb · 2020-05-08T15:39:35.000+10:00
diff --git a/Pipfile b/Pipfile
diff --git a/examples/test.py b/examples/test.py
diff --git a/staxapp/auth.py b/staxapp/auth.py
@@ -8,6 +8,7 @@
 from aws_requests_auth.aws_auth import AWSRequestsAuth
 from botocore import UNSIGNED
 from botocore.client import Config as BotoConfig
+from botocore.exceptions import ClientError
 from warrant import AWSSRP, Cognito
 
 from staxapp.config import Config as JumaConfig
@@ -25,9 +26,13 @@ def __init__(self, config_branch):
 
     def requests_auth(self, username, password):
         if username is None:
-            raise InvalidCredentialsException("Please provide an Access Key to your config")
+            raise InvalidCredentialsException(
+                "Please provide an Access Key to your config"
+            )
         if password is None:
-            raise InvalidCredentialsException("Please provide a Secret Key to your config")
+            raise InvalidCredentialsException(
+                "Please provide a Secret Key to your config"
+            )
 
         id_token = self.id_token_from_cognito(username, password)
         id_creds = self.sts_from_cognito_identity_pool(id_token)
@@ -38,14 +43,15 @@ def requests_auth(self, username, password):
 
         return JumaConfig.auth
 
-    def id_token_from_cognito(self, username=None, password=None):
+    def id_token_from_cognito(self, username=None, password=None, client=None):
         token = None
         if username and password:
-            client = boto3.client(
-                "cognito-idp",
-                region_name=self.aws_region,
-                config=BotoConfig(signature_version=UNSIGNED),
-            )
+            if not client:
+                client = boto3.client(
+                    "cognito-idp",
+                    region_name=self.aws_region,
+                    config=BotoConfig(signature_version=UNSIGNED),
+                )
             aws = AWSSRP(
                 username=username,
                 password=password,
@@ -55,16 +61,21 @@ def id_token_from_cognito(self, username=None, password=None):
             )
             try:
                 tokens = aws.authenticate_user()
-            except client.exceptions.NotAuthorizedException as e:
-                logging.error(e)
-                raise InvalidCredentialsException(message=str(e), detail="Please check your Secret Key is correct")
-            except client.exceptions.UserNotFoundException as e:
-                raise InvalidCredentialsException(message=str(e), detail="Please check your Access Key, that you have created your Api Token and that you are using the right STAX REGION")
-
-            # logging.debug(f"TOKEN: {tokens}")
+            except ClientError as e:
+                if e.response["Error"]["Code"] == "UserNotFoundException":
+                    raise InvalidCredentialsException(
+                        message=str(e), detail="Please check your Secret Key is correct"
+                    )
+                elif e.response["Error"]["Code"] == "NotAuthorizedException":
+                    raise InvalidCredentialsException(
+                        message=str(e),
+                        detail="Please check your Access Key, that you have created your Api Token and that you are using the right STAX REGION",
+                    )
+                else:
+                    raise InvalidCredentialsException(
+                        f"Unexpected Client Error. Error details: {e}"
+                    )
             token = tokens["AuthenticationResult"]["IdToken"]
-        else:
-            token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256")
         return token
 
     def sts_from_cognito_identity_pool(self, token, cognito_client=None):
diff --git a/staxapp/exceptions.py b/staxapp/exceptions.py
@@ -30,6 +30,7 @@ def __init__(self, message):
 
 class InvalidCredentialsException(Exception):
     def __init__(self, message, detail=""):
+        logging.error(message)
         prefix = f"InvalidCredentialsException: "
         if detail:
             prefix = f"{prefix}{detail} - "
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -209,8 +209,8 @@ def testFailedApiException(self):
         )
         with self.assertRaises(ApiException):
             self.Api.get("/test/no/error")
-        
-		# Test an exception which has no json in response
+
+        # Test an exception which has no json in response
         responses.add(
             responses.GET,
             f"{Config.api_base_url()}/test/invalid/json",
@@ -219,16 +219,14 @@ def testFailedApiException(self):
         )
         with self.assertRaises(ApiException):
             self.Api.get("/test/invalid/json")
-		
+
         # Test an exception with no content
         responses.add(
-            responses.GET,
-            f"{Config.api_base_url()}/test/no/content",
-            status=500,
+            responses.GET, f"{Config.api_base_url()}/test/no/content", status=500,
         )
         with self.assertRaises(ApiException):
             self.Api.get("/test/no/content")
-        
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -12,11 +12,12 @@
 
 from botocore import UNSIGNED
 from botocore.client import Config as BotoConfig
-from botocore.stub import Stubber
+from botocore.stub import Stubber, ANY
 
 from staxapp.auth import StaxAuth
 from staxapp.exceptions import InvalidCredentialsException
 
+
 class StaxAuthTests(unittest.TestCase):
     """
     Inherited class to run all unit tests for this module
@@ -30,8 +31,16 @@ def setUp(self):
         )
         self.cognito_stub = Stubber(self.cognito_client)
 
+        self.aws_srp_client = botocore.session.get_session().create_client(
+            "cognito-idp",
+            region_name="ap-southeast-2",
+            config=BotoConfig(signature_version=UNSIGNED),
+        )
+        self.aws_srp_stubber = Stubber(self.aws_srp_client)
+
     def tearDown(self):
         self.cognito_stub.deactivate()
+        self.aws_srp_stubber.deactivate()
 
     def testStaxAuthInit(self):
         """
@@ -45,16 +54,58 @@ def testToken(self):
         Test valid JWT is returned
         """
         sa = StaxAuth("ApiAuth")
-        token = sa.id_token_from_cognito()
-        jwt_token = jwt.decode(token, verify=False)
-        self.assertIn("sub", jwt_token)
+        self.stub_aws_srp(sa, "valid_username")
+        token = sa.id_token_from_cognito(
+            username="valid_username", password="correct", client=self.aws_srp_client
+        )
+        self.assertEqual(token, "valid_token")
+
+    def testCredentialErrors(self):
+        """
+        Test that boto errors are caught and converted to InvalidCredentialExceptions
+      """
+
+        sa = StaxAuth("ApiAuth")
+        # Test with invalid username password
+        self.stub_aws_srp(sa, "bad_password", "UserNotFoundException")
+        user_not_found_success = False
+        try:
+            sa.id_token_from_cognito(
+                username="bad_password", password="wrong", client=self.aws_srp_client
+            )
+        except InvalidCredentialsException as e:
+            self.assertIn("Please check your Secret Key is correct", e.message)
+            user_not_found_success = True
+        self.assertTrue(user_not_found_success)
+
+        # Test with no access
+        self.stub_aws_srp(sa, "no_access", "NotAuthorizedException")
+        no_access_success = False
+        try:
+            sa.id_token_from_cognito(
+                username="no_access", password="wrong", client=self.aws_srp_client
+            )
+        except InvalidCredentialsException as e:
+            self.assertIn(
+                "Please check your Access Key, that you have created your Api Token and that you are using the right STAX REGION",
+                e.message,
+            )
+            no_access_success = True
+        self.assertTrue(no_access_success)
+
+        # Test Unknown Error
+        self.stub_aws_srp(sa, "Unknown", "UnitTesting")
+        with self.assertRaises(InvalidCredentialsException):
+            sa.id_token_from_cognito(
+                username="Unknown", password="wrong", client=self.aws_srp_client
+            )
 
     def testCreds(self):
         """
         Test valid credentials are returned
         """
         sa = StaxAuth("ApiAuth")
-        token = sa.id_token_from_cognito()
+        token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256")
         jwt_token = jwt.decode(token, verify=False)
         self.stub_cognito_creds(jwt_token.get("sub"))
         creds = sa.sts_from_cognito_identity_pool(
@@ -63,19 +114,55 @@ def testCreds(self):
         self.assertIn("Credentials", creds)
         self.assertTrue(creds.get("IdentityId").startswith("ap-southeast-2"))
 
-    def testCredentialErrors(self):
-      """
+    def testAuthErrors(self):
+        """
       Test that errors are thrown when keys are invalid
       """
-      sa = StaxAuth("ApiAuth")
-      # Test with no username
-      with self.assertRaises(InvalidCredentialsException):
-        sa.requests_auth(username=None, password='valid')
-      # Test with no username
-      with self.assertRaises(InvalidCredentialsException):
-        sa.requests_auth(username='valid', password=None)
-      # Test with invalid username password
-      # Todo
+        sa = StaxAuth("ApiAuth")
+        # Test with no username
+        with self.assertRaises(InvalidCredentialsException):
+            sa.requests_auth(username=None, password="valid")
+
+        # Test with no username
+        with self.assertRaises(InvalidCredentialsException):
+            sa.requests_auth(username="valid", password=None)
+
+    def stub_aws_srp(self, stax_auth, username, error_code=None):
+        expected_parameters = {
+            "AuthFlow": "USER_SRP_AUTH",
+            "AuthParameters": {"SRP_A": ANY, "USERNAME": username},
+            "ClientId": stax_auth.client_id,
+        }
+        if error_code:
+            self.aws_srp_stubber.add_client_error(
+                "initiate_auth",
+                service_error_code=error_code,
+                expected_params=expected_parameters,
+            )
+        else:
+            self.aws_srp_stubber.add_response(
+                "initiate_auth",
+                {
+                    "ChallengeParameters": {
+                        "USER_ID_FOR_SRP": "user",
+                        "SALT": "4",
+                        "SRP_B": "5",
+                        "SECRET_BLOCK": "secblock",
+                    },
+                    "ChallengeName": "PASSWORD_VERIFIER",
+                },
+                expected_parameters,
+            )
+            self.aws_srp_stubber.add_response(
+                "respond_to_auth_challenge",
+                {"AuthenticationResult": {"IdToken": "valid_token"},},
+                {
+                    "ClientId": stax_auth.client_id,
+                    "ChallengeName": ANY,
+                    "ChallengeResponses": ANY,
+                },
+            )
+        self.aws_srp_stubber.activate()
 
     def stub_cognito_creds(self, token: str):
         sa = StaxAuth("ApiAuth")
