fix encoding for Python 2 in KMS tests

This commit is contained in:
mattsb42-aws 2019-09-13 14:08:26 -07:00
parent 4f34af95bc
commit 4c7cdec965

View File

@ -1,15 +1,16 @@
# -*- coding: utf-8 -*-
from __future__ import unicode_literals
from datetime import date
from datetime import datetime
from dateutil.tz import tzutc
import base64
import binascii
import os
import re
import boto3
import boto.kms
import botocore.exceptions
import six
import sure # noqa
from boto.exception import JSONResponseError
from boto.kms.exceptions import AlreadyExistsException, NotFoundException
@ -27,6 +28,13 @@ PLAINTEXT_VECTORS = (
)
def _get_encoded_value(plaintext):
if isinstance(plaintext, six.binary_type):
return plaintext
return plaintext.encode("utf-8")
@mock_kms
def test_create_key():
conn = boto3.client("kms", region_name="us-east-1")
@ -245,11 +253,6 @@ def test_decrypt(plaintext):
encrypt_response = client.encrypt(KeyId=key_id, Plaintext=plaintext)
try:
encoded_plaintext = plaintext.encode("utf-8")
except AttributeError:
encoded_plaintext = plaintext
client.create_key(Description="key")
# CiphertextBlob must NOT be base64-encoded
with assert_raises(Exception):
@ -261,7 +264,7 @@ def test_decrypt(plaintext):
with assert_raises(Exception):
base64.b64decode(decrypt_response["Plaintext"], validate=True)
decrypt_response["Plaintext"].should.equal(encoded_plaintext)
decrypt_response["Plaintext"].should.equal(_get_encoded_value(plaintext))
decrypt_response["KeyId"].should.equal(key_arn)
@ -689,13 +692,8 @@ def test_kms_encrypt_boto3(plaintext):
key = client.create_key(Description="key")
response = client.encrypt(KeyId=key["KeyMetadata"]["KeyId"], Plaintext=plaintext)
try:
encoded_plaintext = plaintext.encode("utf-8")
except AttributeError:
encoded_plaintext = plaintext
response = client.decrypt(CiphertextBlob=response["CiphertextBlob"])
response["Plaintext"].should.equal(encoded_plaintext)
response["Plaintext"].should.equal(_get_encoded_value(plaintext))
@mock_kms
@ -918,11 +916,6 @@ def test_re_encrypt_decrypt(plaintext):
EncryptionContext={"encryption": "context"},
)
try:
encoded_plaintext = plaintext.encode("utf-8")
except AttributeError:
encoded_plaintext = plaintext
re_encrypt_response = client.re_encrypt(
CiphertextBlob=encrypt_response["CiphertextBlob"],
SourceEncryptionContext={"encryption": "context"},
@ -941,14 +934,14 @@ def test_re_encrypt_decrypt(plaintext):
CiphertextBlob=encrypt_response["CiphertextBlob"],
EncryptionContext={"encryption": "context"},
)
decrypt_response_1["Plaintext"].should.equal(encoded_plaintext)
decrypt_response_1["Plaintext"].should.equal(_get_encoded_value(plaintext))
decrypt_response_1["KeyId"].should.equal(key_1_arn)
decrypt_response_2 = client.decrypt(
CiphertextBlob=re_encrypt_response["CiphertextBlob"],
EncryptionContext={"another": "context"},
)
decrypt_response_2["Plaintext"].should.equal(encoded_plaintext)
decrypt_response_2["Plaintext"].should.equal(_get_encoded_value(plaintext))
decrypt_response_2["KeyId"].should.equal(key_2_arn)
decrypt_response_1["Plaintext"].should.equal(decrypt_response_2["Plaintext"])