def test_get_flow_azure_runs(monkeypatch):
client = MagicMock(download_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service",
service)
f = Flow("test")
monkeypatch.setattr("cloudpickle.loads", MagicMock(return_value=f))
storage = Azure(container="container")
assert f.name not in storage
storage.add_flow(f)
new_flow = storage.get_flow(f.name)
assert client.download_blob.called
assert f.name in storage
assert isinstance(new_flow, Flow)
assert new_flow.name == "test"
assert len(new_flow.tasks) == 0
state = new_flow.run()
assert state.is_successful()
def test_add_flow_to_azure():
storage = Azure(container="test")
f = Flow("test")
assert f.name not in storage
assert storage.add_flow(f)
assert f.name in storage
def test_get_flow_from_file_azure_runs(monkeypatch):
client = MagicMock(download_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service",
service)
f = Flow("test")
monkeypatch.setattr(
"prefect.storage.azure.extract_flow_from_file",
MagicMock(return_value=f),
)
storage = Azure(container="container", stored_as_script=True)
assert f.name not in storage
flow_location = storage.add_flow(f)
new_flow = storage.get_flow(flow_location)
assert client.download_blob.called
assert f.name in storage
assert isinstance(new_flow, Flow)
assert new_flow.name == "test"
assert len(new_flow.tasks) == 0
state = new_flow.run()
assert state.is_successful()
def test_add_multiple_flows_to_Azure():
storage = Azure(container="container")
f = Flow("test")
g = Flow("testg")
assert f.name not in storage
assert storage.add_flow(f)
assert storage.add_flow(g)
assert f.name in storage
assert g.name in storage
def test_add_flow_to_azure_already_added(monkeypatch):
storage = Azure(container="container")
f = Flow("test")
assert f.name not in storage
assert storage.add_flow(f)
assert f.name in storage
with pytest.raises(ValueError):
storage.add_flow(f)
def test_blob_service_client_property(monkeypatch, secret_name, secret_arg):
connection = MagicMock()
azure = MagicMock(from_connection_string=connection)
monkeypatch.setattr("azure.storage.blob.BlobServiceClient", azure)
with context(secrets={secret_name: secret_arg}):
storage = Azure(container="test", connection_string_secret=secret_name)
client = storage._azure_block_blob_service()
azure_client = storage._azure_block_blob_service
assert storage.connection_string == secret_arg
assert azure_client
connection.assert_called_with(conn_str=secret_arg)
def test_upload_flow_to_azure_blob_name(monkeypatch):
client = MagicMock(upload_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service", service)
storage = Azure(container="container", blob_name="name")
f = Flow("test")
assert storage.add_flow(f)
assert storage.build()
assert service.get_blob_client.call_args[1]["container"] == "container"
assert service.get_blob_client.call_args[1]["blob"] == "name"
def test_upload_flow_to_azure(monkeypatch):
client = MagicMock(upload_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service", service)
storage = Azure(container="container")
f = Flow("test")
assert f.name not in storage
assert storage.add_flow(f)
assert storage.build()
assert client.upload_blob.called
assert f.name in storage
def test_upload_flow_to_azure_blob_overwrite(monkeypatch, overwrite):
client = MagicMock(upload_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service",
service)
storage = Azure(container="container", overwrite=overwrite)
f = Flow("test")
assert storage.add_flow(f)
assert storage.build()
client.upload_blob.assert_called_once_with(unittest.mock.ANY,
overwrite=overwrite)
def test_upload_flow_to_azure(monkeypatch):
client = MagicMock(upload_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service",
service)
storage = Azure(container="container")
f = Flow("test")
assert f.name not in storage
assert storage.add_flow(f)
assert storage.build()
client.upload_blob.assert_called_once_with(flow_to_bytes_pickle(f),
overwrite=False)
assert f.name in storage
def test_get_flow_azure(monkeypatch):
client = MagicMock(download_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service", service)
f = Flow("test")
monkeypatch.setattr("cloudpickle.loads", MagicMock(return_value=f))
storage = Azure(container="container")
assert f.name not in storage
storage.add_flow(f)
assert storage.get_flow(f.name)
assert client.download_blob.called
assert f.name in storage
def test_get_flow_azure_bucket_key(monkeypatch):
client = MagicMock(download_blob=MagicMock())
service = MagicMock(get_blob_client=MagicMock(return_value=client))
monkeypatch.setattr("prefect.storage.Azure._azure_block_blob_service", service)
f = Flow("test")
monkeypatch.setattr("cloudpickle.loads", MagicMock(return_value=f))
storage = Azure(container="container", blob_name="name")
assert f.name not in storage
flow_location = storage.add_flow(f)
assert storage.get_flow(f.name)
assert service.get_blob_client.call_args[1]["container"] == "container"
assert service.get_blob_client.call_args[1]["blob"] == flow_location
def test_build_no_upload_if_file(monkeypatch):
storage = Azure(container="container", stored_as_script=True)
with pytest.raises(ValueError):
storage.build()
storage = Azure(container="container", stored_as_script=True, blob_name="flow.py")
assert storage == storage.build()
def test_blob_service_client_property(monkeypatch, secret_name, secret_arg,
credential):
connection = MagicMock()
azure = MagicMock(from_connection_string=connection)
monkeypatch.setattr("azure.storage.blob.BlobServiceClient", azure)
monkeypatch.setattr("azure.identity.DefaultAzureCredential", credential)
with context(secrets={secret_name: secret_arg}):
storage = Azure(container="test", connection_string_secret=secret_name)
client = storage._azure_block_blob_service()
azure_client = storage._azure_block_blob_service
assert storage.connection_string == secret_arg
assert azure_client
if credential is None:
connection.assert_called_with(conn_str=secret_arg, credential=None)
else:
connection.assert_called_with(conn_str=secret_arg,
credential=credential())
def test_blob_service_client_property(monkeypatch):
connection = MagicMock()
azure = MagicMock(from_connection_string=connection)
monkeypatch.setattr("azure.storage.blob.BlobServiceClient", azure)
storage = Azure(container="test", connection_string="conn")
azure_client = storage._azure_block_blob_service
assert azure_client
connection.assert_called_with(conn_str="conn")
def test_create_azure_storage_init_args():
storage = Azure(
container="test",
connection_string="conn",
blob_name="name",
secrets=["foo"],
)
assert storage
assert storage.flows == dict()
assert storage.container == "test"
assert storage.connection_string == "conn"
assert storage.blob_name == "name"
assert storage.secrets == ["foo"]
def test_serialize_azure_storage():
storage = Azure(container="test")
serialized_storage = storage.serialize()
assert serialized_storage["type"] == "Azure"
def storage(self) -> Storage:
return Azure(self._storage_options)
def test_create_azure_storage():
storage = Azure(container="test")
assert storage
assert storage.logger
"id": "foo"
},
"run_config": run.serialize(),
}),
run,
)
assert env_vars["PREFECT__LOGGING__LEVEL"] == expected_logging_level
@pytest.mark.parametrize(
"storage",
[
Local(directory="test"),
GCS(bucket="test"),
S3(bucket="test"),
Azure(container="test"),
GitLab("test/repo", path="path/to/flow.py"),
Bitbucket(project="PROJECT", repo="test-repo", path="test-flow.py"),
CodeCommit("test/repo", path="path/to/flow.py"),
Webhook(
build_request_kwargs={"url": "test-service/upload"},
build_request_http_method="POST",
get_flow_request_kwargs={"url": "test-service/download"},
get_flow_request_http_method="GET",
),
],
)
def test_local_agent_deploy_processes_valid_storage(storage, monkeypatch):
popen = MagicMock()
monkeypatch.setattr("prefect.agent.local.agent.Popen", popen)
def test_connection_string_property(secret_name, secret_arg):
with context(secrets={secret_name: secret_arg}):
storage = Azure(container="test", connection_string_secret=secret_name)
assert storage.connection_string == secret_arg
if "time0" in list(ds.coords):
cleaned_ds = ds.rename({"time0": "time"})
elif "time1" in list(ds.coords):
cleaned_ds = ds.rename({"time1": "time"})
return cleaned_ds
@task()
def copy_to_azure(file_path: str):
tgt = extract_name_path(file_path)
tgt_map = map_tgt(tgt)
if not zarr_is_complete(tgt_map):
ds = map_and_open_zarr_link(file_path)
cleaned_ds = clean_ds(ds)
copy_cleaned_data(cleaned_ds, tgt_map, overwrite=True)
run_config = KubernetesRun(
cpu_request=3,
memory_request="3Gi",
image="gcr.io/carbonplan/hub-notebook:7252fc3",
labels=["az-eu-west"],
)
storage = Azure("prefect")
with Flow(name="Transfer_ERA5", storage=storage,
run_config=run_config) as flow:
file_pattern_list = create_formatted_links()
copy_to_azure.map(file_pattern_list)
from prefect.storage import Azure
from prefect.run_configs import KubernetesRun
import os
azure_store = Azure(container="prefect")
kubernetes_run = KubernetesRun(
image=os.environ.get("IMAGE_URL"),
job_template_path=os.environ.get("JOB_TEMPLATE_PATH", "./job_template.yaml"),
labels=["prefect"],
)
