Merge pull request #32 from stax-labs/feat/test-clients

chore(tests): Add aws client tests

Author: Mat Lord

staxapp/api.py

@@ -8,10 +8,10 @@ class Api:
     _requests_auth = None
 
     @classmethod
-    def _auth(cls):
+    def _auth(cls, **kwargs):
         if not cls._requests_auth:
             cls._requests_auth = Config.get_auth_class().requests_auth(
-                Config.access_key, Config.secret_key
+                Config.access_key, Config.secret_key, **kwargs
             )
         return cls._requests_auth
 

staxapp/auth.py

@@ -83,18 +83,23 @@ def sts_from_cognito_identity_pool(self, token, cognito_client=None, **kwargs):
                 region_name=self.aws_region,
                 config=BotoConfig(signature_version=UNSIGNED),
             )
-        id = cognito_client.get_id(
-            IdentityPoolId=self.identity_pool,
-            Logins={
-                f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
-            },
-        )
-        id_creds = cognito_client.get_credentials_for_identity(
-            IdentityId=id["IdentityId"],
-            Logins={
-                f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
-            },
-        )
+        try:
+            id = cognito_client.get_id(
+                IdentityPoolId=self.identity_pool,
+                Logins={
+                    f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
+                },
+            )
+            id_creds = cognito_client.get_credentials_for_identity(
+                IdentityId=id["IdentityId"],
+                Logins={
+                    f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
+                },
+            )
+        except ClientError as e:
+            raise InvalidCredentialsException(
+                f"Unexpected Client Error. Error details: {e}"
+            )
         return id_creds
 
     def sigv4_signed_auth_headers(self, id_creds):

staxapp/contract.py

@@ -54,7 +54,7 @@ def validate(cls, data, component):
     @staticmethod
     def default_swagger_template() -> dict:
         # Get the default swagger template from https://api.au1.staxapp.cloud/20190206/public/api-document
-        schema = requests.get(Config.schema_url()).json()
+        schema_response = requests.get(Config.schema_url()).json()
         template = dict(
             openapi="3.0.0",
             info={
@@ -75,7 +75,7 @@ def default_swagger_template() -> dict:
                         "x-amazon-apigateway-authtype": "awsSigv4",
                     }
                 },
-                "schemas": schema,
+                "schemas": schema_response.get("components").get("schemas"),
                 "responses": dict(),
                 "requestBodies": dict(),
             },

tests/test_api.py

@@ -207,8 +207,10 @@ def testFailedApiException(self):
             json=response_dict,
             status=500,
         )
-        with self.assertRaises(ApiException):
+        try:
             self.Api.get("/test/no/error")
+        except ApiException as e:
+            self.assertIn("Api Exception", str(e))
 
         # Test an exception which has no json in response
         responses.add(

tests/test_auth.py

@@ -16,6 +16,7 @@
 from botocore.stub import Stubber, ANY
 from datetime import datetime, timedelta, timezone
 
+from staxapp.api import Api
 from staxapp.auth import StaxAuth, ApiTokenAuth, RootAuth
 from staxapp.config import Config
 from staxapp.exceptions import InvalidCredentialsException
@@ -59,10 +60,20 @@ def testToken(self):
         sa = StaxAuth("ApiAuth")
         self.stub_aws_srp(sa, "valid_username")
         token = sa.id_token_from_cognito(
-            username="valid_username", password="correct", srp_client=self.aws_srp_client
+            username="valid_username",
+            password="correct",
+            srp_client=self.aws_srp_client,
         )
         self.assertEqual(token, "valid_token")
 
+    def testTokenClient(self):
+        """
+        Test the AWSSRP client is invoked and throws an error
+        """
+        sa = StaxAuth("ApiAuth")
+        with self.assertRaises(InvalidCredentialsException):
+            sa.id_token_from_cognito(username="username", password="password")
+
     def testCredentialErrors(self):
         """
         Test that boto errors are caught and converted to InvalidCredentialExceptions
@@ -73,10 +84,12 @@ def testCredentialErrors(self):
         user_not_found_success = False
         try:
             sa.id_token_from_cognito(
-                username="bad_password", password="wrong", srp_client=self.aws_srp_client
+                username="bad_password",
+                password="wrong",
+                srp_client=self.aws_srp_client,
             )
         except InvalidCredentialsException as e:
-            self.assertIn("Please check your Secret Key is correct", e.message)
+            self.assertIn("Please check your Secret Key is correct", str(e))
             user_not_found_success = True
         self.assertTrue(user_not_found_success)
 
@@ -90,7 +103,7 @@ def testCredentialErrors(self):
         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,
+                str(e),
             )
             no_access_success = True
         self.assertTrue(no_access_success)
@@ -116,6 +129,16 @@ def testCreds(self):
         self.assertIn("Credentials", creds)
         self.assertTrue(creds.get("IdentityId").startswith("ap-southeast-2"))
 
+    def testCredsClient(self):
+        """
+        Test the cognito client is invoked and throws an error
+        """
+        sa = StaxAuth("ApiAuth")
+        token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256")
+        jwt_token = jwt.decode(token, verify=False)
+        with self.assertRaises(InvalidCredentialsException):
+            sa.sts_from_cognito_identity_pool(jwt_token.get("sub"))
+
     def testAuthErrors(self):
         """
       Test that errors are thrown when keys are invalid
@@ -289,6 +312,28 @@ def testRootAuth(self):
         )
         self.assertIsNotNone(StaxConfig.auth)
 
+    def testApiAuth(self):
+        """
+        Test auth through the Api class
+        """
+        sa = StaxAuth("ApiAuth")
+        StaxConfig = Config
+        StaxConfig.expiration = None
+        StaxConfig.access_key = "username"
+        StaxConfig.secret_key = "password"
+
+        token = jwt.encode({"sub": "valid_token"}, "secret", algorithm="HS256")
+        jwt_token = jwt.decode(token, verify=False)
+
+        self.stub_cognito_creds(sa, jwt_token.get("sub"))
+        self.stub_aws_srp(sa, "username")
+
+        Api._requests_auth = None
+        Api._auth(
+            srp_client=self.aws_srp_client, cognito_client=self.cognito_client,
+        )
+        self.assertIsNotNone(Api._requests_auth)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_client.py

@@ -36,6 +36,13 @@ def testStaxClient(self):
         self.assertTrue(client._initialized)
         self.assertTrue(client._admin)
 
+    def testInvalidStaxClient(self):
+        """
+        Test an invalid Api class raises an error
+        """
+        with self.assertRaises(ValidationException):
+            StaxClient("fake")
+
     def testLoadLiveSchema(self):
         """
         Test loading live schema

tests/test_contract.py

@@ -19,6 +19,17 @@ class StaxContractTests(unittest.TestCase):
     def setUp(self):
         self.StaxContract = StaxContract
 
+    def testInvalidSchema(self):
+        """
+        Test an error is thrown if no schema is found
+        """
+        sc = StaxContract
+        sc._swagger_doc = None
+        data = {"Name": "Unit", "AccountType": "Test"}
+        component = "accounts.CreateAccount"
+        sc.validate(data, component)
+        self.assertIsNotNone(sc._swagger_doc)
+
     def testDefaultSchema(self):
         """
         Test the default schema is valid