Mocking Cognito in Python (at no cost)

Amazon Cognito is an AWS service for authentication, authorization, and user management for web and mobile apps. It uses JWT tokens as a means of user authentication.

Emulating/mocking up real services is a no-brainer, no one wants tests to affect the production environment in any way. In addition, it allows us to make tests completely independent from the external environment.

There is an excellent tool for emulating AWS applications or Lambdas called localstack. Nevertheless, some service mocks are paid, and as of today Amazon Cognito is such a service.

Moreover, localstack needs to be run as a single container.

This article focuses on finding an alternative to Amazon Cognito mocking that would be free and simpler to use in unit-tests.

Here is where “moto” comes in

Moto is a Python package designed for mocking AWS services. It’s completely free and, unlike localstack, it doesn’t require setting up an additional container, it can simply be used as a Python package, which is a much more convenient solution for unit-tests.

It looks like this is something we are looking for, let’s jump to the code then!

Installing package

Before we start working with the code, it is necessary to install the moto package.

This can be done, e.g. with the pip tool (what a surprise ( ͡° ͜ʖ ͡°) ), like below:

pip install moto

At this moment, the 4.0.1 version has been installed.

That is all, we can use Amazon Cognito mock now.

Bit of a code

The mock we are looking for is mock_cognitoidp.

It’s enough to import it from the moto library and decorate a particular test.

from moto import mock_cognitoidp


@mock_cognitoidp
def test_cognito():
   pass

That is all that has to be done to simulate the functionality of Amazon Cognito.

Real life example

I now turn to an example from real life.

For the purposes of this article, I need to assume a couple of things:

• developer is familiar with JWT tokens,
• developer is familiar with Amazon Cognito JWT validation process,
• developer is familiar with concept of user pool,
• the library used to make use of AWS services is boto3.

In the example below, I will go through the process of testing the validation of the access token received from Amazon Cognito.

To make it simpler, I will not use any fixtures or setUp methods.

Enough idle talk, let’s declare a test case with the following steps:

1. create and set up user pool,
2. get JSON Web Key Set,
3. create Amazon Cognito user,
4. get access token for created user,
5. validate received access token.

Most test runners look for test prefix in the function name, so let’s follow the rules:

def test_cognito_authorization_process():

Lest we forget, let’s mock Amazon Cognito from the beginning, just so we don’t connect to the real AWS service.

from moto import mock_cognitoidp


@mock_cognitoidp
def test_cognito_authorization_process():

Next step would be to declare credentials for the user to be created on Amazon Cognito.

I recommend using faker or hypothesis to generate fake test data, but for the purpose of the example, it will be simply hardcoded.

@mock_cognitoidp
def test_cognito_authorization_process():
     region = "us-west-1"
     username = "test_username"
     password = "SecurePassword1234#$%"  # Password must meet security policies.
     email = "test_mail@test.com"

Now we need a client to connect with the Amazon Cognito. As mentioned in assumptions, for that purpose boto3 will be used.

import boto3


cognito_client = boto3.client("cognito-idp", region_name=region)

Any communication with Amazon Cognito will be done via `cognito_client` object.

Next, we need to create a user pool, which is simply a users’ directory on Amazon Cognito.

user_pool_id = cognito_client.create_user_pool(PoolName="TestUserPool")["UserPool"]["Id"]

User pool client is also necessary for the authorization process.

app_client = cognito_client.create_user_pool_client(
    UserPoolId=user_pool_id, ClientName="TestAppClient"
)

That’s all that is needed to register a user on Amazon Cognito. We can proceed from here now:

cognito_client.sign_up(
	ClientId=app_client["UserPoolClient"]["ClientId"],
	Username=username,
	Password=password,
	UserAttributes=[
    	    {"Name": "email", "Value": email},
	],
)

By default, Amazon Cognito requires the user to be activated. In this test case, it will be simulated via admin manual confirmation.

cognito_client.admin_confirm_sign_up(UserPoolId=user_pool_id, Username=username)

