def test_boto3_client_is_created_with_session(self, mocked_boto_client):
""" Tests the fix for #3925 """
task = S3Upload("test")
task.run("key")
assert (
"session.Session()" in mocked_boto_client.return_value._extract_mock_name()
)
def test_raises_on_invalid_compression_method(self, monkeypatch):
task = S3Upload("test")
client = MagicMock()
boto3 = MagicMock(client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.utilities.aws.boto3", boto3)
with pytest.raises(ValueError, match="gz_fake"):
task.run(b"data", compression="gz_fake")
def test_gzip_compression(self, mocked_boto_client):
task = S3Upload("bucket")
byte_string = b"col1,col2,col3\nfake,data,1\nfalse,info,2\n"
task.run(byte_string, key="key", compression="gzip")
args, kwargs = mocked_boto_client.upload_fileobj.call_args_list[0]
gzip_data_stream = args[0]
assert gzip.decompress(gzip_data_stream.read()) == byte_string
def test_generated_key_is_str(self, mocked_boto_client):
task = S3Upload(bucket="test")
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(AWS_CREDENTIALS={
"ACCESS_KEY": "42",
"SECRET_ACCESS_KEY": "99"
})):
task.run(data="")
assert type(
mocked_boto_client.upload_fileobj.call_args[1]["Key"]) == str
def test_creds_default_to_environment(self, monkeypatch):
task = S3Upload(bucket="bob")
client = MagicMock()
boto3 = MagicMock(client=client)
monkeypatch.setattr("prefect.tasks.aws.s3.boto3", boto3)
task.run(data="")
kwargs = client.call_args[1]
assert kwargs == {
"aws_access_key_id": None,
"aws_secret_access_key": None
}
def test_generated_key_is_str(self, monkeypatch):
task = S3Upload(bucket="test")
client = MagicMock()
boto3 = MagicMock(client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.utilities.aws.boto3", boto3)
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(AWS_CREDENTIALS={
"ACCESS_KEY": "42",
"SECRET_ACCESS_KEY": "99"
})):
task.run(data="")
assert type(client.upload_fileobj.call_args[1]["Key"]) == str
def test_gzip_compression(self, monkeypatch):
task = S3Upload("bucket")
byte_string = b"col1,col2,col3\nfake,data,1\nfalse,info,2\n"
client = MagicMock()
boto3 = MagicMock(client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.utilities.aws.boto3", boto3)
task.run(byte_string, key="key", compression="gzip")
args, kwargs = client.upload_fileobj.call_args_list[0]
gzip_data_stream = args[0]
assert gzip.decompress(gzip_data_stream.read()) == byte_string
def test_creds_are_pulled_from_secret(self, monkeypatch):
task = S3Upload(bucket="bob")
client = MagicMock()
boto3 = MagicMock(client=client)
monkeypatch.setattr("prefect.tasks.aws.s3.boto3", boto3)
with set_temporary_config({"cloud.use_local_secrets": True}):
with prefect.context(secrets=dict(AWS_CREDENTIALS={
"ACCESS_KEY": "42",
"SECRET_ACCESS_KEY": "99"
})):
task.run(data="")
kwargs = client.call_args[1]
assert kwargs == {
"aws_access_key_id": "42",
"aws_secret_access_key": "99"
}
def test_raises_if_bucket_not_eventually_provided(self):
task = S3Upload()
with pytest.raises(ValueError, match="bucket"):
task.run(data="")
def test_initialization_passes_to_task_constructor(self):
task = S3Upload(name="test", tags=["AWS"])
assert task.name == "test"
assert task.tags == {"AWS"}
def test_initialization(self):
task = S3Upload()
def test_initialization(self):
S3Upload()
def test_raises_on_invalid_compression_method(self, mocked_boto_client):
task = S3Upload("test")
with pytest.raises(ValueError, match="gz_fake"):
task.run(b"data", compression="gz_fake")
def test_raises_if_bucket_not_eventually_provided(self):
task = S3Upload()
with pytest.raises(ValueError) as exc:
task.run(data="")
assert "bucket" in str(exc.value)
def test_initialization(self):
task = S3Upload()
assert task.aws_credentials_secret == "AWS_CREDENTIALS"
def test_initialization(self):
task = S3Upload()
assert task.aws_credentials_secret is None
