def test_gcs_init(self, google_client):
result = GCSResult(bucket="bob")
assert result.value is None
assert result.bucket == "bob"
assert google_client.called is False
result.gcs_bucket()
assert google_client.return_value.bucket.call_args[0][0] == "bob"
def test_gcs_writes_binary_string(self, google_client):
blob = MagicMock()
google_client.return_value.bucket = MagicMock(return_value=MagicMock(
blob=MagicMock(return_value=blob)))
result = GCSResult(bucket="foo", location="nothing/here.txt")
new_result = result.write(None)
assert blob.upload_from_string.called
assert isinstance(blob.upload_from_string.call_args[0][0], bytes)
def test_gcs_uses_custom_secret_name(self, google_client):
result = GCSResult(bucket="foo", credentials_secret="TEST_SECRET")
with prefect.context(secrets=dict(TEST_SECRET=94611)):
with set_temporary_config({"cloud.use_local_secrets": True}):
result.gcs_bucket()
assert google_client.call_args[1]["credentials"] == 94611
def test_gcs_writes_to_blob_using_rendered_template_name(
self, google_client):
bucket = MagicMock()
google_client.return_value.bucket = MagicMock(return_value=bucket)
result = GCSResult(bucket="foo", filepath="{thing}/here.txt")
new_result = result.write("so-much-data", thing=42)
assert bucket.blob.called
assert bucket.blob.call_args[0][0] == "42/here.txt"
def test_gcs_reads_and_updates_location(self, google_client):
bucket = MagicMock()
bucket.blob.return_value.download_as_bytes.return_value = b""
google_client.return_value.bucket = MagicMock(return_value=bucket)
result = GCSResult(bucket="foo", location="{thing}/here.txt")
new_result = result.read("path/to/my/stuff.txt")
assert new_result.location == "path/to/my/stuff.txt"
assert new_result.value is None
def test_gcs_writes_binary_string(self, google_client):
blob = MagicMock()
google_client.return_value.bucket = MagicMock(return_value=MagicMock(
blob=MagicMock(return_value=blob)))
result = GCSResult(bucket="foo", filepath_template="nothing/here.txt")
result.value = None
new_result = result.format()
new_result.write()
assert blob.upload_from_string.called
assert isinstance(blob.upload_from_string.call_args[0][0], str)
def main(register, run):
if register:
schedule = Schedule(clocks=[CronClock("1 19 * * *")])
else:
schedule = None
result = GCSResult(bucket='uuazed-prefect')
with Flow("numerai-reports", schedule, result=result) as flow:
filenames = fetch()
upload_to_gcs(filenames)
flow.storage = Docker(
registry_url="gcr.io/numerai-171710",
python_dependencies=['pandas', 'numerapi', 'pyarrow'],
files={
os.path.abspath("data.py"): "numerai_reports/data.py",
os.path.abspath("settings.py"): "numerai_reports/settings.py",
os.path.abspath("utils.py"): "numerai_reports/utils.py",
},
env_vars={"PYTHONPATH": "$PYTHONPATH:/"},
secrets=["GCP_CREDENTIALS"])
if register:
flow.register(project_name="numerai", labels=["docker"])
if run:
flow.run()
def test_gcs_result_is_pickleable(self, google_client, monkeypatch):
class gcs_bucket:
def __init__(self, *args, **kwargs):
pass
def __getstate__(self):
raise ValueError("I cannot be pickled.")
result = GCSResult("foo")
res = cloudpickle.loads(cloudpickle.dumps(result))
assert isinstance(res, GCSResult)
def __init__(
self, bucket: str, key: str = None, project: str = None, **kwargs: Any
) -> None:
self.flows = dict() # type: Dict[str, str]
self._flows = dict() # type: Dict[str, "Flow"]
self.bucket = bucket
self.key = key
self.project = project
result = GCSResult(bucket=bucket)
super().__init__(result=result, **kwargs)
def __init__(
self,
bucket: str,
key: str = None,
project: str = None,
secrets: List[str] = None,
) -> None:
self.flows = dict() # type: Dict[str, str]
self._flows = dict() # type: Dict[str, "Flow"]
self.bucket = bucket
self.key = key
self.project = project
result = GCSResult(bucket=bucket)
super().__init__(result=result, secrets=secrets)
def __init__(self,
bucket: str,
key: str = None,
project: str = None,
stored_as_script: bool = False,
local_script_path: str = None,
**kwargs: Any) -> None:
self.bucket = bucket
self.key = key
self.project = project
self.local_script_path = local_script_path or prefect.context.get(
"local_script_path", None)
result = GCSResult(bucket=bucket)
super().__init__(result=result,
stored_as_script=stored_as_script,
**kwargs)
def test_gcs_init_with_value(self):
result = GCSResult(value=3, bucket="bob")
assert result.value == 3
def test_gcs_exists_fails_if_format_not_called_first(self, google_client):
result = GCSResult(bucket="foo", filepath_template="nothing/here.txt")
with pytest.raises(ValueError):
result.exists()
