def test_azure_service_init_uses_custom_secrets(self, azure_service):
handler = AzureResultHandler(container="bob",
azure_credentials_secret="MY_FOO")
with prefect.context(secrets=dict(
MY_FOO=dict(ACCOUNT_NAME=1, ACCOUNT_KEY=999))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_service()
assert handler.container == "bob"
assert azure_service.call_args[1] == {
"account_name": 1,
"account_key": 999
}
def test_azure_service_init_uses_secrets(self, azure_service):
handler = AzureResultHandler(container="bob")
assert handler.container == "bob"
assert azure_service.called is False
with prefect.context(secrets=dict(
AZ_CREDENTIALS=dict(ACCOUNT_NAME="1", ACCOUNT_KEY="42"))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_service()
assert azure_service.call_args[1] == {
"account_name": "1",
"account_key": "42"
}
def test_azure_service_writes_to_blob_prefixed_by_date_suffixed_by_prefect(
self, azure_service):
handler = AzureResultHandler(container="foo")
with prefect.context(secrets=dict(
AZ_CREDENTIALS=dict(ACCOUNT_NAME=1, ACCOUNT_KEY=42))):
with set_temporary_config({"cloud.use_local_secrets": True}):
uri = handler.write("so-much-data")
used_uri = azure_service.return_value.create_blob_from_text.call_args[
1]["blob_name"]
assert used_uri == uri
assert used_uri.startswith(pendulum.now("utc").format("Y/M/D"))
assert used_uri.endswith("prefect_result")
def test_azure_service_init_uses_connection_string_over_secret(
self, azure_service):
handler = AzureResultHandler(container="bob",
azure_credentials_secret="MY_FOO",
connection_string="TEST")
with prefect.context(secrets=dict(
MY_FOO=dict(ACCOUNT_NAME=1, ACCOUNT_KEY=999))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_service()
assert handler.container == "bob"
assert azure_service.call_args[1] == {
"connection_string": "TEST",
}
def test_serialize(self):
handler = AzureResultHandler(container="my-container",
azure_credentials_secret="FOO")
serialized = ResultHandlerSchema().dump(handler)
assert serialized["type"] == "AzureResultHandler"
assert serialized["container"] == "my-container"
assert serialized["azure_credentials_secret"] == "FOO"
def test_roundtrip_never_loads_client(self, monkeypatch):
schema = ResultHandlerSchema()
def raise_me(*args, **kwargs):
raise SyntaxError("oops")
monkeypatch.setattr(AzureResultHandler, "initialize_service", raise_me)
handler = schema.load(
schema.dump(
AzureResultHandler(container="container3",
azure_credentials_secret="FOO")))
assert isinstance(handler, AzureResultHandler)
assert handler.container == "container3"
assert handler.azure_credentials_secret == "FOO"
def __init__(self,
container: str,
connection_string: str = None,
blob_name: str = None) -> None:
self.flows = dict() # type: Dict[str, str]
self._flows = dict() # type: Dict[str, "prefect.core.flow.Flow"]
self.connection_string = connection_string or os.getenv(
"AZURE_STORAGE_CONNECTION_STRING")
self.container = container
self.blob_name = blob_name
result_handler = AzureResultHandler(
connection_string=connection_string, container=container)
super().__init__(result_handler=result_handler)
def test_azure_service_handler_is_pickleable(self):
class service:
def __init__(self, *args, **kwargs):
pass
def __getstate__(self):
raise ValueError("I cannot be pickled.")
with patch.dict(
"sys.modules",
{"azure.storage.blob": MagicMock(BlockBlobService=service)}):
with prefect.context(secrets=dict(
AZ_CREDENTIALS=dict(ACCOUNT_NAME=1, ACCOUNT_KEY=42))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler = AzureResultHandler(container="foo")
res = cloudpickle.loads(cloudpickle.dumps(handler))
assert isinstance(res, AzureResultHandler)
def test_roundtrip(self):
schema = ResultHandlerSchema()
handler = schema.load(
schema.dump(AzureResultHandler(container="container3")))
assert isinstance(handler, AzureResultHandler)
assert handler.container == "container3"
