def test_basic_yield_multiple_materializations(): SomeDagsterType = create_any_type( name='SomeType', output_materialization_config=yield_two_materializations) @lambda_solid(output_def=OutputDefinition(SomeDagsterType)) def return_one(): return 1 pipeline_def = PipelineDefinition(name='single_int_output_pipeline', solid_defs=[return_one]) result = execute_pipeline(pipeline_def, environment_dict={ 'solids': { 'return_one': { 'outputs': [{ 'result': 2 }] } } }) assert result.success event_types = [event.event_type_value for event in result.event_list] assert 2 == (sum([ True for event_type in event_types if event_type == DagsterEventType.STEP_MATERIALIZATION.value ]))
def test_user_error_boundary_storage_plugin(): class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init @classmethod def compatible_with_storage_def(cls, system_storage_def): return True @classmethod def set_object(cls, intermediate_store, obj, context, runtime_type, paths): raise UserError() @classmethod def get_object(cls, intermediate_store, context, runtime_type, paths): raise UserError() CustomDagsterType = create_any_type(name='CustomType', auto_plugins=[CustomStoragePlugin]) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return 'hello' @pipeline def plugin_pipeline(): output_solid() pipeline_result = execute_pipeline(plugin_pipeline, {'storage': { 'filesystem': {} }}, raise_on_error=False) assert not pipeline_result.success
def test_basic_yield_multiple_materializations(): SomeDagsterType = create_any_type(name="SomeType", materializer=yield_two_materializations) @lambda_solid(output_def=OutputDefinition(SomeDagsterType)) def return_one(): return 1 pipeline_def = PipelineDefinition(name="single_int_output_pipeline", solid_defs=[return_one]) result = execute_pipeline( pipeline_def, run_config={"solids": { "return_one": { "outputs": [{ "result": 2 }] } }}) assert result.success event_types = [event.event_type_value for event in result.event_list] assert 2 == (sum([ True for event_type in event_types if event_type == DagsterEventType.ASSET_MATERIALIZATION.value ]))
def test_user_error_boundary_storage_plugin(): class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init @classmethod def compatible_with_storage_def(cls, intermediate_storage_def): return True @classmethod def set_intermediate_object(cls, intermediate_storage, context, dagster_type, step_output_handle, value): raise UserError() @classmethod def get_intermediate_object(cls, intermediate_storage, context, dagster_type, step_output_handle): raise UserError() CustomDagsterType = create_any_type(name="CustomType", auto_plugins=[CustomStoragePlugin]) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return "hello" @pipeline def plugin_pipeline(): output_solid() pipeline_result = execute_pipeline(plugin_pipeline, {"storage": { "filesystem": {} }}, raise_on_error=False) assert not pipeline_result.success
def test_user_error_boundary_output_materialization(): @output_materialization_config(String) def materialize(context, *_args, **_kwargs): raise UserError() CustomDagsterType = create_any_type( name='CustomType', output_materialization_config=materialize) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return 'hello' @pipeline def output_materialization_pipeline(): output_solid() pipeline_result = execute_pipeline( output_materialization_pipeline, {'solids': { 'output_solid': { 'outputs': [{ 'result': 'hello' }] } }}, raise_on_error=False, ) assert not pipeline_result.success
def define_materialization_pipeline(should_require_resources=True, resources_initted=None): if resources_initted is None: resources_initted = {} @resource def resource_a(_): resources_initted['a'] = True yield 'A' @dagster_type_materializer( String, required_resource_keys={'a'} if should_require_resources else set()) def materialize(context, *_args, **_kwargs): assert context.resources.a == 'A' return AssetMaterialization('hello') CustomDagsterType = create_any_type(name='CustomType', materializer=materialize) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return 'hello' @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})]) def output_pipeline(): output_solid() return output_pipeline
def test_user_error_boundary_output_materialization(): @dagster_type_materializer(String) def materialize(context, *_args, **_kwargs): raise UserError() CustomDagsterType = create_any_type(name="CustomType", materializer=materialize) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return "hello" @pipeline def output_materialization_pipeline(): output_solid() pipeline_result = execute_pipeline( output_materialization_pipeline, {"solids": { "output_solid": { "outputs": [{ "result": "hello" }] } }}, raise_on_error=False, ) assert not pipeline_result.success
def define_plugin_pipeline( should_require_resources=True, resources_initted=None, compatible_storage=True, mode_defines_resource=True, ): if resources_initted is None: resources_initted = {} class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init @classmethod def compatible_with_storage_def(cls, _): return compatible_storage @classmethod def set_intermediate_object(cls, intermediate_storage, context, dagster_type, step_output_handle, value): assert context.resources.a == "A" return intermediate_storage.set_intermediate_object( dagster_type, step_output_handle, value) @classmethod def get_intermediate_object(cls, intermediate_storage, context, dagster_type, step_output_handle): assert context.resources.a == "A" return intermediate_storage.get_intermediate_object( dagster_type, step_output_handle) @classmethod def required_resource_keys(cls): return {"a"} if should_require_resources else set() @resource def resource_a(_): resources_initted["a"] = True yield "A" CustomDagsterType = create_any_type(name="CustomType", auto_plugins=[CustomStoragePlugin]) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return "hello" if mode_defines_resource: mode_defs = [ModeDefinition(resource_defs={"a": resource_a})] else: mode_defs = [ModeDefinition()] @pipeline(mode_defs=mode_defs) def plugin_pipeline(): output_solid() return plugin_pipeline
def test_basic_bad_output_materialization(): SomeDagsterType = create_any_type(name="SomeType", materializer=return_int) @lambda_solid(output_def=OutputDefinition(SomeDagsterType)) def return_one(): return 1 pipeline_def = PipelineDefinition(name="single_int_output_pipeline", solid_defs=[return_one]) with pytest.raises( DagsterInvariantViolationError, match="You must return an AssetMaterialization" ): execute_pipeline( pipeline_def, run_config={"solids": {"return_one": {"outputs": [{"result": 2}]}}} )
def define_plugin_pipeline(should_require_resources=True, resources_initted=None, compatible_storage=True): if resources_initted is None: resources_initted = {} class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init @classmethod def compatible_with_storage_def(cls, system_storage_def): return compatible_storage @classmethod def set_object(cls, intermediate_store, obj, context, dagster_type, paths): assert context.resources.a == 'A' return intermediate_store.set_object(obj, context, dagster_type, paths) @classmethod def get_object(cls, intermediate_store, context, dagster_type, paths): assert context.resources.a == 'A' return intermediate_store.get_object(context, dagster_type, paths) @classmethod def required_resource_keys(cls): return {'a'} if should_require_resources else set() @resource def resource_a(_): resources_initted['a'] = True yield 'A' CustomDagsterType = create_any_type(name='CustomType', auto_plugins=[CustomStoragePlugin]) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return 'hello' @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})]) def plugin_pipeline(): output_solid() return plugin_pipeline
def test_materialize_missing_resource_fails(): @dagster_type_materializer(String, required_resource_keys={"a"}) def materialize(context, *_args, **_kwargs): assert context.resources.a == "A" return AssetMaterialization("hello") CustomType = create_any_type(name="CustomType", materializer=materialize) @solid(output_defs=[OutputDefinition(CustomType, "custom_type")]) def custom_type_solid(_): return "A" with pytest.raises( DagsterInvalidDefinitionError, match='required by the materializer on type "CustomType"'): @pipeline def _type_check_pipeline(): custom_type_solid()
def test_basic_bad_output_materialization(): SomeDagsterType = create_any_type(name='SomeType', materializer=return_int) @lambda_solid(output_def=OutputDefinition(SomeDagsterType)) def return_one(): return 1 pipeline_def = PipelineDefinition(name='single_int_output_pipeline', solid_defs=[return_one]) with pytest.raises(DagsterInvariantViolationError, match='You must return an AssetMaterialization'): execute_pipeline(pipeline_def, run_config={ 'solids': { 'return_one': { 'outputs': [{ 'result': 2 }] } } })
def define_composite_materialization_pipeline( should_require_resources=True, resources_initted=None ): if resources_initted is None: resources_initted = {} @resource def resource_a(_): resources_initted['a'] = True yield 'A' @output_materialization_config( String, required_resource_keys={'a'} if should_require_resources else set() ) def materialize(context, *_args, **_kwargs): assert context.resources.a == 'A' return Materialization('hello') CustomDagsterType = create_any_type( name='CustomType', output_materialization_config=materialize ) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return 'hello' wrap_solid = CompositeSolidDefinition( name="wrap_solid", solid_defs=[output_solid], output_mappings=[OutputDefinition(CustomDagsterType).mapping_from('output_solid')], ) @pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})]) def output_pipeline(): wrap_solid() return output_pipeline
def define_composite_materialization_pipeline(should_require_resources=True, resources_initted=None): if resources_initted is None: resources_initted = {} @resource def resource_a(_): resources_initted["a"] = True yield "A" @dagster_type_materializer( String, required_resource_keys={"a"} if should_require_resources else set()) def materialize(context, *_args, **_kwargs): assert context.resources.a == "A" return AssetMaterialization("hello") CustomDagsterType = create_any_type(name="CustomType", materializer=materialize) @solid(output_defs=[OutputDefinition(CustomDagsterType)]) def output_solid(_context): return "hello" wrap_solid = CompositeSolidDefinition( name="wrap_solid", solid_defs=[output_solid], output_mappings=[ OutputDefinition(CustomDagsterType).mapping_from("output_solid") ], ) @pipeline(mode_defs=[ModeDefinition(resource_defs={"a": resource_a})]) def output_pipeline(): wrap_solid() return output_pipeline
assert result.success with open(filename_one, 'r') as ff: value = json.loads(ff.read()) assert value == {'value': 1} with open(filename_two, 'r') as ff: value = json.loads(ff.read()) assert value == {'value': 1} @output_materialization_config(Int) def return_int(*_args, **_kwargs): return 1 SomeDagsterType = create_any_type(name='SomeType', output_materialization_config=return_int) def test_basic_bad_output_materialization(): @lambda_solid(output_def=OutputDefinition(SomeDagsterType)) def return_one(): return 1 pipeline_def = PipelineDefinition(name='single_int_output_pipeline', solid_defs=[return_one]) with pytest.raises(DagsterInvariantViolationError, match='You must return a Materialization'): execute_pipeline( pipeline_def, environment_dict={'solids': {'return_one': {'outputs': [{'result': 2}]}}} )
adls2_resource, create_adls2_client, ) from dagster_azure.blob import create_blob_client class UppercaseSerializationStrategy(SerializationStrategy): # pylint: disable=no-init def serialize(self, value, write_file_obj): return write_file_obj.write(bytes(value.upper().encode("utf-8"))) def deserialize(self, read_file_obj): return read_file_obj.read().decode("utf-8").lower() LowercaseString = create_any_type( "LowercaseString", serialization_strategy=UppercaseSerializationStrategy("uppercase"), ) nettest = pytest.mark.nettest def define_inty_pipeline(should_throw=True): @lambda_solid def return_one(): return 1 @lambda_solid(input_defs=[InputDefinition("num", Int)], output_def=OutputDefinition(Int)) def add_one(num): return num + 1
from dagster.core.types.dagster_type import String as RuntimeString from dagster.core.types.dagster_type import create_any_type, resolve_dagster_type from dagster.core.utils import make_new_run_id from dagster.utils.test import yield_empty_pipeline_context class UppercaseSerializationStrategy(SerializationStrategy): # pylint: disable=no-init def serialize(self, value, write_file_obj): return write_file_obj.write(bytes(value.upper().encode('utf-8'))) def deserialize(self, read_file_obj): return read_file_obj.read().decode('utf-8').lower() LowercaseString = create_any_type( 'LowercaseString', serialization_strategy=UppercaseSerializationStrategy('uppercase'), ) def aws_credentials_present(): return os.getenv('AWS_ACCESS_KEY_ID') and os.getenv( 'AWS_SECRET_ACCESS_KEY') nettest = pytest.mark.nettest def define_inty_pipeline(should_throw=True): @lambda_solid def return_one(): return 1