def test_invalid_data_type_raises_error(self, monkeypatch):
task = GCSUpload(bucket="test", blob="blobber")
blob = MagicMock()
client = MagicMock()
client.return_value = MagicMock(get_bucket=MagicMock(
return_value=MagicMock(blob=blob)))
monkeypatch.setattr("prefect.tasks.gcp.storage.get_storage_client",
client)
with pytest.raises(
TypeError,
match="data must be str, bytes or BytesIO: got .* instead"):
task.run([1, 2, 3])
def test_bucket_doesnt_exist_can_be_created_on_upload(self, monkeypatch):
task = GCSUpload(bucket="test", create_bucket=True)
client = MagicMock()
client.return_value = MagicMock(
get_bucket=MagicMock(side_effect=NotFound("no bucket"))
)
monkeypatch.setattr("prefect.tasks.gcp.storage.get_storage_client", client)
task.run(data="empty", credentials={})
task.run(data="empty", bucket="run", credentials={})
assert client.return_value.create_bucket.called
assert client.return_value.create_bucket.call_args_list[0][0][0] == "test"
assert client.return_value.create_bucket.call_args_list[1][0][0] == "run"
def test_blob_name_can_be_overwritten_at_runtime_by_upload(self, monkeypatch):
task = GCSUpload(bucket="test", blob="blobber")
blob = MagicMock()
client = MagicMock()
client.return_value = MagicMock(
get_bucket=MagicMock(return_value=MagicMock(blob=blob))
)
monkeypatch.setattr("prefect.tasks.gcp.storage.get_storage_client", client)
task.run(data="empty", credentials={})
task.run(data="empty", blob="run-time", credentials={})
first, second = blob.call_args_list
assert first[0] == ("blobber",)
assert second[0] == ("run-time",)
def test_upload_initializes_attr_from_kwargs(self, attr):
task = GCSUpload(bucket="bucket", **{attr: "my-value"})
assert task.bucket == "bucket"
assert getattr(task, attr) == "my-value"
