Skip to content

Commit

Permalink
Fix S3Hook transfer config arguments validation (#25544)
Browse files Browse the repository at this point in the history
* Fix S3Hook transfer config arguments validation

* Add proper exception type
  • Loading branch information
Taragolis authored Aug 5, 2022
1 parent eceb4cc commit d4f560b
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 6 deletions.
6 changes: 3 additions & 3 deletions airflow/providers/amazon/aws/hooks/s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,12 +128,12 @@ def __init__(
kwargs['client_type'] = 's3'
kwargs['aws_conn_id'] = aws_conn_id

if extra_args and not isinstance(extra_args, dict):
raise ValueError(f"transfer_config_args '{extra_args!r}' must be of type {dict}")
if transfer_config_args and not isinstance(transfer_config_args, dict):
raise TypeError(f"transfer_config_args expected dict, got {type(transfer_config_args).__name__}.")
self.transfer_config = TransferConfig(**transfer_config_args or {})

if extra_args and not isinstance(extra_args, dict):
raise ValueError(f"extra_args '{extra_args!r}' must be of type {dict}")
raise TypeError(f"extra_args expected dict, got {type(extra_args).__name__}.")
self._extra_args = extra_args or {}

super().__init__(*args, **kwargs)
Expand Down
9 changes: 6 additions & 3 deletions tests/providers/amazon/aws/hooks/test_s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,16 +55,19 @@ def test_get_conn(self):
hook = S3Hook()
assert hook.get_conn() is not None

@mock_s3
def test_use_threads_default_value(self):
hook = S3Hook()
assert hook.transfer_config.use_threads is True

@mock_s3
def test_use_threads_set_value(self):
hook = S3Hook(transfer_config_args={"use_threads": False})
assert hook.transfer_config.use_threads is False

@pytest.mark.parametrize("transfer_config_args", [1, True, '{"use_threads": false}'])
def test_transfer_config_args_invalid(self, transfer_config_args):
with pytest.raises(TypeError, match="transfer_config_args expected dict, got .*"):
S3Hook(transfer_config_args=transfer_config_args)

def test_parse_s3_url(self):
parsed = S3Hook.parse_s3_url("s3://test/this/is/not/a-real-key.txt")
assert parsed == ("test", "this/is/not/a-real-key.txt"), "Incorrect parsing of the s3 url"
Expand Down Expand Up @@ -509,7 +512,7 @@ def test_generate_presigned_url(self, s3_bucket):
assert {"AWSAccessKeyId", "Signature", "Expires"}.issubset(set(params.keys()))

def test_should_throw_error_if_extra_args_is_not_dict(self):
with pytest.raises(ValueError):
with pytest.raises(TypeError, match="extra_args expected dict, got .*"):
S3Hook(extra_args=1)

def test_should_throw_error_if_extra_args_contains_unknown_arg(self, s3_bucket):
Expand Down

0 comments on commit d4f560b

Please sign in to comment.