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)