Skip to content

Commit

Permalink
refactor unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
davidlm committed Jul 25, 2023
1 parent 30d7c12 commit b9334df
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 107 deletions.
1 change: 1 addition & 0 deletions botocore/compress.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def _should_compress_request(config, request_dict, operation_model):
body_size = _get_body_size(request_dict['body'])
min_size = config.request_min_compression_size_bytes
return min_size <= body_size

return False


Expand Down
269 changes: 162 additions & 107 deletions tests/unit/test_compress.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ def _make_op(
return op


OP_NO_COMPRESSION = _make_op()
OP_WITH_COMPRESSION = _make_op(request_compression={'encodings': ['gzip']})
OP_UNKNOWN_COMPRESSION = _make_op(request_compression={'encodings': ['foo']})
OP_MULTIPLE_COMPRESSIONS = _make_op(
Expand Down Expand Up @@ -71,35 +72,49 @@ def _make_op(
)


def request_dict():
def _request_dict(body=None, headers=None):
if body is None:
body = b''
if headers is None:
headers = {}

return {
'body': REQUEST_BODY,
'headers': {'foo': 'bar'},
'body': body,
'headers': headers,
}


def request_dict_with_content_encoding_header():
return {
'body': REQUEST_BODY,
'headers': {'foo': 'bar', 'Content-Encoding': 'identity'},
}
def default_request_dict():
return _request_dict(REQUEST_BODY)


DECOMPRESSION_METHOD_MAP = {'gzip': gzip.decompress}
def request_dict_string():
return _request_dict(REQUEST_BODY.decode('utf-8'))


def request_dict_bytearray():
return _request_dict(bytearray(REQUEST_BODY))

def _assert_compression(body, maybe_compressed_body, encoding):
if hasattr(body, 'read'):
body = body.read()
maybe_compressed_body = maybe_compressed_body.read()
if isinstance(body, str):
body = body.encode('utf-8')
if isinstance(body, dict):
body = urlencode(body, doseq=True, encoding='utf-8').encode('utf-8')
decompress_method = DECOMPRESSION_METHOD_MAP.get(
encoding, lambda body: body

def request_dict_with_content_encoding_header():
return _request_dict(
REQUEST_BODY, {'foo': b'bar', 'Content-Encoding': 'identity'}
)
assert decompress_method(maybe_compressed_body) == body


def request_dict_string_io():
return _request_dict(io.StringIO(REQUEST_BODY.decode('utf-8')))


def request_dict_bytes_io():
return _request_dict(io.BytesIO(REQUEST_BODY))


def request_dict_dict():
return _request_dict({'foo': 'bar'})


DECOMPRESSION_METHOD_MAP = {'gzip': gzip.decompress}


def _bad_compression(body):
Expand All @@ -110,151 +125,191 @@ def _bad_compression(body):
MOCK_COMPRESSION.update(COMPRESSION_MAPPING)


def _assert_compression_body(original_body, compressed_body, encoding):
if hasattr(original_body, 'read'):
original_body = original_body.read()
compressed_body = compressed_body.read()
if isinstance(original_body, str):
original_body = original_body.encode('utf-8')
decompress = DECOMPRESSION_METHOD_MAP[encoding]
assert original_body == decompress(compressed_body)


def _assert_compression_header(request_dict, encoding):
assert (
'headers' in request_dict
and 'Content-Encoding' in request_dict['headers']
and encoding in request_dict['headers']['Content-Encoding']
)


@pytest.mark.parametrize(
'config, request_dict, operation_model, is_compressed, encoding',
'config, request_dict, operation_model, encoding',
[
(
Config(
disable_request_compression=True,
request_min_compression_size_bytes=1000,
),
{'body': b'foo', 'headers': {}},
OP_WITH_COMPRESSION,
False,
None,
),
(
COMPRESSION_CONFIG_128_BYTES,
request_dict(),
default_request_dict(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
Config(
disable_request_compression=False,
request_min_compression_size_bytes=256,
),
request_dict(),
OP_WITH_COMPRESSION,
False,
None,
),
(
Config(request_min_compression_size_bytes=128),
request_dict(),
COMPRESSION_CONFIG_128_BYTES,
default_request_dict(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
COMPRESSION_CONFIG_128_BYTES,
request_dict(),
default_request_dict(),
OP_MULTIPLE_COMPRESSIONS,
True,
'gzip',
),
(
DEFAULT_COMPRESSION_CONFIG,
request_dict(),
default_request_dict(),
STREAMING_OP_WITH_COMPRESSION,
True,
'gzip',
),
(
DEFAULT_COMPRESSION_CONFIG,
request_dict(),
STREAMING_OP_WITH_COMPRESSION_REQUIRES_LENGTH,
False,
None,
),
(
DEFAULT_COMPRESSION_CONFIG,
request_dict(),
_make_op(),
False,
None,
),
(
COMPRESSION_CONFIG_128_BYTES,
request_dict(),
OP_UNKNOWN_COMPRESSION,
False,
None,
request_dict_bytearray(),
OP_WITH_COMPRESSION,
'gzip',
),
(
COMPRESSION_CONFIG_128_BYTES,
{'body': REQUEST_BODY.decode(), 'headers': {}},
request_dict_with_content_encoding_header(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
COMPRESSION_CONFIG_128_BYTES,
{'body': bytearray(REQUEST_BODY), 'headers': {}},
request_dict_string(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
COMPRESSION_CONFIG_128_BYTES,
{'body': io.BytesIO(REQUEST_BODY), 'headers': {}},
request_dict_bytes_io(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
COMPRESSION_CONFIG_128_BYTES,
{'body': io.StringIO(REQUEST_BODY.decode()), 'headers': {}},
request_dict_string_io(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
],
)
def test_maybe_compress(
config,
request_dict,
operation_model,
encoding,
):
original_body = request_dict['body']
maybe_compress_request(config, request_dict, operation_model)
_assert_compression_body(original_body, request_dict['body'], encoding)
_assert_compression_header(request_dict, encoding)


def test_maybe_compress_dict():
request_dict = request_dict_dict()
original_body = request_dict['body']
encoded_body = urlencode(
original_body, doseq=True, encoding='utf-8'
).encode('utf-8')
encoding = 'gzip'
maybe_compress_request(
COMPRESSION_CONFIG_1_BYTE, request_dict, OP_WITH_COMPRESSION
)
_assert_compression_body(encoded_body, request_dict['body'], encoding)
_assert_compression_header(request_dict, encoding)


@pytest.mark.parametrize(
'config, request_dict, operation_model',
[
(
Config(
disable_request_compression=True,
request_min_compression_size_bytes=1000,
),
default_request_dict(),
OP_WITH_COMPRESSION,
),
(
Config(
disable_request_compression=False,
request_min_compression_size_bytes=256,
),
default_request_dict(),
OP_WITH_COMPRESSION,
),
(
DEFAULT_COMPRESSION_CONFIG,
default_request_dict(),
STREAMING_OP_WITH_COMPRESSION_REQUIRES_LENGTH,
),
(
DEFAULT_COMPRESSION_CONFIG,
default_request_dict(),
OP_NO_COMPRESSION,
),
(
COMPRESSION_CONFIG_128_BYTES,
request_dict_with_content_encoding_header(),
default_request_dict(),
OP_UNKNOWN_COMPRESSION,
False,
'foo',
),
(
COMPRESSION_CONFIG_128_BYTES,
request_dict_with_content_encoding_header(),
DEFAULT_COMPRESSION_CONFIG,
request_dict_string(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
COMPRESSION_CONFIG_1_BYTE,
{'body': {'foo': 'bar'}, 'headers': {}},
DEFAULT_COMPRESSION_CONFIG,
request_dict_bytearray(),
OP_WITH_COMPRESSION,
True,
'gzip',
),
(
COMPRESSION_CONFIG_128_BYTES,
{'body': {'foo': 'bar'}, 'headers': {}},
DEFAULT_COMPRESSION_CONFIG,
request_dict_bytes_io(),
OP_WITH_COMPRESSION,
False,
None,
),
(
DEFAULT_COMPRESSION_CONFIG,
request_dict_string_io(),
OP_WITH_COMPRESSION,
),
(
COMPRESSION_CONFIG_128_BYTES,
request_dict_with_content_encoding_header(),
OP_UNKNOWN_COMPRESSION,
),
],
)
def test_maybe_compress(
def test_no_compression(
config,
request_dict,
operation_model,
is_compressed,
encoding,
):
ce_header = request_dict['headers'].get('Content-Encoding')
original_body = request_dict['body']
maybe_compress_request(config, request_dict, operation_model)
_assert_compression(original_body, request_dict['body'], encoding)
assert (
'headers' in request_dict
and 'Content-Encoding' in request_dict['headers']
and encoding in request_dict['headers']['Content-Encoding']
) == is_compressed
assert request_dict['body'] is original_body
assert ce_header == request_dict['headers'].get('Content-Encoding')


def test_dict_no_compression():
request_dict = request_dict_dict()
original_body = request_dict['body']
maybe_compress_request(
COMPRESSION_CONFIG_128_BYTES, request_dict, OP_WITH_COMPRESSION
)
assert request_dict['body'] == urlencode(
original_body, doseq=True, encoding='utf-8'
).encode('utf-8')


@pytest.mark.parametrize('body', [1, object(), None, True, 1.0])
Expand All @@ -267,16 +322,16 @@ def test_maybe_compress_bad_types(body):


@pytest.mark.parametrize(
'body',
[io.StringIO('foo'), io.BytesIO(b'foo')],
'request_dict',
[request_dict_string_io(), request_dict_bytes_io()],
)
def test_body_streams_position_reset(body):
def test_body_streams_position_reset(request_dict):
maybe_compress_request(
COMPRESSION_CONFIG_1_BYTE,
{'body': body, 'headers': {}},
COMPRESSION_CONFIG_128_BYTES,
request_dict,
OP_WITH_COMPRESSION,
)
assert body.tell() == 0
assert request_dict['body'].tell() == 0


def test_only_compress_once():
Expand All @@ -285,4 +340,4 @@ def test_only_compress_once():
maybe_compress_request(
COMPRESSION_CONFIG_128_BYTES, request_dict, OP_WITH_COMPRESSION
)
_assert_compression(REQUEST_BODY, request_dict['body'], 'gzip')
_assert_compression_body(REQUEST_BODY, request_dict['body'], 'gzip')

0 comments on commit b9334df

Please sign in to comment.