How do I mock an asset's value when using materialize?

[sryza]

Is there a way to write a test for a dagster pipeline with a mock value for the first asset?

Keywords: input value input_values

[sryza]

One way to do this is with a function that generates a source asset with an IO manager that returns the value you want to mock in.

from dagster import IOManager, SourceAsset

def mock_asset(asset_key, value):
    class MockIOManager(IOManager):
        def load_input(self, context):
            return value

        def handle_output(self, context, obj):
            raise NotImplementedError()

    return SourceAsset(key=asset_key, io_manager_def=MockIOManager())

Example usage:

from dagster import asset, materialize

@asset
def asset2(asset1):
    return asset1 + 1


result = materialize(assets=[asset2, mock_asset("asset1", 5)])

assert result.asset_value(asset2.key) == 6

[AlanBENlabs]

When using this to mock an asset with a partitions definition, will the mock_asset() also have to include the partitions definition as one of the arguments and pass it to SourceAsset? Or how will this change when you have a partitions definition set on asset1 and asset2 in the pipeline you are testing?

[sryza]

It shouldn't make a big difference:

from dagster import IOManager, SourceAsset


def mock_asset(asset_key, value):
    class MockIOManager(IOManager):
        def load_input(self, context):
            return value

        def handle_output(self, context, obj):
            raise NotImplementedError()

    return SourceAsset(key=asset_key, io_manager_def=MockIOManager())


from dagster import StaticPartitionsDefinition, asset, materialize


@asset(partitions_def=StaticPartitionsDefinition(["foo", "bar"]))
def asset2(asset1):
    return asset1 + 1


result = materialize(assets=[asset2, mock_asset("asset1", 5)], partition_key="foo")

assert result.asset_value(asset2.key) == 6

[GinaStavropoulou]

How can I do this using AssetSpec instead of SourceAsset (since it is going to be deprecated)?

[sryza]

Alas this is a little trickier with AssetSpec, though we'd like to make it easier in the future. For now, I'd recommend just using @asset:

from dagster import IOManager, SourceAsset

def mock_asset(asset_key, value):
    class MockIOManager(IOManager):
        def load_input(self, context):
            return value

        def handle_output(self, context, obj):
            raise NotImplementedError()

    @asset(key=asset_key, io_manager_def=MockIOManager())
    def _asset():
        raise NotImplementedError()

    return _asset

When executing, make sure to only select the asset you want to materialize:

from dagster import asset, materialize

@asset
def asset2(asset1):
    return asset1 + 1


result = materialize(assets=[asset2, mock_asset("asset1", 5)], selection=[asset2])

assert result.asset_value(asset2.key) == 6