Repository: getmoto/moto (python)

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?


Pull request description:
CognitoIDP: ID-token now exposes custom attributes
Fixes #7209 

Patch:
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/moto/sagemakerruntime/models.py b/moto/sagemakerruntime/models.py
index 26928bcf8..872ca4464 100644
--- a/moto/sagemakerruntime/models.py
+++ b/moto/sagemakerruntime/models.py
@@ -1,6 +1,8 @@
+import json
 from typing import Dict, List, Tuple
 
 from moto.core import BackendDict, BaseBackend
+from moto.moto_api._internal import mock_random as random
 
 
 class SageMakerRuntimeBackend(BaseBackend):
@@ -8,6 +10,7 @@ class SageMakerRuntimeBackend(BaseBackend):
 
     def __init__(self, region_name: str, account_id: str):
         super().__init__(region_name, account_id)
+        self.async_results: Dict[str, Dict[str, str]] = {}
         self.results: Dict[str, Dict[bytes, Tuple[str, str, str, str]]] = {}
         self.results_queue: List[Tuple[str, str, str, str]] = []
 
@@ -62,5 +65,37 @@ class SageMakerRuntimeBackend(BaseBackend):
             )
         return self.results[endpoint_name][unique_repr]
 
+    def invoke_endpoint_async(self, endpoint_name: str, input_location: str) -> str:
+        if endpoint_name not in self.async_results:
+            self.async_results[endpoint_name] = {}
+        if input_location in self.async_results[endpoint_name]:
+            return self.async_results[endpoint_name][input_location]
+        if self.results_queue:
+            body, _type, variant, attrs = self.results_queue.pop(0)
+        else:
+            body = "body"
+            _type = "content_type"
+            variant = "invoked_production_variant"
+            attrs = "custom_attributes"
+        json_data = {
+            "Body": body,
+            "ContentType": _type,
+            "InvokedProductionVariant": variant,
+            "CustomAttributes": attrs,
+        }
+        output = json.dumps(json_data).encode("utf-8")
+
+        output_bucket = f"sagemaker-output-{random.uuid4()}"
+        output_location = "response.json"
+        from moto.s3.models import s3_backends
+
+        s3_backend = s3_backends[self.account_id]["global"]
+        s3_backend.create_bucket(output_bucket, region_name=self.region_name)
+        s3_backend.put_object(output_bucket, output_location, value=output)
+        self.async_results[endpoint_name][
+            input_location
+        ] = f"s3://{output_bucket}/{output_location}"
+        return self.async_results[endpoint_name][input_location]
+
 
 sagemakerruntime_backends = BackendDict(SageMakerRuntimeBackend, "sagemaker-runtime")
diff --git a/moto/sagemakerruntime/responses.py b/moto/sagemakerruntime/responses.py
index ee1b86608..c30c4b2b6 100644
--- a/moto/sagemakerruntime/responses.py
+++ b/moto/sagemakerruntime/responses.py
@@ -3,6 +3,7 @@ import json
 
 from moto.core.common_types import TYPE_RESPONSE
 from moto.core.responses import BaseResponse
+from moto.moto_api._internal import mock_random as random
 
 from .models import SageMakerRuntimeBackend, sagemakerruntime_backends
 
@@ -43,3 +44,14 @@ class SageMakerRuntimeResponse(BaseResponse):
         if custom_attributes:
             headers["X-Amzn-SageMaker-Custom-Attributes"] = custom_attributes
         return 200, headers, body
+
+    def invoke_endpoint_async(self) -> TYPE_RESPONSE:
+        endpoint_name = self.path.split("/")[2]
+        input_location = self.headers.get("X-Amzn-SageMaker-InputLocation")
+        inference_id = self.headers.get("X-Amzn-SageMaker-Inference-Id")
+        location = self.sagemakerruntime_backend.invoke_endpoint_async(
+            endpoint_name, input_location
+        )
+        resp = {"InferenceId": inference_id or str(random.uuid4())}
+        headers = {"X-Amzn-SageMaker-OutputLocation": location}
+        return 200, headers, json.dumps(resp)
diff --git a/moto/sagemakerruntime/urls.py b/moto/sagemakerruntime/urls.py
index 01c753437..f8b6976ae 100644
--- a/moto/sagemakerruntime/urls.py
+++ b/moto/sagemakerruntime/urls.py
@@ -10,5 +10,6 @@ response = SageMakerRuntimeResponse()
 
 
 url_paths = {
+    "{0}/endpoints/(?P<name>[^/]+)/async-invocations$": response.dispatch,
     "{0}/endpoints/(?P<name>[^/]+)/invocations$": response.dispatch,
 }


Test patch:
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
diff --git a/tests/test_sagemakerruntime/test_sagemakerruntime.py b/tests/test_sagemakerruntime/test_sagemakerruntime.py
index 60b256d5a..9ccece62a 100644
--- a/tests/test_sagemakerruntime/test_sagemakerruntime.py
+++ b/tests/test_sagemakerruntime/test_sagemakerruntime.py
@@ -1,7 +1,10 @@
+import json
+
 import boto3
 import requests
 
-from moto import mock_sagemakerruntime, settings
+from moto import mock_s3, mock_sagemakerruntime, settings
+from moto.s3.utils import bucket_and_name_from_url
 
 # See our Development Tips on writing tests for hints on how to write good tests:
 # http://docs.getmoto.org/en/latest/docs/contributing/development_tips/tests.html
@@ -54,3 +57,49 @@ def test_invoke_endpoint():
         EndpointName="asdf", Body="qwer", Accept="sth", TargetModel="tm"
     )
     assert body["Body"].read() == b"second body"
+
+
+@mock_s3
+@mock_sagemakerruntime
+def test_invoke_endpoint_async():
+    client = boto3.client("sagemaker-runtime", region_name="us-east-1")
+    base_url = (
+        "localhost:5000" if settings.TEST_SERVER_MODE else "motoapi.amazonaws.com"
+    )
+
+    sagemaker_result = {
+        "results": [
+            {
+                "Body": "first body",
+                "ContentType": "text/xml",
+                "InvokedProductionVariant": "prod",
+                "CustomAttributes": "my_attr",
+            },
+            {"Body": "second body"},
+        ]
+    }
+    requests.post(
+        f"http://{base_url}/moto-api/static/sagemaker/endpoint-results",
+        json=sagemaker_result,
+    )
+
+    # Return the first item from the list
+    body = client.invoke_endpoint_async(EndpointName="asdf", InputLocation="qwer")
+    first_output_location = body["OutputLocation"]
+
+    # Same input -> same output
+    body = client.invoke_endpoint_async(EndpointName="asdf", InputLocation="qwer")
+    assert body["OutputLocation"] == first_output_location
+
+    # Different input -> second item
+    body = client.invoke_endpoint_async(
+        EndpointName="asdf", InputLocation="asf", InferenceId="sth"
+    )
+    second_output_location = body["OutputLocation"]
+    assert body["InferenceId"] == "sth"
+
+    s3 = boto3.client("s3", "us-east-1")
+    bucket_name, obj = bucket_and_name_from_url(second_output_location)
+    resp = s3.get_object(Bucket=bucket_name, Key=obj)
+    resp = json.loads(resp["Body"].read().decode("utf-8"))
+    assert resp["Body"] == "second body"