Repository: getmoto/moto (python)

Issue:
KMS sign fails with alias ARN
During execution of

```python
@mock_aws
def main():
    kms: KMSClient = boto3.client("kms", region_name="eu-west-1")
    key = kms.create_key(
        KeyUsage="SIGN_VERIFY",
        KeySpec="ECC_NIST_P256",
    )
    key_id = key["KeyMetadata"]["KeyId"]
    kms.create_alias(
        AliasName=f"alias/foo",
        TargetKeyId=key_id,
    )
    kms.sign(
        KeyId=f"arn:aws:kms:eu-west-1:123456789012:alias/foo",
        MessageType="RAW",
        SigningAlgorithm="ECDSA_SHA_256",
        Message=bytes(8),
    )
```

an exception is raised:
```
Traceback (most recent call last):
  File "/home/phil/repos/cbcf4-cdk-library/src/moto_bug.py", line 27, in <module>
    main()
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/core/models.py", line 122, in wrapper
    result = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/src/moto_bug.py", line 20, in main
    kms.sign(
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/client.py", line 553, in _api_call
    return self._make_api_call(operation_name, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/client.py", line 989, in _make_api_call
    http, parsed_response = self._make_request(
                            ^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/client.py", line 1015, in _make_request
    return self._endpoint.make_request(operation_model, request_dict)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/endpoint.py", line 119, in make_request
    return self._send_request(request_dict, operation_model)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/endpoint.py", line 202, in _send_request
    while self._needs_retry(
          ^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/endpoint.py", line 354, in _needs_retry
    responses = self._event_emitter.emit(
                ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/hooks.py", line 412, in emit
    return self._emitter.emit(aliased_event_name, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/hooks.py", line 256, in emit
    return self._emit(event_name, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/hooks.py", line 239, in _emit
    response = handler(**kwargs)
               ^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/retryhandler.py", line 207, in __call__
    if self._checker(**checker_kwargs):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/retryhandler.py", line 284, in __call__
    should_retry = self._should_retry(
                   ^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/retryhandler.py", line 307, in _should_retry
    return self._checker(
           ^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/retryhandler.py", line 363, in __call__
    checker_response = checker(
                       ^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/retryhandler.py", line 247, in __call__
    return self._check_caught_exception(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/retryhandler.py", line 416, in _check_caught_exception
    raise caught_exception
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/endpoint.py", line 278, in _do_get_response
    responses = self._event_emitter.emit(event_name, request=request)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/hooks.py", line 412, in emit
    return self._emitter.emit(aliased_event_name, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/hooks.py", line 256, in emit
    return self._emit(event_name, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/botocore/hooks.py", line 239, in _emit
    response = handler(**kwargs)
               ^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/core/botocore_stubber.py", line 38, in __call__
    response = self.process_request(request)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/core/botocore_stubber.py", line 88, in process_request
    status, headers, body = method_to_execute(
                            ^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/core/responses.py", line 290, in dispatch
    return cls()._dispatch(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/core/responses.py", line 501, in _dispatch
    return self.call_action()
           ^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/core/responses.py", line 587, in call_action
    response = method()
               ^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/kms/responses.py", line 645, in sign
    key_id, signature, signing_algorithm = self.kms_backend.sign(
                                           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/kms/models.py", line 681, in sign
    key = self.describe_key(key_id)
          ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/kms/models.py", line 381, in describe_key
    key = self.keys[self.get_key_id(key_id)]
                    ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/phil/repos/cbcf4-cdk-library/venv/lib64/python3.12/site-packages/moto/kms/models.py", line 398, in get_key_id
    if key_id.startswith("arn:") and ":key/" in key_id:
       ^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'startswith'
```

It seems like the alias isn't resolved correctly when passed along as an ARN. Tried with Moto 5.0.13 and 5.0.14dev62 and Python 3.12.4.

Pull request:
KMS: sign() now supports Alias ARNs
Fixes #8090 

Patch:
diff --git a/moto/kms/models.py b/moto/kms/models.py
index 9f35cd893..5e9f3f6f0 100644
--- a/moto/kms/models.py
+++ b/moto/kms/models.py
@@ -372,22 +372,13 @@ class KmsBackend(BaseBackend):
             self.keys.pop(key_id)
 
     def describe_key(self, key_id: str) -> Key:
