Skip to content

Commit 7013f40

Browse files
authored
Merge pull request #98 from mattsb42-aws/read
Properly handle negative values to _EncryptionStream.read()
2 parents 66d06e4 + 870707c commit 7013f40

File tree

3 files changed

+73
-39
lines changed

3 files changed

+73
-39
lines changed

src/aws_encryption_sdk/streaming_client.py

+4
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,10 @@ def read(self, b=None):
209209
:returns: Processed (encrypted or decrypted) bytes from source stream
210210
:rtype: bytes
211211
"""
212+
# Any negative value for b is interpreted as a full read
213+
if b is not None and b < 0:
214+
b = None
215+
212216
_LOGGER.debug("Stream read called, requesting %s bytes", b)
213217
output = io.BytesIO()
214218
if not self._message_prepped:

test/unit/test_streaming_client_encryption_stream.py

+68-39
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,11 @@
1111
# ANY KIND, either express or implied. See the License for the specific
1212
# language governing permissions and limitations under the License.
1313
"""Unit test suite for aws_encryption_sdk.streaming_client._EncryptionStream"""
14+
import copy
1415
import io
15-
import unittest
1616

1717
import attr
1818
import pytest
19-
import six
2019
from mock import MagicMock, PropertyMock, call, patch, sentinel
2120

2221
import aws_encryption_sdk.exceptions
@@ -45,20 +44,22 @@ def _read_bytes(self, b):
4544
return self.config.mock_read_bytes
4645

4746

48-
class TestEncryptionStream(unittest.TestCase):
49-
def setUp(self):
50-
self.mock_source_stream = MagicMock()
51-
self.mock_source_stream.__class__ = io.IOBase
52-
self.mock_source_stream.tell.side_effect = (10, 500)
47+
class TestEncryptionStream(object):
48+
def _mock_key_provider(self):
49+
mock_key_provider = MagicMock()
50+
mock_key_provider.__class__ = MasterKeyProvider
51+
return mock_key_provider
5352

54-
self.mock_key_provider = MagicMock()
55-
self.mock_key_provider.__class__ = MasterKeyProvider
53+
def _mock_source_stream(self):
54+
mock_source_stream = MagicMock()
55+
mock_source_stream.__class__ = io.IOBase
56+
mock_source_stream.tell.side_effect = (10, 500)
57+
return mock_source_stream
5658

57-
self.mock_line_length = MagicMock()
58-
self.mock_line_length.__class__ = int
59-
60-
self.mock_source_length = MagicMock()
61-
self.mock_source_length.__class__ = int
59+
@pytest.fixture(autouse=True)
60+
def apply_fixtures(self):
61+
self.mock_key_provider = self._mock_key_provider()
62+
self.mock_source_stream = self._mock_source_stream()
6263

6364
def test_read_bytes_enforcement(self):
6465
class TestStream(_EncryptionStream):
@@ -67,19 +68,23 @@ class TestStream(_EncryptionStream):
6768
def _prep_message(self):
6869
pass
6970

70-
with six.assertRaisesRegex(self, TypeError, "Can't instantiate abstract class TestStream"):
71+
with pytest.raises(TypeError) as excinfo:
7172
TestStream()
7273

74+
excinfo.match("Can't instantiate abstract class TestStream")
75+
7376
def test_prep_message_enforcement(self):
7477
class TestStream(_EncryptionStream):
7578
_config_class = MockClientConfig
7679

7780
def _read_bytes(self):
7881
pass
7982

80-
with six.assertRaisesRegex(self, TypeError, "Can't instantiate abstract class TestStream"):
83+
with pytest.raises(TypeError) as excinfo:
8184
TestStream()
8285

86+
excinfo.match("Can't instantiate abstract class TestStream")
87+
8388
def test_config_class_enforcement(self):
8489
class TestStream(_EncryptionStream):
8590
def _read_bytes(self):
@@ -88,29 +93,32 @@ def _read_bytes(self):
8893
def _prep_message(self):
8994
pass
9095

91-
with six.assertRaisesRegex(self, TypeError, "Can't instantiate abstract class TestStream"):
96+
with pytest.raises(TypeError) as excinfo:
9297
TestStream()
9398

99+
excinfo.match("Can't instantiate abstract class TestStream")
100+
94101
def test_new_with_params(self):
102+
mock_int_sentinel = MagicMock(__class__=int)
95103
mock_stream = MockEncryptionStream(
96104
source=self.mock_source_stream,
97105
key_provider=self.mock_key_provider,
98106
mock_read_bytes=sentinel.read_bytes,
99-
line_length=self.mock_line_length,
100-
source_length=self.mock_source_length,
107+
line_length=io.DEFAULT_BUFFER_SIZE,
108+
source_length=mock_int_sentinel,
101109
)
102110
assert mock_stream.config == MockClientConfig(
103111
source=self.mock_source_stream,
104112
key_provider=self.mock_key_provider,
105113
mock_read_bytes=sentinel.read_bytes,
106-
line_length=self.mock_line_length,
107-
source_length=self.mock_source_length,
114+
line_length=io.DEFAULT_BUFFER_SIZE,
115+
source_length=mock_int_sentinel,
108116
)
109117
assert mock_stream.bytes_read == 0
110118
assert mock_stream.output_buffer == b""
111119
assert not mock_stream._message_prepped
112120
assert mock_stream.source_stream is self.mock_source_stream
113-
assert mock_stream._stream_length is self.mock_source_length
121+
assert mock_stream._stream_length is mock_int_sentinel
114122
assert mock_stream.line_length == io.DEFAULT_BUFFER_SIZE
115123

116124
def test_new_with_config(self):
@@ -154,7 +162,8 @@ class CustomUnknownError(Exception):
154162
mock_stream = MockEncryptionStream(
155163
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
156164
)
157-
with self.assertRaises(CustomUnknownError):
165+
166+
with pytest.raises(CustomUnknownError):
158167
mock_stream.__exit__(None, None, None)
159168

160169
def test_stream_length(self):
@@ -173,9 +182,12 @@ def test_stream_length_unsupported(self):
173182
mock_stream = MockEncryptionStream(
174183
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
175184
)
176-
with six.assertRaisesRegex(self, aws_encryption_sdk.exceptions.NotSupportedError, "Unexpected exception!"):
185+
186+
with pytest.raises(aws_encryption_sdk.exceptions.NotSupportedError) as excinfo:
177187
mock_stream.stream_length # pylint: disable=pointless-statement
178188

189+
excinfo.match("Unexpected exception!")
190+
179191
def test_header_property(self):
180192
mock_prep_message = MagicMock()
181193
mock_stream = MockEncryptionStream(
@@ -205,31 +217,37 @@ def test_read_closed(self):
205217
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
206218
)
207219
mock_stream.close()
208-
with six.assertRaisesRegex(self, ValueError, "I/O operation on closed file"):
220+
221+
with pytest.raises(ValueError) as excinfo:
209222
mock_stream.read()
210223

211-
def test_read_b(self):
224+
excinfo.match("I/O operation on closed file")
225+
226+
@pytest.mark.parametrize("bytes_to_read", range(1, 11))
227+
def test_read_b(self, bytes_to_read):
212228
mock_stream = MockEncryptionStream(
213229
source=io.BytesIO(VALUES["data_128"]),
214230
key_provider=self.mock_key_provider,
215231
mock_read_bytes=sentinel.read_bytes,
216232
)
233+
data = b"1234567890"
217234
mock_stream._read_bytes = MagicMock()
218-
mock_stream.output_buffer = b"1234567890"
219-
test = mock_stream.read(5)
220-
mock_stream._read_bytes.assert_called_once_with(5)
221-
assert test == b"12345"
222-
assert mock_stream.output_buffer == b"67890"
223-
224-
def test_read_all(self):
235+
mock_stream.output_buffer = copy.copy(data)
236+
test = mock_stream.read(bytes_to_read)
237+
mock_stream._read_bytes.assert_called_once_with(bytes_to_read)
238+
assert test == data[:bytes_to_read]
239+
assert mock_stream.output_buffer == data[bytes_to_read:]
240+
241+
@pytest.mark.parametrize("bytes_to_read", (None, -1, -99))
242+
def test_read_all(self, bytes_to_read):
225243
mock_stream = MockEncryptionStream(
226244
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
227245
)
228246
mock_stream._stream_length = 5
229247
mock_stream.output_buffer = b"1234567890"
230248
mock_stream.source_stream = MagicMock()
231249
type(mock_stream.source_stream).closed = PropertyMock(side_effect=(False, False, True))
232-
test = mock_stream.read()
250+
test = mock_stream.read(bytes_to_read)
233251
assert test == b"1234567890"
234252

235253
def test_read_all_empty_source(self):
@@ -262,23 +280,32 @@ def test_writelines(self):
262280
mock_stream = MockEncryptionStream(
263281
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
264282
)
265-
with six.assertRaisesRegex(self, NotImplementedError, "writelines is not available for this object"):
283+
284+
with pytest.raises(NotImplementedError) as excinfo:
266285
mock_stream.writelines(None)
267286

287+
excinfo.match("writelines is not available for this object")
288+
268289
def test_write(self):
269290
mock_stream = MockEncryptionStream(
270291
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
271292
)
272-
with six.assertRaisesRegex(self, NotImplementedError, "write is not available for this object"):
293+
294+
with pytest.raises(NotImplementedError) as excinfo:
273295
mock_stream.write(None)
274296

297+
excinfo.match("write is not available for this object")
298+
275299
def test_seek(self):
276300
mock_stream = MockEncryptionStream(
277301
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
278302
)
279-
with six.assertRaisesRegex(self, NotImplementedError, "seek is not available for this object"):
303+
304+
with pytest.raises(NotImplementedError) as excinfo:
280305
mock_stream.seek(None)
281306

307+
excinfo.match("seek is not available for this object")
308+
282309
def test_readline(self):
283310
test_line = "TEST_LINE_AAAA"
284311
test_line_length = len(test_line)
@@ -319,7 +346,8 @@ def test_next_stream_closed(self):
319346
source=self.mock_source_stream, key_provider=self.mock_key_provider, mock_read_bytes=sentinel.read_bytes
320347
)
321348
mock_stream.close()
322-
with self.assertRaises(StopIteration):
349+
350+
with pytest.raises(StopIteration):
323351
mock_stream.next()
324352

325353
def test_next_source_stream_closed_and_buffer_empty(self):
@@ -328,7 +356,8 @@ def test_next_source_stream_closed_and_buffer_empty(self):
328356
)
329357
self.mock_source_stream.closed = True
330358
mock_stream.output_buffer = b""
331-
with self.assertRaises(StopIteration):
359+
360+
with pytest.raises(StopIteration):
332361
mock_stream.next()
333362

334363
@patch("aws_encryption_sdk.streaming_client._EncryptionStream.closed", new_callable=PropertyMock)

tox.ini

+1
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ commands =
5454
accept: {[testenv:base-command]commands} test/ -m accept
5555
examples: {[testenv:base-command]commands} examples/test/ -m examples
5656
all: {[testenv:base-command]commands} test/ examples/test/
57+
manual: {[testenv:base-command]commands}
5758

5859
# Verify that local tests work without environment variables present
5960
[testenv:nocmk]

0 commit comments

Comments
 (0)