test_aws_encryption_sdk.py

# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Unit test suite for high-level functions in aws_encryption_sdk module"""
import unittest

from mock import MagicMock, patch, sentinel
import pytest
import six

import aws_encryption_sdk
import aws_encryption_sdk.internal.defaults

pytestmark = [pytest.mark.unit, pytest.mark.local]


class TestAwsEncryptionSdk(unittest.TestCase):

    def setUp(self):
        # Set up StreamEncryptor patch
        self.mock_stream_encryptor_patcher = patch('aws_encryption_sdk.StreamEncryptor')
        self.mock_stream_encryptor = self.mock_stream_encryptor_patcher.start()
        self.mock_stream_encryptor_instance = MagicMock()
        self.mock_stream_encryptor_instance.read.return_value = sentinel.ciphertext
        self.mock_stream_encryptor_instance.header = sentinel.header
        self.mock_stream_encryptor.return_value = self.mock_stream_encryptor_instance
        self.mock_stream_encryptor_instance.__enter__.return_value = self.mock_stream_encryptor_instance
        # Set up StreamDecryptor patch
        self.mock_stream_decryptor_patcher = patch('aws_encryption_sdk.StreamDecryptor')
        self.mock_stream_decryptor = self.mock_stream_decryptor_patcher.start()
        self.mock_stream_decryptor_instance = MagicMock()
        self.mock_stream_decryptor_instance.read.return_value = sentinel.plaintext
        self.mock_stream_decryptor_instance.header = sentinel.header
        self.mock_stream_decryptor.return_value = self.mock_stream_decryptor_instance
        self.mock_stream_decryptor_instance.__enter__.return_value = self.mock_stream_decryptor_instance

    def tearDown(self):
        self.mock_stream_encryptor_patcher.stop()
        self.mock_stream_decryptor_patcher.stop()

    def test_encrypt(self):
        test_ciphertext, test_header = aws_encryption_sdk.encrypt(
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        self.mock_stream_encryptor.called_once_with(
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        assert test_ciphertext is sentinel.ciphertext
        assert test_header is sentinel.header

    def test_decrypt(self):
        test_plaintext, test_header = aws_encryption_sdk.decrypt(
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        self.mock_stream_encryptor.called_once_with(
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        assert test_plaintext is sentinel.plaintext
        assert test_header is sentinel.header

    def test_stream_encryptor_e(self):
        test = aws_encryption_sdk.stream(
            mode='e',
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        assert test is self.mock_stream_encryptor_instance
        self.mock_stream_encryptor.assert_called_once_with(
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )

    def test_stream_encryptor_encrypt(self):
        test = aws_encryption_sdk.stream(
            mode='ENCRYPT',
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        assert test is self.mock_stream_encryptor_instance

    def test_stream_decryptor_d(self):
        test = aws_encryption_sdk.stream(
            mode='d',
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        assert test is self.mock_stream_decryptor_instance
        self.mock_stream_decryptor.assert_called_once_with(
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )

    def test_stream_decryptor_decrypt(self):
        test = aws_encryption_sdk.stream(
            mode='DECRYPT',
            a=sentinel.a,
            b=sentinel.b,
            c=sentinel.b
        )
        assert test is self.mock_stream_decryptor_instance

    def test_stream_unknown(self):
        with six.assertRaisesRegex(self, ValueError, 'Unsupported mode: *'):
            aws_encryption_sdk.stream(
                mode='ERROR',
                a=sentinel.a,
                b=sentinel.b,
                c=sentinel.b
            )