https://dagster.io/ logo
#ask-community
Title
# ask-community
c

Chris Histe

09/28/2022, 2:01 PM
[Relates to unit testing and
Context
] Hello there, is there an equivalent of
build_init_resource_context
for
StepExecutionContext
. I’m unit testing a custom IO Manager and getting an error when accessing
_context_.has_partition_key
(InputContext). It seems you can pass a
StepExecutionContext
to
build_input_context
but I’m not sure how to instantiate that class.
Copy code
dagster._core.errors.DagsterInvariantViolationError: Attempting to access step_context, but it was not provided when constructing the InputContext
The error
The documentation says we should not instantiante
StepExecutionContext
so I assume there must a helper function for test.
c

claire

09/28/2022, 11:25 PM
Unfortunately, there currently is not a helper function for fetching a
StepExecutionContext
. The limitation here is that some information, including the partition key, is only available on the pipeline run object (and this information is then passed to the
StepInputContext
). Because a input context constructed via
build_input_context
is not associated with a pipeline run, these fields cause an error.
Something you can do to test your IO manager is doing something like:
Copy code
partitions_def = StaticPartitionsDefinition(["a", "b", "c", "d"])

class MyIOManager(IOManager):
    def handle_output(self, context, obj):
        assert context.asset_partition_key == "b"

    def load_input(self, context):
        ...

@asset(partitions_def=partitions_def)
def my_asset(context):
    assert context.asset_partitions_def_for_output() == partitions_def

my_job = build_assets_job(
    "my_job",
    assets=[my_asset],
    resource_defs={"io_manager": IOManagerDefinition.hardcoded_io_manager(MyIOManager())},
)
result = my_job.execute_in_process(partition_key="b")
c

Chris Histe

09/28/2022, 11:38 PM
Hey Claire, thanks. Your example is similar to my code but I'm trying to unit test my io manager not manually test it. Are you suggesting I create an @asset in my unit test?
c

claire

09/28/2022, 11:39 PM
You could create an
@asset
in each unit test, or you could share the same
@asset
/ job across all unit tests and provide different partition keys, depending your needs
c

Chris Histe

09/29/2022, 3:37 PM
It’s a bit weird to me to have an asset in my unit test. I’m trying to test my IO Manager. It seems a bit overkill to update all my test to use an asset for one line of code. Isn’t there a better way to mock
_context_.has_partition_key
?
c

claire

09/29/2022, 7:18 PM
Hmm.. I realized we actually allow for mocking the partition key on `build_output_context` but not on
build_input_context
. Maybe I can put together a PR quickly to allow for the same functionality on
build_input_context
🙂
This update should make it out in next week's release (1.0.12).
c

Chris Histe

09/29/2022, 7:45 PM
Oh that’s awesome thanks 🙂
I have a similar question could it be possible to configure the run_id value in unit test through
build_init_resource_context
. I have
dagster_run_id = _context_.run_id
in my code and it’s causing some errors because it’s
None
during my unit test So as a workaround I’m doing the folllowing but it’s also a smell
Copy code
if "PYTEST_CURRENT_TEST" in os.environ:
        dagster_run_id = "unit-testing"
    else:
        dagster_run_id = context.run_id
c

claire

09/29/2022, 7:51 PM
We don't currently have any functionality to do that, but that seems like a reasonable request. I can file an issue, since it seems like there is an umbrella of configuration we'd ideally like to be able to support for unit testing.
c

Chris Histe

09/29/2022, 7:53 PM
I would love that 🙂
Would it be possible get notified when and if that code is released?
c

claire

09/29/2022, 8:03 PM
I put out an issue here, feel free to comment if you have any thoughts on other functionalities we should try to support: https://github.com/dagster-io/dagster/issues/9838
If you mean the pull request above, that will be released next Thursday 10/6. If you mean additional functionalities for unit testing, not sure--it seems like we might need to think of higher-level abstractions to make unit testing easier as a whole. Though you can subscribe to the issue to follow any progress.
c

Chris Histe

09/29/2022, 8:06 PM
Alright thanks again 🙂 I will comment that issue
It’s always a delightful experience to talk to the Dagster team!
🌈 1
72 Views