Skip to content

Migrate unit/test_streaming_client_stream_encryptor.py from unittest to pytest #127

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Dec 17, 2018
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 19 additions & 17 deletions test/unit/test_streaming_client_stream_encryptor.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
# language governing permissions and limitations under the License.
"""Unit test suite for aws_encryption_sdk.streaming_client.StreamEncryptor"""
import io
import unittest

import pytest
import six
Expand All @@ -36,8 +35,9 @@
pytestmark = [pytest.mark.unit, pytest.mark.local]


class TestStreamEncryptor(unittest.TestCase):
def setUp(self):
class TestStreamEncryptor(object):
@pytest.fixture(autouse=True)
def apply_fixtures(self):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We also need to run the old tearDown actions after the tests are done. This is what is likely causing most of the test failures.

This should work by just running them in this method after a yield.

@pytest.fixture(autouse=True)
def apply_fixtures(self):
    do_stuff_from_old_setup()
    yield
    do_stuff_from_old_teardown()

https://docs.pytest.org/en/latest/fixture.html#autouse-fixtures-xunit-setup-on-steroids

# Create mock key provider
self.mock_key_provider = MagicMock()
self.mock_key_provider.__class__ = MasterKeyProvider
Expand Down Expand Up @@ -144,8 +144,8 @@ def setUp(self):
# Set up serialize_frame patch
self.mock_serialize_frame_patcher = patch("aws_encryption_sdk.streaming_client.serialize_frame")
self.mock_serialize_frame = self.mock_serialize_frame_patcher.start()

def tearDown(self):
yield
# Run tearDown
self.mock_content_type_patcher.stop()
self.mock_validate_frame_length_patcher.stop()
self.mock_message_id_patcher.stop()
Expand Down Expand Up @@ -173,14 +173,15 @@ def test_init(self):
assert test_encryptor.content_type is sentinel.content_type

def test_init_non_framed_message_too_large(self):
with six.assertRaisesRegex(self, SerializationError, "Source too large for non-framed message"):
with pytest.raises(SerializationError) as excinfo:
StreamEncryptor(
source=io.BytesIO(self.plaintext),
key_provider=self.mock_key_provider,
frame_length=0,
algorithm=self.mock_algorithm,
source_length=aws_encryption_sdk.internal.defaults.MAX_NON_FRAMED_SIZE + 1,
)
excinfo.match("Source too large for non-framed message")

def test_prep_message_no_master_keys(self):
self.mock_key_provider.master_keys_for_encryption.return_value = sentinel.primary_master_key, set()
Expand All @@ -190,9 +191,9 @@ def test_prep_message_no_master_keys(self):
frame_length=self.mock_frame_length,
source_length=5,
)

with six.assertRaisesRegex(self, MasterKeyProviderError, "No Master Keys available from Master Key Provider"):
with pytest.raises(MasterKeyProviderError) as excinfo:
test_encryptor._prep_message()
excinfo.match("No Master Keys available from Master Key Provider")

def test_prep_message_primary_master_key_not_in_master_keys(self):
self.mock_key_provider.master_keys_for_encryption.return_value = (
Expand All @@ -206,8 +207,9 @@ def test_prep_message_primary_master_key_not_in_master_keys(self):
source_length=5,
)

with six.assertRaisesRegex(self, MasterKeyProviderError, "Primary Master Key not in provided Master Keys"):
with pytest.raises(MasterKeyProviderError) as excinfo:
test_encryptor._prep_message()
excinfo.match("Primary Master Key not in provided Master Keys")

def test_prep_message_algorithm_change(self):
self.mock_encryption_materials.algorithm = Algorithm.AES_256_GCM_IV12_TAG16
Expand All @@ -217,13 +219,11 @@ def test_prep_message_algorithm_change(self):
algorithm=Algorithm.AES_128_GCM_IV12_TAG16,
source_length=128,
)

with six.assertRaisesRegex(
self,
ActionNotAllowedError,
"Cryptographic materials manager provided algorithm suite differs from algorithm suite in request.*",
):
with pytest.raises(ActionNotAllowedError) as excinfo:
test_encryptor._prep_message()
excinfo.match(
"Cryptographic materials manager provided algorithm suite differs from algorithm suite in request.*"
)

@patch("aws_encryption_sdk.streaming_client.EncryptionMaterialsRequest")
@patch("aws_encryption_sdk.streaming_client.derive_data_encryption_key")
Expand Down Expand Up @@ -397,8 +397,9 @@ def test_read_bytes_to_non_framed_body_too_large(self):
test_encryptor.bytes_read = aws_encryption_sdk.internal.defaults.MAX_NON_FRAMED_SIZE
test_encryptor._StreamEncryptor__unframed_plaintext_cache = pt_stream

with six.assertRaisesRegex(self, SerializationError, "Source too large for non-framed message"):
with pytest.raises(SerializationError) as excinfo:
test_encryptor._read_bytes_to_non_framed_body(5)
excinfo.match("Source too large for non-framed message")

def test_read_bytes_to_non_framed_body_close(self):
test_encryptor = StreamEncryptor(source=io.BytesIO(self.plaintext), key_provider=self.mock_key_provider)
Expand Down Expand Up @@ -487,8 +488,9 @@ def test_read_bytes_unsupported_type(self, mock_read_non_framed, mock_read_frame
test_encryptor._encryption_materials = self.mock_encryption_materials
test_encryptor._header = MagicMock()
test_encryptor.content_type = None
with six.assertRaisesRegex(self, NotSupportedError, "Unsupported content type"):
with pytest.raises(NotSupportedError) as excinfo:
test_encryptor._read_bytes(5)
excinfo.match("Unsupported content type")
assert not mock_read_non_framed.called
assert not mock_read_framed.called

Expand Down