https://dagster.io/ logo
Join Slack
Powered by
[Relates to unit testing and `Context`] Hello the...
# ask-community
c
[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
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
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
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
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
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
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
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
I would love that 🙂
Would it be possible get notified when and if that code is released?
c
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
Alright thanks again 🙂 I will comment that issue
It’s always a delightful experience to talk to the Dagster team!
🌈 1
95 Views
Open in Slack
PreviousNext
Powered by