def test_in_notebook_manager_bad_complex_load_parameter():
run_id = str(uuid.uuid4())
marshal_dir = tempfile.mkdtemp()
try:
with tempfile.NamedTemporaryFile() as output_log_file:
dagstermill.deregister_repository()
context_dict = {
'run_id': run_id,
'mode': 'default',
'solid_def_name': 'hello_world',
'pipeline_name': 'hello_world_pipeline',
'marshal_dir': marshal_dir,
'environment_config': {},
'output_log_path': output_log_file.name,
'input_name_type_dict': {'a': SerializableRuntimeType.NONE},
'output_name_type_dict': {},
}
with pytest.raises(
DagstermillError,
match=(
'If Dagstermill solids have inputs that require serialization strategies that '
),
):
dagstermill.populate_context(context_dict)
finally:
shutil.rmtree(marshal_dir)
def test_in_notebook_manager_load_parameter():
run_id = str(uuid.uuid4())
marshal_dir = tempfile.mkdtemp()
try:
with tempfile.NamedTemporaryFile() as output_log_file:
repository_def = define_example_repository()
dagstermill.register_repository(repository_def)
context_dict = {
'run_id': run_id,
'mode': 'default',
'solid_def_name': 'add_two_numbers',
'pipeline_name': 'test_add_pipeline',
'marshal_dir': marshal_dir,
'environment_config': {},
'output_log_path': output_log_file.name,
'input_name_type_dict': {},
'output_name_type_dict': {},
}
dagstermill.populate_context(context_dict)
dagstermill.load_parameter('a', 7)
finally:
shutil.rmtree(marshal_dir)
def test_in_notebook_manager_bad_load_parameter():
class Foo:
pass
run_id = str(uuid.uuid4())
marshal_dir = tempfile.mkdtemp()
try:
with tempfile.NamedTemporaryFile() as output_log_file:
repository_def = define_example_repository()
dagstermill.register_repository(repository_def)
context_dict = {
'run_id': run_id,
'mode': 'default',
'solid_def_name': 'hello_world',
'pipeline_name': 'hello_world_pipeline',
'marshal_dir': marshal_dir,
'environment_config': {},
'output_log_path': output_log_file.name,
'input_name_type_dict': {},
'output_name_type_dict': {},
}
dagstermill.populate_context(context_dict)
with pytest.raises(KeyError):
dagstermill.load_parameter('garble', Foo())
finally:
shutil.rmtree(marshal_dir)
def test_in_notebook_manager_load_parameter_pickleable():
run_id = str(uuid.uuid4())
marshal_dir = tempfile.mkdtemp()
dagstermill.deregister_repository()
try:
with tempfile.NamedTemporaryFile() as output_log_file:
with tempfile.NamedTemporaryFile() as pickle_file:
pickle.dump(7, pickle_file)
pickle_file.seek(0)
context_dict = {
'run_id': run_id,
'mode': 'default',
'solid_def_name': 'add_two_numbers',
'pipeline_name': 'test_add_pipeline',
'marshal_dir': marshal_dir,
'environment_config': {},
'output_log_path': output_log_file.name,
'input_name_type_dict': {
'a': SerializableRuntimeType.PICKLE_SERIALIZABLE,
'b': SerializableRuntimeType.ANY,
},
'output_name_type_dict': {},
}
dagstermill.populate_context(context_dict)
dagstermill.load_parameter('a', pickle_file.name)
with pytest.raises(
DagstermillError, match='loading parameter b resulted in an error'
):
dagstermill.load_parameter('b', 9)
finally:
shutil.rmtree(marshal_dir)