The setup of the environment is done, it’s time to execute the actual process being tested, i.e. retrieve the access token and get claims from it.

Amazon Cognito mock is running locally, so it is necessary to use a trick in order to initialize a user authorization process. It’s necessary to use the admin_initiate_auth method and not initiate_auth. This is because initiate_auth is a client/browser side API call, whereas admin_initiate_auth is meant to be run on the server side.

access_token = cognito_client.admin_initiate_auth(
	UserPoolId=user_pool_id,
	ClientId=app_client["UserPoolClient"]["ClientId"],
	AuthFlow="ADMIN_NO_SRP_AUTH",
	AuthParameters={"USERNAME": username, "PASSWORD": password},
)["AuthenticationResult"]["AccessToken"]

Extracting claims from the access token can be fully implemented in many ways, e.g. with the jose package. No less important is the prior validation of the data contained in the claims, so it is good practice to follow the process described in Amazon Cognito’s official documentation.

import requests


def fetch_public_keys(region: str, user_pool_id: str) -> dict:
	keys_url = f"https://cognito-idp.{region}.amazonaws.com/{user_pool_id}/.well-known/jwks.json"
	response = requests.get(keys_url).json()
	return response["keys"]


def validate_token(token: str, keys: dict, valid_iss: str) -> dict:
	"""
	For now, let's return some fake claims data.
	"""

	return {
    	"sub": "9cc71104-d2b8-4a84-8269-a55b95f5bd23",
    	"iss": "https://cognito-idp.us-west-1.amazonaws.com/us-west-1_4ujjMwsfK",
    	"client_id": "4pon8gi9t6f6dllka6ap2ihcad",
    	"origin_jti": "65e3bb61-b3ba-4156-8ebd-ae6db607310f",
    	"event_id": "2e63932a-3583-4265-949c-05f52bb56467",
    	"token_use": "access",
    	"auth_time": 1662022792,
    	"exp": 1662026392,
    	"iat": 1662022792,
    	"jti": "d5d3f3d9-c02d-41e7-a9da-4443278d61cf",
    	"username": "test_username",
	}




	user_pool_jwks_keys = fetch_public_keys(region=region, user_pool_id=user_pool_id)
	user_pool_valid_iss = f"https://cognito-idp.us-west-1.amazonaws.com/{user_pool_id}"

	claims = validate_token(
    	token=access_token, keys=user_pool_jwks_keys, valid_iss=user_pool_valid_iss
	)

The final step is to assert actual behavior with the expected one.

 assert claims["username"] == username

And that’s it, together we went through a test case which set up the Amazon Cognito environment, created the user, and – at the very end – validated the data obtained from the access token, and all this without using the real Amazon Cognito service.

Here is the entire example on GitHub gist. 

There are also odd issues

Unfortunately, we noticed two odd issues with mock_cognitoidp 4.0.1 version.

The first problem is that the mock returns the email instead of the username in token claims.

It’s a discrepancy in claims returned by the mock and those assumed by Amazon Cognito documentation.

Recently I’ve created an issue for that bug on moto github: https://github.com/spulec/moto/issues/5279.

Temporarily, before this problem is officially fixed, it can be hidden with unittest mock. More specifically, we can patch the claims returned by our custom function to validate the token and return validated claims to those given in the documentation.

The second problem is that the function to fetch public keys must use the requests package to be mocked by mock_cognitoidp.

More details can be found here: https://github.com/spulec/moto/issues/4601.

An alternative solution to this problem would be to implement a function, only for test cases, which uses a requests package to retrieve the public keys, but only when the production version of this function uses a different package than requests.

Conclusion

Amazon Cognito is a major player in the world of authorization, so it is important to find a way to properly test the integration of the application under development with this service. Both localstack and moto have their advantages and disadvantages. It’s crucial to choose the appropriate package for our needs since it may turn out that we do not need all the localstack features and moto is entirely sufficient.