def test_add_upstream_modifies_signature(backup_spec_with_functions):
dag = DAGSpec('pipeline.yaml').to_dag()
dag.render()
fn = dag['raw'].source.primitive
params = dag['raw'].params.to_json_serializable()
dev = CallableInteractiveDeveloper(fn, params)
# add an upstream reference...
nb = dev.to_nb()
nb.cells[-1]['source'] += '\nupstream["some_task"]'
dev.overwrite(nb)
# source must be updated...
source = Path('my_tasks', 'raw', 'functions.py').read_text()
top_lines = '\n'.join(source.splitlines()[:5])
expected = (
'from pathlib import Path\n\n\n'
'def function(product, upstream):\n Path(str(product)).touch()')
assert expected == top_lines
# if we save again, nothing should change
dev.overwrite(nb)
source = Path('my_tasks', 'raw', 'functions.py').read_text()
top_lines = '\n'.join(source.splitlines()[:5])
assert expected == top_lines
def test_remove_upstream_modifies_signature(backup_spec_with_functions):
# by the time we reach this test, my_tasks.raw.functions has alread been
# loaded (previous test), so we force reload to avoid wrongfully reading
# the modified source code in the raw task
from my_tasks.raw import functions
importlib.reload(functions)
dag = DAGSpec('pipeline.yaml').to_dag()
dag.render()
fn = dag['clean'].source.primitive
params = dag['clean'].params.to_json_serializable()
dev = CallableInteractiveDeveloper(fn, params)
nb = dev.to_nb()
# delete upstream reference
del nb.cells[-2]
dev.overwrite(nb)
source = Path('my_tasks', 'clean', 'functions.py').read_text()
top_lines = '\n'.join(source.splitlines()[:5])
expected = ('# adding this to make sure relative imports work '
'fine\nfrom .util import util_touch\n\n\n'
'def function(product):')
assert top_lines == expected
def test_develop_spec_with_local_functions(task_name,
backup_spec_with_functions):
"""
Check we can develop functions defined locally, the sample project includes
relative imports, which should work when generating the temporary notebook
"""
dag = DAGSpec('pipeline.yaml').to_dag()
dag.render()
fn = dag[task_name].source.primitive
params = dag[task_name].params.to_json_serializable()
if sys.platform == 'win32':
# edge case, wee need this to correctly parametrize the notebook
# when running the test on windows
params['product'] = str(params['product']).replace('\\', '\\\\')
with CallableInteractiveDeveloper(fn, params) as tmp_nb:
pm.execute_notebook(tmp_nb, tmp_nb)