-        # allow the different methods (alias, ARN :key/, keyId, ARN alias) to
-        # describe key not just KeyId
-        key_id = self.get_key_id(key_id)
-        if r"alias/" in str(key_id).lower():
-            key_id = self.get_key_id_from_alias(key_id)  # type: ignore[assignment]
-
-        key = self.keys[self.get_key_id(key_id)]
+        key = self.keys[self.any_id_to_key_id(key_id)]
 
         if key.multi_region:
-            if key.arn == key.multi_region_configuration["PrimaryKey"]["Arn"]:
-                return self.keys[self.get_key_id(key_id)]
-            else:
+            if key.arn != key.multi_region_configuration["PrimaryKey"]["Arn"]:
                 key.multi_region_configuration["MultiRegionKeyType"] = "REPLICA"
-                return key
 
-        return self.keys[self.get_key_id(key_id)]
+        return key
 
     def list_keys(self) -> Iterable[Key]:
         return self.keys.values()
diff --git a/moto/moto_api/_internal/models.py b/moto/moto_api/_internal/models.py
index ba5241aa4..2ab8861ce 100644
--- a/moto/moto_api/_internal/models.py
+++ b/moto/moto_api/_internal/models.py
@@ -84,6 +84,18 @@ class MotoAPIBackend(BaseBackend):
         backend = sagemakerruntime_backends[account_id][region]
         backend.results_queue.append((body, content_type, prod_variant, custom_attrs))
 
+    def set_sagemaker_async_result(
+        self,
+        is_failure: bool,
+        data: str,
+        account_id: str,
+        region: str,
+    ) -> None:
+        from moto.sagemakerruntime.models import sagemakerruntime_backends
+
+        backend = sagemakerruntime_backends[account_id][region]
+        backend.async_results_queue.append((is_failure, data))
+
     def set_rds_data_result(
         self,
         records: Optional[List[List[Dict[str, Any]]]],
diff --git a/moto/moto_api/_internal/responses.py b/moto/moto_api/_internal/responses.py
index fdd873577..dd5712475 100644
--- a/moto/moto_api/_internal/responses.py
+++ b/moto/moto_api/_internal/responses.py
@@ -215,6 +215,29 @@ class MotoAPIResponse(BaseResponse):
             )
         return 201, {}, ""
 
