def test_build_plain_function(self):
nb_file = get_tmp_notebook(notebook_source)
builder = NbComponentBuilder('op1', inject_notebook_path=nb_file.name)
func_source = builder.build_component_function_source()
annotation = "def op1() -> NamedTuple('TaskOutput', [('mlpipeline_ui_metadata', 'UI_metadata'), ('mlpipeline_metrics', 'Metrics'), ]):"
self.assertTrue(annotation in func_source)
def test_failing_without_parameters_tag(self):
nb_file = get_tmp_notebook(invalid_notebook_source)
runner = KFNotebookRunner(nb_file.name, kernel_name=kernel)
with self.assertRaises(Exception):
runner.run()
nb_file.close()
def test_enabled_notebook_inputs(self):
nb_file = get_tmp_notebook(notebook_source)
runner = KFNotebookRunner(nb_file.name, kernel_name=kernel)
runner.run()
self.assertRegex(runner.notebook_html_output, r'.*TEST-COMMENT.*')
nb_file.close()
def test_build_function_with_input_output_artifacts(self):
nb_file = get_tmp_notebook(notebook_source)
builder = NbComponentBuilder('op1', inject_notebook_path=nb_file.name)
builder.add_input_artifact('a_in')
builder.add_output_artifact('a_out')
func = builder.build_component_function()
self.assertEqual(type(func.__annotations__['a_in']), type(InputPath()))
self.assertEqual(type(func.__annotations__['a_out']), type(OutputPath()))
def test_injectind_and_overwriting_inputs(self):
nb_file = get_tmp_notebook(notebook_source)
runner = KFNotebookRunner(nb_file.name,
inject_params={'a': 1},
kernel_name=kernel)
runner.run()
self.assertEqual(runner.outputs['a'], 1)
nb_file.close()
def test_build_function_with_input_output_params(self):
nb_file = get_tmp_notebook(notebook_source)
builder = NbComponentBuilder('op1', inject_notebook_path=nb_file.name)
builder.add_input_param('a', int, default_value=1)
builder.add_output_param('x', float)
func = builder.build_component_function()
self.assertEqual(func.__annotations__['a'], int)
self.assertEqual(list(func.__annotations__['return'].__annotations__.keys()),
['mlpipeline_ui_metadata', 'mlpipeline_metrics', 'x'])
self.assertEqual(func.__annotations__['return'].__annotations__['x'], float)
def test_executes_without_params(self):
nb_file = get_tmp_notebook(notebook_source)
runner = KFNotebookRunner(nb_file.name, kernel_name=kernel)
runner.run()
self.assertTrue(runner.notebook_html_output)
self.assertEqual(runner.outputs, {'a': 11})
self.assertEqual(runner.metrics, {'accuracy': 1})
nb_pref = re.sub('.py$', '', nb_file.name)
self.assertTrue(os.path.exists(nb_pref + '_inject_output.ipynb'))
self.assertTrue(os.path.exists(nb_pref + '_inject_output_out.ipynb'))
nb_file.close()
def test_injecting_notebook_code(self):
nb_file = get_tmp_notebook('NOTEBOOK SOURCE CODE')
builder = NbComponentBuilder('op1', inject_notebook_path=nb_file.name)
self.assertTrue('NOTEBOOK SOURCE CODE' in builder.extra_code_builder.get_code)
def test_build_function(self):
nb_file = get_tmp_notebook(notebook_source)
builder = NbComponentBuilder('op1', inject_notebook_path=nb_file.name)
x = builder.build_component_function()
self.assertEquals(x.__name__, 'op1')