Repository: getmoto/moto

Issue:
Cognito ID/Access Token do not contain user pool custom attributes
Hi,

I notice here that user pool [custom attributes](https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-settings-attributes.html#user-pool-settings-custom-attributes) are not injected into the JWT 

https://github.com/getmoto/moto/blob/624de34d82a1b2c521727b14a2173380e196f1d8/moto/cognitoidp/models.py#L526

As described in the [AWS Cognito docs](https://docs.aws.amazon.com/cognito/latest/developerguide/amazon-cognito-user-pools-using-the-id-token.html):

> "The ID token can also contain custom attributes that you define in your user pool. Amazon Cognito writes custom attribute values to the ID token as strings regardless of attribute type."

Would it be possible to add the custom attributes to this method?



Discussion:
Hi @davidtwomey, welcome to Moto! I'll raise a PR with this feature in a bit.

Patch:
diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -568,6 +568,9 @@ def add_custom_attributes(self, custom_attributes: List[Dict[str, str]]) -> None
     def create_id_token(self, client_id: str, username: str) -> Tuple[str, int]:
         extra_data = self.get_user_extra_data_by_client_id(client_id, username)
         user = self._get_user(username)
+        for attr in user.attributes:
+            if attr["Name"].startswith("custom:"):
+                extra_data[attr["Name"]] = attr["Value"]
         if len(user.groups) > 0:
             extra_data["cognito:groups"] = [group.group_name for group in user.groups]
         id_token, expires_in = self.create_jwt(


Test patch:
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -514,8 +514,8 @@ def test_add_custom_attributes_existing_attribute():
 def test_create_user_pool_default_id_strategy():
     conn = boto3.client("cognito-idp", "us-west-2")
 
-    first_pool = conn.create_user_pool(PoolName=str("default-pool"))
-    second_pool = conn.create_user_pool(PoolName=str("default-pool"))
+    first_pool = conn.create_user_pool(PoolName="default-pool")
+    second_pool = conn.create_user_pool(PoolName="default-pool")
 
     assert first_pool["UserPool"]["Id"] != second_pool["UserPool"]["Id"]
 
@@ -528,8 +528,8 @@ def test_create_user_pool_hash_id_strategy_with_equal_pool_name():
 
     conn = boto3.client("cognito-idp", "us-west-2")
 
-    first_pool = conn.create_user_pool(PoolName=str("default-pool"))
-    second_pool = conn.create_user_pool(PoolName=str("default-pool"))
+    first_pool = conn.create_user_pool(PoolName="default-pool")
+    second_pool = conn.create_user_pool(PoolName="default-pool")
 
     assert first_pool["UserPool"]["Id"] == second_pool["UserPool"]["Id"]
 
@@ -542,8 +542,8 @@ def test_create_user_pool_hash_id_strategy_with_different_pool_name():
 
     conn = boto3.client("cognito-idp", "us-west-2")
 
-    first_pool = conn.create_user_pool(PoolName=str("first-pool"))
-    second_pool = conn.create_user_pool(PoolName=str("second-pool"))
+    first_pool = conn.create_user_pool(PoolName="first-pool")
+    second_pool = conn.create_user_pool(PoolName="second-pool")
 
     assert first_pool["UserPool"]["Id"] != second_pool["UserPool"]["Id"]
 
@@ -557,7 +557,7 @@ def test_create_user_pool_hash_id_strategy_with_different_attributes():
     conn = boto3.client("cognito-idp", "us-west-2")
 
     first_pool = conn.create_user_pool(
-        PoolName=str("default-pool"),
+        PoolName="default-pool",
         Schema=[
             {
                 "Name": "first",
@@ -566,7 +566,7 @@ def test_create_user_pool_hash_id_strategy_with_different_attributes():
         ],
     )
     second_pool = conn.create_user_pool(
-        PoolName=str("default-pool"),
+        PoolName="default-pool",
         Schema=[
             {
                 "Name": "second",
@@ -1545,12 +1545,16 @@ def test_group_in_access_token():
 
 
 @mock_cognitoidp
-def test_group_in_id_token():
+def test_other_attributes_in_id_token():
     conn = boto3.client("cognito-idp", "us-west-2")
 
     username = str(uuid.uuid4())
     temporary_password = "P2$Sword"
-    user_pool_id = conn.create_user_pool(PoolName=str(uuid.uuid4()))["UserPool"]["Id"]
+    user_pool_id = conn.create_user_pool(
+        PoolName=str(uuid.uuid4()),
+        Schema=[{"Name": "myattr", "AttributeDataType": "String"}],
+    )["UserPool"]["Id"]
+
     user_attribute_name = str(uuid.uuid4())
     user_attribute_value = str(uuid.uuid4())
     group_name = str(uuid.uuid4())
@@ -1566,7 +1570,10 @@ def test_group_in_id_token():
         UserPoolId=user_pool_id,
         Username=username,
         TemporaryPassword=temporary_password,
-        UserAttributes=[{"Name": user_attribute_name, "Value": user_attribute_value}],
+        UserAttributes=[
+            {"Name": user_attribute_name, "Value": user_attribute_value},
+            {"Name": "custom:myattr", "Value": "some val"},
+        ],
     )
 
     conn.admin_add_user_to_group(
@@ -1596,6 +1603,7 @@ def test_group_in_id_token():
 
     claims = jwt.get_unverified_claims(result["AuthenticationResult"]["IdToken"])
     assert claims["cognito:groups"] == [group_name]
+    assert claims["custom:myattr"] == "some val"
 
 
 @mock_cognitoidp