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"