From f47080a4202f8de6efdadf9e17d0e6aeb705a157 Mon Sep 17 00:00:00 2001
From: Bert Blommers <info@bertblommers.nl>
Date: Sat, 13 Jan 2024 23:05:45 +0000
Subject: [PATCH] CognitoIDP: ID-token now exposes custom attributes (#7212)

---
 moto/cognitoidp/models.py                |  3 +++
 tests/test_cognitoidp/test_cognitoidp.py | 30 +++++++++++++++---------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
index 48462be09..0e16230e6 100644
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -568,6 +568,9 @@ class CognitoIdpUserPool(BaseModel):
     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(
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
index ed2851096..6df344398 100644
--- 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