Can't effectively test code that uses `aws-encryption-sdk` with stubs

[philsnow]

I'm trying to write some code that consumes aws-encryption-sdk, specifically KMSMasterKeyProvider. Everything works fine but I would like to add tests with stubbed out KMS requests+responses (using botocore.stub.Stubber'ed boto3 clients).

My test setup creates a botocore session with fake credentials, e.g.

@pytest.fixture
def botocore_session():
    session = botocore.session.get_session()
    session.set_credentials("invalid-access-key", "invalid-secret-key")
    session.set_stream_logger('', 'DEBUG')
    return session

and then passes that session down into the KMSMasterKeyProvider constructor. However, I can't also pass in a boto3 KMS client object for it to use. It creates its own clients here

aws-encryption-sdk-python/src/aws_encryption_sdk/key_providers/kms.py

Lines 163 to 167 in 0f4dc6e

	if region_name not in self._regional_clients:
	session = boto3.session.Session(botocore_session=self.config.botocore_session)
	client = session.client("kms", region_name=region_name, config=self._user_agent_adding_config)
	self._register_client(client, region_name)
	self._regional_clients[region_name] = client

I could create a KMSMasterKeyProvider object in my test setup, reach into its ._regional_clients and add my stubbed KMS client, and inject that into my code. Is that the best approach for now? Is there some other approach to stubbing that I should try?

[mattsb42-aws]

As of v1.3.14 that was released last week, moto[1] supports all of the necessary KMS APIs. Have you tried using that?

ex:

import aws_encryption_sdk
import boto3
import pytest
from aws_encryption_sdk.key_providers.kms import KMSMasterKey
from moto import mock_kms


@pytest.fixture
def fake_cmk():
    with mock_kms():
        kms = boto3.client("kms", region_name="us-west-2")
        response = kms.create_key()
        yield response["KeyMetadata"]["Arn"]


def test_my_stuff(fake_cmk):
    kms_master_key = KMSMasterKey(key_id=fake_cmk)
    ciphertext, header = aws_encryption_sdk.encrypt(
        source=b"my plaintext data!",
        key_provider=kms_master_key,
    )

[1] http://docs.getmoto.org/

[philsnow]

Dealing with testing has always been the part of working with boto3 that I dread.

I've looked at using moto off and on for a long time, and I always get put off by the fact that it doesn't have full coverage of services/actions. I'll look at it for this particular case @mattsb42-aws (and thanks for the recommendation, I really hadn't considered it, but it might work for this case), but the issue/question still stands for anybody who is trying to test aws-encryption-sdk who needs to also stub one of the services/actions that moto doesn't yet support.

[philsnow]

(It turned out moto wouldn't work for me without some changes I didn't want to have to make, because it doesn't support KMS CreateAlias yet)

In case anybody else is trying to do this with botocore.stub.Stubber, one more tip which might save you some aggravation:

aws-encryption-sdk-python/src/aws_encryption_sdk/key_providers/kms.py

Line 128 in 0f4dc6e

self.add_master_keys_from_list(self.config.key_ids)

calls

aws-encryption-sdk-python/src/aws_encryption_sdk/key_providers/kms.py

Line 195 in 0f4dc6e

return KMSMasterKey(config=KMSMasterKeyConfig(key_id=key_id, client=self._client(_key_id)))

which calls

aws-encryption-sdk-python/src/aws_encryption_sdk/key_providers/kms.py

Lines 177 to 184 in 0f4dc6e

	def _client(self, key_id):
	"""Returns a Boto3 KMS client for the appropriate region.
	:param str key_id: KMS CMK ID
	"""
	region_name = _region_from_key_id(key_id, self.default_region)
	self.add_regional_client(region_name)
	return self._regional_clients[region_name]

, which instantiates boto3 clients for regions corresponding to the cmk_ids that you pass to the constructor of KMSMasterKeyProvider.

If you want it to use your stubbed boto3 KMS client, you need to create the KMSMasterKeyProvider with an empty list of cmk_ids (the default), add your clients to the internal dict, and then add your CMK:

master_key_provider = aws_encryption_sdk.KMSMasterKeyProvider()
if kms_client_mock is not None:
  master_key_provider._register_client(kms_client_mock, 'us-east-1')
  master_key_provider._regional_clients['us-east-1'] = kms_client_mock
# Have to add the master key after instantiating,
# otherwise the __init__ creates a boto3 client itself and
# we don't get to use our mocked one
master_key_provider.add_master_key(cmk_id)

[mattsb42-aws]

moto does support CreateAlias

import aws_encryption_sdk
import boto3
import pytest
from aws_encryption_sdk.key_providers.kms import KMSMasterKey, KMSMasterKeyProvider
from moto import mock_kms

MY_ALIAS_NAME = "alias/awesome"
REGION = "us-west-2"


@pytest.fixture
def fake_cmk():
    with mock_kms():
        kms = boto3.client("kms", region_name=REGION)
        response = kms.create_key()
        cmk_arn = response["KeyMetadata"]["Arn"]
        kms.create_alias(
            AliasName=MY_ALIAS_NAME, TargetKeyId=cmk_arn,
        )
        yield


def test_with_master_key_provider(fake_cmk):
    kms_master_key = KMSMasterKeyProvider(key_ids=[MY_ALIAS_NAME], region_names=[REGION])
    ciphertext, header = aws_encryption_sdk.encrypt(
        source=b"my plaintext data!", key_provider=kms_master_key,
    )


def test_with_master_key(fake_cmk):
    kms_master_key = KMSMasterKey(
        key_id=MY_ALIAS_NAME, client=boto3.client("kms", region_name=REGION)
    )
    ciphertext, header = aws_encryption_sdk.encrypt(
        source=b"my plaintext data!", key_provider=kms_master_key,
    )

[mattsb42-aws]

I'm personally not a fan of the stubber for reasons I explain here[1], but if you prefer to use them, be aware that the _register_client method and _regional_clients attribute are not part of our supported public API. They are probably unlikely to change because we are not going to do any further feature development on master key providers, but be aware that we do not guarantee that those APIs will continue to exist and do the same things.

Once we finish adding keyring support[2], the AWS KMS keyring[3] will expose an interface in the form of client suppliers that, while not designed to work with stubs, will give you a supported API that you can use to inject stubs.

[1] boto/boto3#2123 (comment)
[2] #146
[3] https://github.com/awslabs/aws-encryption-sdk-specification/blob/master/framework/kms-keyring.md