def test_dagspec_initialization_from_yaml_and_env(tmp_nbs, monkeypatch):
"""
DAGSpec can be initialized with a path to a spec or a dictionary, but
they have a slightly different behavior. This ensure the cli passes
the path, instead of a dictionary
"""
mock_DAGSpec = Mock(wraps=parsers.DAGSpec)
mock_default_path_to_env = Mock(wraps=parsers.default.path_to_env)
mock_EnvDict = Mock(wraps=parsers.EnvDict)
monkeypatch.setattr(sys, 'argv', ['python'])
monkeypatch.setattr(parsers, 'DAGSpec', mock_DAGSpec)
monkeypatch.setattr(parsers.default, 'path_to_env',
mock_default_path_to_env)
monkeypatch.setattr(parsers, 'EnvDict', mock_EnvDict)
parser = CustomParser()
with parser:
pass
dag, args = _custom_command(parser)
# ensure called using the path to the yaml spec
mock_DAGSpec.assert_called_once_with('pipeline.yaml',
env=EnvDict({'sample': False}))
# and EnvDict initialized from env.yaml
mock_EnvDict.assert_called_once_with(str(Path('env.yaml').resolve()))
def main():
parser = CustomParser(description='Show pipeline status')
with parser:
# this command has no static args
pass
dag, args = _custom_command(parser)
print(dag.status())
def main():
parser = CustomParser(description='Make a pipeline report')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the report, defaults to pipeline.html',
default='pipeline.html')
dag, args = _custom_command(parser)
dag.to_markup(path=args.output)
print('Report saved at:', args.output)
def main():
parser = CustomParser(description='Plot a pipeline')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the plot, defaults to pipeline.png',
default='pipeline.png')
dag, args = _custom_command(parser)
dag.plot(output=args.output)
print('Plot saved at:', args.output)
def test_entry_point_from_factory_in_environment_variable(
backup_test_pkg, monkeypatch):
monkeypatch.setattr(sys, 'argv', ['python'])
monkeypatch.setenv('ENTRY_POINT', 'test_pkg.entry.plain_function')
parser = CustomParser()
with parser:
pass
dag, _ = _custom_command(parser)
assert isinstance(dag, DAG)
def main(render_only=False):
parser = CustomParser(description='Build pipeline')
with parser:
parser.add_argument('--force',
'-f',
help='Force execution by ignoring status',
action='store_true',
default=False)
parser.add_argument('--skip-upstream',
'-su',
help='Skip building upstream dependencies. '
'Only applicable when using --partially',
action='store_true',
default=False)
parser.add_argument(
'--partially',
'-p',
help='Build a pipeline partially until certain task',
default=None)
parser.add_argument(
'--debug',
'-d',
help='Drop a debugger session if an exception happens',
action='store_true',
default=False)
dag, args = _custom_command(parser)
# when using the parallel executor from the CLI, ensure we print progress
# to stdout
if isinstance(dag.executor, Parallel):
dag.executor.print_progress = True
if render_only:
dag.render()
else:
if args.partially:
report = dag.build_partially(args.partially,
force=args.force,
debug=args.debug,
skip_upstream=args.skip_upstream)
else:
report = dag.build(force=args.force, debug=args.debug)
if report:
print(report)
return dag
def main():
parser = CustomParser(description='Get task information')
with parser:
parser.add_argument('task_name')
parser.add_argument('--source',
'-s',
help='Print task source code',
action='store_true')
parser.add_argument(
'--build',
'-b',
help='Build task (default if no other option is passed)',
action='store_true')
parser.add_argument('--force',
'-f',
help='Force task build (ignore up-to-date status)',
action='store_true')
parser.add_argument('--status',
'-st',
help='Get task status',
action='store_true')
parser.add_argument('--on-finish',
'-of',
help='Only execute on_finish hook',
action='store_true')
dag, args = _custom_command(parser)
dag.render()
task = dag[args.task_name]
if args.source:
print(task.source)
if args.status:
print(task.status())
if args.on_finish:
task._run_on_finish()
# task if build by default, but when --source or --status are passed,
# the --build flag is required
no_flags = not any((args.build, args.status, args.source, args.on_finish))
if no_flags or args.build:
task.build(force=args.force)
def main():
parser = CustomParser(description='Call an entry point '
'(pipeline.yaml or dotted path to factory)')
with parser:
# this command has no static args
pass
dag, _ = _custom_command(parser)
try:
dag.render()
except Exception:
print('Your dag failed to render, but you can still inspect the '
'object to debug it.\n')
# NOTE: do not use embed here, we must use start_ipytho, see here:
# https://github.com/ipython/ipython/issues/8918
start_ipython(argv=[], user_ns={'dag': dag})
def main():
parser = CustomParser(description='Plot a pipeline')
with parser:
parser.add_argument(
'--output',
'-o',
help='Where to save the plot, defaults to pipeline.png',
default=None)
dag, args = _custom_command(parser)
if args.output is not None:
output = args.output
else:
name = extract_name(args.entry_point)
output = 'pipeline.png' if name is None else f'pipeline.{name}.png'
dag.plot(output=output)
print('Plot saved at:', output)
def test_dagspec_initialization_from_yaml(tmp_nbs_nested, monkeypatch):
"""
DAGSpec can be initialized with a path to a spec or a dictionary, but
they have a slightly different behavior. This checks that we initialize
with the path
"""
mock = Mock(wraps=parsers.DAGSpec)
monkeypatch.setattr(sys, 'argv', ['python'])
monkeypatch.setattr(parsers, 'DAGSpec', mock)
parser = CustomParser()
with parser:
pass
dag, args = _custom_command(parser)
mock.assert_called_once_with('pipeline.yaml')
def main(render_only=False):
parser = CustomParser(description='Build pipeline')
with parser:
parser.add_argument('--force',
'-f',
help='Force execution by ignoring status',
action='store_true',
default=False)
parser.add_argument(
'--partially',
'-p',
help='Build a pipeline partially until certain task',
default=None)
parser.add_argument(
'--debug',
'-d',
help='Drop a debugger session if an exception happens',
action='store_true',
default=False)
dag, args = _custom_command(parser)
if render_only:
dag.render()
else:
if args.partially:
report = dag.build_partially(args.partially,
force=args.force,
debug=args.debug)
else:
report = dag.build(force=args.force, debug=args.debug)
if report:
print(report)
return dag
def __call__(self, parser):
dag, args = _custom_command(parser)
self.dag_mock = MagicMock(wraps=dag)
return self.dag_mock, args