def __init__(self, name=None, clients=None, executor='serial'):
self._G = nx.DiGraph()
self.name = name or 'No name'
self._logger = logging.getLogger(__name__)
self._clients = DAGClients(clients)
self.__exec_status = DAGStatus.WaitingRender
if executor == 'serial':
self._executor = executors.Serial()
elif executor == 'parallel':
self._executor = executors.Parallel()
elif isinstance(executor, executors.abc.Executor):
self._executor = executor
else:
raise TypeError(
'executor must be "serial", "parallel" or '
'an instance of executors.Executor, got type {}'.format(
type(executor)))
self.on_finish = None
self.on_failure = None
self.on_render = None
self._available_callback_kwargs = {'dag': self}
self._params = DAGConfiguration()
# task access differ using .dag.differ
self.differ = self._params.differ
self.serializer = None
self.unserializer = None
def test_error_does_not_suggest_if_key_does_not_exist():
clients = DAGClients()
with pytest.raises(KeyError) as excinfo:
clients['sqlscript']
assert "Did you mean 'SQLScript'?" not in str(excinfo.value)
def test_error_if_setting_invalid_key():
clients = DAGClients()
with pytest.raises(ValueError) as excinfo:
clients[object] = Mock()
assert ('DAG client keys must be Tasks or '
'Products, value <class \'object\'> is not') == str(excinfo.value)
def test_error_setitem_invalid_str_with_typo(typo, expected):
clients = DAGClients()
mock = Mock()
with pytest.raises(ValueError) as excinfo:
clients[typo] = mock
assert f"Did you mean {expected!r}?" in str(excinfo.value)
def test_setitem_and_getitem_with_str():
clients = DAGClients()
mock = Mock()
clients['SQLScript'] = mock
assert clients[SQLScript] is mock
assert clients['SQLScript'] is mock
def test_error_setitem_invalid_str():
clients = DAGClients()
mock = Mock()
with pytest.raises(ValueError) as excinfo:
clients['invalid_name'] = mock
expected = (f"Could not set DAG-level client {mock!r}. 'invalid_name' "
"is not a valid Task or Product class name")
assert str(excinfo.value) == expected
def test_error_getitem_invalid_str_with_typo(typo, expected, class_):
clients = DAGClients()
mock = Mock()
clients[class_] = mock
with pytest.raises(KeyError) as excinfo:
clients[typo]
expected = f"{typo!r}. Did you mean {expected!r}?"
assert expected in str(excinfo.value)
def test_initializes_dotted_path_spec(tmp_directory, tmp_imports):
Path('my_testing_clients.py').write_text("""
from ploomber.clients import LocalStorageClient
def get_client():
return LocalStorageClient('backup', path_to_project_root='.')
""")
clients = DAGClients()
# this happens when using the spec API and lazy load is turned on
clients[File] = DottedPath('my_testing_clients.get_client')
client = clients[File]
assert isinstance(client, LocalStorageClient)
assert clients[File] is client
def test_repr():
clients = DAGClients()
clients[SQLScript] = 1
expected = "DAGClients({<class 'ploomber.tasks.sql.SQLScript'>: 1})"
assert repr(clients) == expected
def test_iter():
assert list(DAGClients()) == []
