un-mock tests to work correctly with updated prep_data_stream implementation

Author: mattsb42-aws

test/unit/test_streaming_client_encryption_stream.py

@@ -20,6 +20,7 @@
 
 import aws_encryption_sdk.exceptions
 from aws_encryption_sdk.internal.defaults import LINE_LENGTH
+from aws_encryption_sdk.internal.utils.streams import InsistentReaderBytesIO
 from aws_encryption_sdk.key_providers.base import MasterKeyProvider
 from aws_encryption_sdk.streaming_client import _ClientConfig, _EncryptionStream
 
@@ -107,17 +108,19 @@ def test_new_with_params(self):
             line_length=io.DEFAULT_BUFFER_SIZE,
             source_length=mock_int_sentinel,
         )
-        assert mock_stream.config == MockClientConfig(
-            source=self.mock_source_stream,
-            key_provider=self.mock_key_provider,
-            mock_read_bytes=sentinel.read_bytes,
-            line_length=io.DEFAULT_BUFFER_SIZE,
-            source_length=mock_int_sentinel,
-        )
+
+        assert mock_stream.config.source == self.mock_source_stream
+        assert isinstance(mock_stream.config.source, InsistentReaderBytesIO)
+        assert mock_stream.config.key_provider is self.mock_key_provider
+        assert mock_stream.config.mock_read_bytes is sentinel.read_bytes
+        assert mock_stream.config.line_length == io.DEFAULT_BUFFER_SIZE
+        assert mock_stream.config.source_length is mock_int_sentinel
+
         assert mock_stream.bytes_read == 0
         assert mock_stream.output_buffer == b""
         assert not mock_stream._message_prepped
-        assert mock_stream.source_stream is self.mock_source_stream
+        assert mock_stream.source_stream == self.mock_source_stream
+        assert isinstance(mock_stream.source_stream, InsistentReaderBytesIO)
         assert mock_stream._stream_length is mock_int_sentinel
         assert mock_stream.line_length == io.DEFAULT_BUFFER_SIZE
 

test/unit/test_streaming_client_stream_decryptor.py

@@ -261,7 +261,7 @@ def test_prep_non_framed(self):
         test_decryptor._prep_non_framed()
 
         self.mock_deserialize_non_framed_values.assert_called_once_with(
-            stream=self.mock_input_stream, header=self.mock_header, verifier=sentinel.verifier
+            stream=test_decryptor.source_stream, header=self.mock_header, verifier=sentinel.verifier
         )
         assert test_decryptor.body_length == len(VALUES["data_128"])
         self.mock_get_aad_content_string.assert_called_once_with(

test/unit/test_streaming_client_stream_encryptor.py

@@ -58,8 +58,6 @@ def setUp(self):
 
         self.mock_master_key = MagicMock(__class__=MasterKey)
 
-        self.mock_input_stream = MagicMock(__class__=io.IOBase)
-
         self.mock_frame_length = MagicMock(__class__=int)
 
         self.mock_algorithm = MagicMock(__class__=Algorithm)
@@ -178,7 +176,7 @@ def tearDown(self):
 
     def test_init(self):
         test_encryptor = StreamEncryptor(
-            source=self.mock_input_stream,
+            source=io.BytesIO(self.plaintext),
             key_provider=self.mock_key_provider,
             frame_length=self.mock_frame_length,
             algorithm=self.mock_algorithm,
@@ -190,7 +188,7 @@ def test_init(self):
     def test_init_non_framed_message_too_large(self):
         with six.assertRaisesRegex(self, SerializationError, "Source too large for non-framed message"):
             StreamEncryptor(
-                source=self.mock_input_stream,
+                source=io.BytesIO(self.plaintext),
                 key_provider=self.mock_key_provider,
                 frame_length=0,
                 algorithm=self.mock_algorithm,
@@ -200,7 +198,7 @@ def test_init_non_framed_message_too_large(self):
     def test_prep_message_no_master_keys(self):
         self.mock_key_provider.master_keys_for_encryption.return_value = sentinel.primary_master_key, set()
         test_encryptor = StreamEncryptor(
-            source=self.mock_input_stream,
+            source=io.BytesIO(self.plaintext),
             key_provider=self.mock_key_provider,
             frame_length=self.mock_frame_length,
             source_length=5,
@@ -215,7 +213,7 @@ def test_prep_message_primary_master_key_not_in_master_keys(self):
             self.mock_master_keys_set,
         )
         test_encryptor = StreamEncryptor(
-            source=self.mock_input_stream,
+            source=io.BytesIO(self.plaintext),
             key_provider=self.mock_key_provider,
             frame_length=self.mock_frame_length,
             source_length=5,
@@ -227,7 +225,7 @@ def test_prep_message_primary_master_key_not_in_master_keys(self):
     def test_prep_message_algorithm_change(self):
         self.mock_encryption_materials.algorithm = Algorithm.AES_256_GCM_IV12_TAG16
         test_encryptor = StreamEncryptor(
-            source=self.mock_input_stream,
+            source=io.BytesIO(self.plaintext),
             materials_manager=self.mock_materials_manager,
             algorithm=Algorithm.AES_128_GCM_IV12_TAG16,
             source_length=128,
@@ -255,7 +253,7 @@ def test_prep_message_framed_message(
     ):
         mock_rostream.return_value = sentinel.plaintext_rostream
         test_encryptor = StreamEncryptor(
-            source=self.mock_input_stream,
+            source=io.BytesIO(self.plaintext),
             materials_manager=self.mock_materials_manager,
             frame_length=self.mock_frame_length,
             source_length=5,
@@ -359,7 +357,7 @@ def test_write_header(self):
     @patch("aws_encryption_sdk.streaming_client.non_framed_body_iv")
     def test_prep_non_framed(self, mock_non_framed_iv):
         self.mock_serialize_non_framed_open.return_value = b"1234567890"
-        test_encryptor = StreamEncryptor(source=self.mock_input_stream, key_provider=self.mock_key_provider)
+        test_encryptor = StreamEncryptor(source=io.BytesIO(self.plaintext), key_provider=self.mock_key_provider)
         test_encryptor.signer = sentinel.signer
         test_encryptor._encryption_materials = self.mock_encryption_materials
         test_encryptor._header = MagicMock()
@@ -411,7 +409,7 @@ def test_read_bytes_to_non_framed_body_too_large(self):
             test_encryptor._read_bytes_to_non_framed_body(5)
 
     def test_read_bytes_to_non_framed_body_close(self):
-        test_encryptor = StreamEncryptor(source=self.mock_input_stream, key_provider=self.mock_key_provider)
+        test_encryptor = StreamEncryptor(source=io.BytesIO(self.plaintext), key_provider=self.mock_key_provider)
         test_encryptor.signer = MagicMock()
         test_encryptor._encryption_materials = self.mock_encryption_materials
         test_encryptor.encryptor = MagicMock()
@@ -420,7 +418,9 @@ def test_read_bytes_to_non_framed_body_close(self):
         test_encryptor.encryptor.tag = sentinel.tag
         self.mock_serialize_non_framed_close.return_value = b"789"
         self.mock_serialize_footer.return_value = b"0-="
+
         test = test_encryptor._read_bytes_to_non_framed_body(len(self.plaintext) + 1)
+
         test_encryptor.signer.update.assert_has_calls(calls=(call(b"123"), call(b"456")), any_order=False)
         assert test_encryptor.source_stream.closed
         test_encryptor.encryptor.finalize.assert_called_once_with()