3 files changed, 22 insertions, 5 deletions
diff --git a/python/tink/integration/awskms/_aws_kms_client.py b/python/tink/integration/awskms/_aws_kms_client.py
index be368b078..d17d70b7c 100644
--- a/python/tink/integration/awskms/_aws_kms_client.py
+++ b/python/tink/integration/awskms/_aws_kms_client.py
@@ -63,11 +63,6 @@ class _AwsKmsAead(aead.Aead):
           CiphertextBlob=ciphertext,
           EncryptionContext=_encryption_context(associated_data),
       )
-      if response['KeyId'] != self.key_arn:
-        raise tink.TinkError(
-            'invalid key id: got %s, want %s'
-            % (self.key_arn, response['KeyId'])
-        )
       return response['Plaintext']
     except exceptions.ClientError as e:
       raise tink.TinkError(e)
diff --git a/python/tink/integration/awskms/_aws_kms_client_test.py b/python/tink/integration/awskms/_aws_kms_client_test.py
index 129f2de9e..1968c7bac 100644
--- a/python/tink/integration/awskms/_aws_kms_client_test.py
+++ b/python/tink/integration/awskms/_aws_kms_client_test.py
@@ -27,10 +27,17 @@ from tink.testing import helper
 
 CREDENTIAL_PATH = os.path.join(helper.tink_py_testdata_path(),
                                'aws/credentials.ini')
+
 KEY_URI = ('aws-kms://arn:aws:kms:us-east-2:235739564943:key/'
            '3ee50705-5a82-4f5b-9753-05c4f473922f')
+
+# An alias for KEY_URI.
+KEY_ALIAS_URI = ('aws-kms://arn:aws:kms:us-east-2:235739564943:alias/'
+                 'unit-and-integration-testing')
+
 KEY_URI_2 = ('aws-kms://arn:aws:kms:us-east-2:235739564943:key/'
              'b3ca2efd-a8fb-47f2-b541-7e20f8c5cd11')
+
 GCP_KEY_URI = ('gcp-kms://projects/tink-test-infrastructure/locations/global/'
                'keyRings/unit-and-integration-testing/cryptoKeys/aead-key')
 
@@ -41,6 +48,7 @@ class AwsKmsClientTest(absltest.TestCase):
     aws_client = awskms.AwsKmsClient(KEY_URI, CREDENTIAL_PATH)
 
     self.assertEqual(aws_client.does_support(KEY_URI), True)
+    self.assertEqual(aws_client.does_support(KEY_ALIAS_URI), False)
     self.assertEqual(aws_client.does_support(KEY_URI_2), False)
     self.assertEqual(aws_client.does_support(GCP_KEY_URI), False)
 
@@ -48,6 +56,7 @@ class AwsKmsClientTest(absltest.TestCase):
     aws_client = awskms.AwsKmsClient('', CREDENTIAL_PATH)
 
     self.assertEqual(aws_client.does_support(KEY_URI), True)
+    self.assertEqual(aws_client.does_support(KEY_ALIAS_URI), True)
     self.assertEqual(aws_client.does_support(KEY_URI_2), True)
     self.assertEqual(aws_client.does_support(GCP_KEY_URI), False)
 
diff --git a/python/tink/integration/awskms/_aws_kms_integration_test.py b/python/tink/integration/awskms/_aws_kms_integration_test.py
index b670a4e7e..eeb1e210c 100644
--- a/python/tink/integration/awskms/_aws_kms_integration_test.py
+++ b/python/tink/integration/awskms/_aws_kms_integration_test.py
@@ -69,6 +69,19 @@ class AwsKmsAeadTest(absltest.TestCase):
     ciphertext = aws_aead.encrypt(plaintext, b'')
     self.assertEqual(plaintext, aws_aead.decrypt(ciphertext, b''))
 
+  def test_encrypt_decrypt_with_key_alias(self):
+    aws_client = awskms.AwsKmsClient(KEY_ALIAS_URI, CREDENTIAL_PATH)
+    aws_aead = aws_client.get_aead(KEY_ALIAS_URI)
+
+    plaintext = b'hello'
+    associated_data = b'world'
+    ciphertext = aws_aead.encrypt(plaintext, associated_data)
+    self.assertEqual(plaintext, aws_aead.decrypt(ciphertext, associated_data))
+
+    plaintext = b'hello'
+    ciphertext = aws_aead.encrypt(plaintext, b'')
+    self.assertEqual(plaintext, aws_aead.decrypt(ciphertext, b''))
+
   def test_corrupted_ciphertext(self):
     aws_client = awskms.AwsKmsClient(KEY_URI, CREDENTIAL_PATH)
     aws_aead = aws_client.get_aead(KEY_URI)