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"], )