def test_get_encryption_materials(patch_for_dcmm_encrypt):
    encryption_context = {"a": "b"}
    mock_request = MagicMock(algorithm=None,
                             encryption_context=encryption_context)
    cmm = build_cmm()

    test = cmm.get_encryption_materials(request=mock_request)

    cmm.master_key_provider.master_keys_for_encryption.assert_called_once_with(
        encryption_context=encryption_context,
        plaintext_rostream=mock_request.plaintext_rostream,
        plaintext_length=mock_request.plaintext_length,
    )
    cmm._generate_signing_key_and_update_encryption_context.assert_called_once_with(
        cmm.algorithm, encryption_context)
    aws_encryption_sdk.materials_managers.default.prepare_data_keys.assert_called_once_with(
        primary_master_key=cmm.master_key_provider.master_keys_for_encryption.
        return_value[0],
        master_keys=cmm.master_key_provider.master_keys_for_encryption.
        return_value[1],
        algorithm=cmm.algorithm,
        encryption_context=encryption_context,
    )
    assert isinstance(test, EncryptionMaterials)
    assert test.algorithm is cmm.algorithm
    assert test.data_encryption_key == RawDataKey.from_data_key(
        patch_for_dcmm_encrypt[0][0])
    assert test.encrypted_data_keys == patch_for_dcmm_encrypt[0][1]
    assert test.encryption_context == encryption_context
    assert test.signing_key == patch_for_dcmm_encrypt[1]
def _data_key_to_raw_data_key(data_key):
    # type: (Union[DataKey, RawDataKey, None]) -> Union[RawDataKey, None]
    """Convert a :class:`DataKey` into a :class:`RawDataKey`."""
    if isinstance(data_key, RawDataKey) or data_key is None:
        return data_key

    return RawDataKey.from_data_key(data_key=data_key)
def test_decrypt_materials(mocker, patch_for_dcmm_decrypt):
    mock_request = MagicMock()
    cmm = build_cmm()

    test = cmm.decrypt_materials(request=mock_request)

    cmm.master_key_provider.decrypt_data_key_from_list.assert_called_once_with(
        encrypted_data_keys=mock_request.encrypted_data_keys,
        algorithm=mock_request.algorithm,
        encryption_context=mock_request.encryption_context,
    )
    cmm._load_verification_key_from_encryption_context.assert_called_once_with(
        algorithm=mock_request.algorithm,
        encryption_context=mock_request.encryption_context)
    assert test.data_key == RawDataKey.from_data_key(
        cmm.master_key_provider.decrypt_data_key_from_list.return_value)
    assert test.verification_key == patch_for_dcmm_decrypt