def test_get_flow_from_gcs_as_file(self, google_client):
f1 = Flow("flow-1")
f2 = Flow("flow-2")
flow_content = textwrap.dedent("""
from prefect import Flow
f1 = Flow('flow-1')
f2 = Flow('flow-2')
""")
blob_mock = MagicMock(download_as_bytes=MagicMock(
return_value=flow_content))
bucket_mock = MagicMock(get_blob=MagicMock(return_value=blob_mock))
google_client.return_value.get_bucket = MagicMock(
return_value=bucket_mock)
storage = GCS(bucket="awesome-bucket",
key="a-place",
stored_as_script=True)
storage.add_flow(f1)
storage.add_flow(f2)
assert storage.get_flow(f1.name).name == f1.name
assert storage.get_flow(f2.name).name == f2.name
bucket_mock.get_blob.assert_called_with("a-place")
assert blob_mock.download_as_bytes.call_count == 2
def test_get_flow_from_gcs(self, google_client):
f = Flow("awesome-flow")
flow_content = cloudpickle.dumps(f)
blob_mock = MagicMock(download_as_bytes=MagicMock(return_value=flow_content))
bucket_mock = MagicMock(get_blob=MagicMock(return_value=blob_mock))
google_client.return_value.get_bucket = MagicMock(return_value=bucket_mock)
with pytest.raises(ValueError):
storage = GCS(bucket="test")
storage.get_flow()
storage = GCS(bucket="awesome-bucket", key="a-place")
storage.add_flow(f)
fetched_flow = storage.get_flow("a-place")
assert fetched_flow.name == f.name
bucket_mock.get_blob.assert_called_with("a-place")
assert blob_mock.download_as_bytes.call_count == 1
def test_get_flow_from_gcs_as_file(self, google_client):
f = Flow("awesome-flow")
flow_content = """from prefect import Flow\nf=Flow('awesome-flow')"""
blob_mock = MagicMock(download_as_bytes=MagicMock(return_value=flow_content))
bucket_mock = MagicMock(get_blob=MagicMock(return_value=blob_mock))
google_client.return_value.get_bucket = MagicMock(return_value=bucket_mock)
storage = GCS(bucket="awesome-bucket", key="a-place", stored_as_script=True)
storage.add_flow(f)
fetched_flow = storage.get_flow("a-place")
assert fetched_flow.name == f.name
bucket_mock.get_blob.assert_called_with("a-place")
assert blob_mock.download_as_bytes.call_count == 1