+    def set_sagemaker_async_result(
+        self,
+        request: Any,
+        full_url: str,  # pylint: disable=unused-argument
+        headers: Any,
+    ) -> TYPE_RESPONSE:
+        from .models import moto_api_backend
+
+        body = self._get_body(headers, request)
+        account_id = body.get("account_id", DEFAULT_ACCOUNT_ID)
+        region = body.get("region", "us-east-1")
+
+        for result in body.get("results", []):
+            data = result["data"]
+            is_failure = result.get("is_failure", False)
+            moto_api_backend.set_sagemaker_async_result(
+                data=data,
+                is_failure=is_failure,
+                account_id=account_id,
+                region=region,
+            )
+        return 201, {}, ""
+
     def set_sagemaker_result(
         self,
         request: Any,
diff --git a/moto/moto_api/_internal/urls.py b/moto/moto_api/_internal/urls.py
index 7821f11b5..9f3e66f18 100644
--- a/moto/moto_api/_internal/urls.py
+++ b/moto/moto_api/_internal/urls.py
@@ -20,6 +20,7 @@ url_paths = {
     "{0}/moto-api/static/lambda-simple/response": response_instance.set_lambda_simple_result,
     "{0}/moto-api/static/resilience-hub-assessments/response": response_instance.set_resilience_result,
     "{0}/moto-api/static/sagemaker/endpoint-results": response_instance.set_sagemaker_result,
+    "{0}/moto-api/static/sagemaker/async-endpoint-results": response_instance.set_sagemaker_async_result,
     "{0}/moto-api/static/timestream/query-results": response_instance.set_timestream_result,
     "{0}/moto-api/static/rds-data/statement-results": response_instance.set_rds_data_result,
     "{0}/moto-api/state-manager/get-transition": response_instance.get_transition,
diff --git a/moto/sagemakerruntime/models.py b/moto/sagemakerruntime/models.py
index 45c05c13e..d17863193 100644
--- a/moto/sagemakerruntime/models.py
+++ b/moto/sagemakerruntime/models.py
@@ -10,9 +10,10 @@ 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.async_results: Dict[str, Dict[str, Tuple[str, str]]] = {}
         self.results: Dict[str, Dict[bytes, Tuple[str, str, str, str]]] = {}
         self.results_queue: List[Tuple[str, str, str, str]] = []
+        self.async_results_queue: List[Tuple[bool, str]] = []
 
     def invoke_endpoint(
         self, endpoint_name: str, unique_repr: bytes
@@ -65,35 +66,81 @@ class SageMakerRuntimeBackend(BaseBackend):
             )
         return self.results[endpoint_name][unique_repr]
 
-    def invoke_endpoint_async(self, endpoint_name: str, input_location: str) -> str:
+    def invoke_endpoint_async(
+        self, endpoint_name: str, input_location: str
+    ) -> Tuple[str, str]:
+        """
+        This call will return static data by default.
+
+        You can use a dedicated API to override this, by configuring a queue of expected results.
+        Subsequent requests using the same `InputLocation` will return the same result.
+        Other requests using a different `InputLocation` will take the next result from the queue.
+        Please note that for backward compatibility, if the async queue is empty, the sync queue will be used, if configured.
+
+        Configuring this queue by making an HTTP request to `/moto-api/static/sagemaker/async-endpoint-results`.
+        An example invocation looks like this:
+
+        .. sourcecode:: python
+
+            expected_results = {
+                "account_id": "123456789012",  # This is the default - can be omitted
+                "region": "us-east-1",  # This is the default - can be omitted
+                "results": [
+                    {
+                        "data": json.dumps({"first": "output"}),
+                    },
+                    {
+                        "is_failure": True,
+                        "data": "second inference failed",
+                    },
+                    # other results as required
+                ],
+            }
+            requests.post(
+                "http://motoapi.amazonaws.com/moto-api/static/sagemaker/async-endpoint-results",
+                json=expected_results,
+            )
+
+            client = boto3.client("sagemaker-runtime", region_name="us-east-1")
+            details = client.invoke_endpoint_async(EndpointName="asdf", InputLocation="qwer")
+
+        """
         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:
+        if self.async_results_queue:
+            is_failure, data = self.async_results_queue.pop(0)
+        elif self.results_queue:  # Backward compatibility
+            is_failure = False
             body, _type, variant, attrs = self.results_queue.pop(0)
+            data = json.dumps(
+                {
+                    "Body": body,
+                    "ContentType": _type,
+                    "InvokedProductionVariant": variant,
+                    "CustomAttributes": attrs,
+                }
+            )
         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")
+            is_failure = False
+            data = json.dumps({"default": "response"})
 
         output_bucket = f"sagemaker-output-{random.uuid4()}"
         output_location = "response.json"
+        failure_location = "failure.json"
+        result_location = failure_location if is_failure else output_location
         from moto.s3.models import s3_backends
 
         s3_backend = s3_backends[self.account_id][self.partition]
         s3_backend.create_bucket(output_bucket, region_name=self.region_name)
-        s3_backend.put_object(output_bucket, output_location, value=output)
+        s3_backend.put_object(
+            output_bucket, result_location, value=data.encode("utf-8")
+        )
+
         self.async_results[endpoint_name][input_location] = (
-            f"s3://{output_bucket}/{output_location}"
+            f"s3://{output_bucket}/{output_location}",
+            f"s3://{output_bucket}/{failure_location}",
         )
         return self.async_results[endpoint_name][input_location]
 
diff --git a/moto/sagemakerruntime/responses.py b/moto/sagemakerruntime/responses.py
index c30c4b2b6..f409325c5 100644
--- a/moto/sagemakerruntime/responses.py
+++ b/moto/sagemakerruntime/responses.py
@@ -49,9 +49,14 @@ class SageMakerRuntimeResponse(BaseResponse):
         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
+        output_location, failure_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}
+        headers = {
+            "X-Amzn-SageMaker-OutputLocation": output_location,
+            "X-Amzn-SageMaker-FailureLocation": failure_location,
+        }
         return 200, headers, json.dumps(resp)


Test patch:
diff --git a/tests/test_kms/test_kms_boto3.py b/tests/test_kms/test_kms_boto3.py
index 455ef368b..d1ccaf402 100644
--- a/tests/test_kms/test_kms_boto3.py
+++ b/tests/test_kms/test_kms_boto3.py
@@ -1192,6 +1192,30 @@ def test_sign_happy(plaintext):
     assert sign_response["KeyId"] == key_arn
 
 
+@mock_aws
+def test_sign_using_alias():
+    client = boto3.client("kms", region_name="us-west-2")
+
+    key = client.create_key(
+        Description="sk", KeyUsage="SIGN_VERIFY", KeySpec="RSA_2048"
+    )
+    key_id = key["KeyMetadata"]["KeyId"]
+    client.create_alias(AliasName="alias/my-alias", TargetKeyId=key_id)
+    client.sign(
+        KeyId=key_id, Message="plaintext", SigningAlgorithm="RSASSA_PSS_SHA_256"
+    )
+
+    key_id = "alias/my-alias"
+    client.sign(
+        KeyId=key_id, Message="plaintext", SigningAlgorithm="RSASSA_PSS_SHA_256"
+    )
+
+    key_id = f"arn:aws:kms:us-west-2:{ACCOUNT_ID}:alias/my-alias"
+    client.sign(
+        KeyId=key_id, Message="plaintext", SigningAlgorithm="RSASSA_PSS_SHA_256"
+    )
+
+
 @mock_aws
 def test_sign_invalid_signing_algorithm():
     client = boto3.client("kms", region_name="us-west-2")
diff --git a/tests/test_sagemakerruntime/test_sagemakerruntime.py b/tests/test_sagemakerruntime/test_sagemakerruntime.py
index 0d6633a8b..d30ad1dae 100644
--- a/tests/test_sagemakerruntime/test_sagemakerruntime.py
+++ b/tests/test_sagemakerruntime/test_sagemakerruntime.py
@@ -68,37 +68,76 @@ def test_invoke_endpoint_async():
 
     sagemaker_result = {
         "results": [
+            {"data": json.dumps({"first": "output"})},
             {
-                "Body": "first body",
-                "ContentType": "text/xml",
-                "InvokedProductionVariant": "prod",
-                "CustomAttributes": "my_attr",
+                "is_failure": True,
+                "data": "second failure",
             },
-            {"Body": "second body"},
         ]
     }
     requests.post(
-        f"http://{base_url}/moto-api/static/sagemaker/endpoint-results",
+        f"http://{base_url}/moto-api/static/sagemaker/async-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"]
+    first_failure_location = body["FailureLocation"]
 
     # Same input -> same output
     body = client.invoke_endpoint_async(EndpointName="asdf", InputLocation="qwer")
     assert body["OutputLocation"] == first_output_location
+    assert body["FailureLocation"] == first_failure_location
+
+    s3 = boto3.client("s3", "us-east-1")
+    bucket_name, obj = bucket_and_name_from_url(first_output_location)
+    resp = s3.get_object(Bucket=bucket_name, Key=obj)
+    resp = json.loads(resp["Body"].read().decode("utf-8"))
+    assert resp == {"first": "output"}
 
     # Different input -> second item
     body = client.invoke_endpoint_async(
         EndpointName="asdf", InputLocation="asf", InferenceId="sth"
     )
-    second_output_location = body["OutputLocation"]
+    second_failure_location = body["FailureLocation"]
     assert body["InferenceId"] == "sth"
 
+    bucket_name, obj = bucket_and_name_from_url(second_failure_location)
+    resp = s3.get_object(Bucket=bucket_name, Key=obj)
+    resp = resp["Body"].read().decode("utf-8")
+    assert resp == "second failure"
+
+
+@mock_aws
+def test_invoke_endpoint_async_should_read_sync_queue_if_async_not_configured():
+    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": "support sync queue for backward compatibility"},
+        ]
+    }
+    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"]
+    first_failure_location = body["FailureLocation"]
+
+    # Same input -> same output
+    body = client.invoke_endpoint_async(EndpointName="asdf", InputLocation="qwer")
+    assert body["OutputLocation"] == first_output_location
+    assert body["FailureLocation"] == first_failure_location
+
     s3 = boto3.client("s3", "us-east-1")
-    bucket_name, obj = bucket_and_name_from_url(second_output_location)
+    bucket_name, obj = bucket_and_name_from_url(first_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"
+    assert resp["Body"] == "support sync queue for backward compatibility"