def test_solid_invocation_with_config(): @solid(config_schema={"foo": str}) def solid_requires_config(context): assert context.solid_config["foo"] == "bar" return 5 # Ensure that error is raised when attempting to execute and no context is provided with pytest.raises( DagsterInvalidInvocationError, match= "Compute function of solid 'solid_requires_config' has context argument, but no " "context was provided when invoking.", ): solid_requires_config() # Ensure that alias is accounted for in error message with pytest.raises( DagsterInvalidInvocationError, match= "Compute function of solid 'aliased_solid_requires_config' has context argument, but no " "context was provided when invoking.", ): solid_requires_config.alias("aliased_solid_requires_config")() # Ensure that error is raised when we attempt to invoke with a None context with pytest.raises( DagsterInvalidInvocationError, match= 'solid "solid_requires_config" has required config schema, but no context was ' "provided.", ): solid_requires_config(None) # Ensure that error is raised when context does not have the required config. context = build_solid_context() with pytest.raises( DagsterInvalidConfigError, match="Error in config for solid", ): solid_requires_config(context) # Ensure that error is raised when attempting to execute and no context is provided, even when # configured with pytest.raises( DagsterInvalidInvocationError, match= "Compute function of solid 'configured_solid' has context argument, but no " "context was provided when invoking.", ): solid_requires_config.configured({"foo": "bar"}, name="configured_solid")() # Ensure that if you configure the solid, you can provide a none-context. result = solid_requires_config.configured({"foo": "bar"}, name="configured_solid")(None) assert result == 5 result = solid_requires_config( build_solid_context(solid_config={"foo": "bar"})) assert result == 5
def test_solid_invocation_with_cm_resource(): teardown_log = [] @resource def cm_resource(_): try: yield "foo" finally: teardown_log.append("collected") @solid(required_resource_keys={"cm_resource"}) def solid_requires_cm_resource(context): return context.resources.cm_resource # Attempt to use solid context as fxn with cm resource should fail context = build_solid_context(resources={"cm_resource": cm_resource}) with pytest.raises(DagsterInvariantViolationError): solid_requires_cm_resource(context) del context assert teardown_log == ["collected"] # Attempt to use solid context as cm with cm resource should succeed with build_solid_context( resources={"cm_resource": cm_resource}) as context: assert solid_requires_cm_resource(context) == "foo" assert teardown_log == ["collected", "collected"]
def test_solid_invocation_context_arg(): @solid def basic_solid(context): context.log.info("yay") basic_solid(None) basic_solid(build_solid_context()) basic_solid(context=None) basic_solid(context=build_solid_context())
def test_solid_invocation_dict_config(): @solid(config_schema=dict) def solid_requires_dict(context): assert context.solid_config == {"foo": "bar"} return context.solid_config assert solid_requires_dict(build_solid_context(solid_config={"foo": "bar"})) == {"foo": "bar"} @solid(config_schema=Noneable(dict)) def solid_noneable_dict(context): return context.solid_config assert solid_noneable_dict(build_solid_context()) is None assert solid_noneable_dict(None) is None
def test_solid_invocation_empty_run_config(): @solid def basic_solid(context): assert context.run_config is not None assert context.run_config == {"resources": {}} basic_solid(context=build_solid_context())
def test_context_mapping_key(): _observed = [] @solid def observe_key(context, _dep=None): _observed.append(context.get_mapping_key()) @solid(output_defs=[DynamicOutputDefinition()]) def emit(): yield DynamicOutput(1, mapping_key="key_1") yield DynamicOutput(2, mapping_key="key_2") @pipeline def test(): observe_key() emit().map(observe_key) result = execute_pipeline(test) assert result.success assert _observed == [None, "key_1", "key_2"] # test standalone doesn't throw as well _observed = [] observe_key(build_solid_context()) assert _observed == [None]
def test_solid_invocation_kitchen_sink_config(): @solid( config_schema={ "str_field": str, "int_field": int, "list_int": [int], "list_list_int": [[int]], "dict_field": {"a_string": str}, "list_dict_field": [{"an_int": int}], "selector_of_things": Selector( {"select_list_dict_field": [{"an_int": int}], "select_int": int} ), "optional_list_of_optional_string": Noneable([Noneable(str)]), } ) def kitchen_sink(context): return context.solid_config solid_config_one = { "str_field": "kjf", "int_field": 2, "list_int": [3], "list_list_int": [[1], [2, 3]], "dict_field": {"a_string": "kdjfkd"}, "list_dict_field": [{"an_int": 2}, {"an_int": 4}], "selector_of_things": {"select_int": 3}, "optional_list_of_optional_string": ["foo", None], } assert kitchen_sink(build_solid_context(solid_config=solid_config_one)) == solid_config_one
def test_solid_invocation_run_config_with_config(): @solid(config_schema={"foo": str}) def basic_solid(context): assert context.run_config assert context.run_config["solids"] == {"basic_solid": {"config": {"foo": "bar"}}} basic_solid(build_solid_context(solid_config={"foo": "bar"}))
def test_solid_invocation_default_config(): @solid(config_schema={ "foo": Field(str, is_required=False, default_value="bar") }) def solid_requires_config(context): assert context.solid_config["foo"] == "bar" return context.solid_config["foo"] assert solid_requires_config(None) == "bar" @solid(config_schema=Field(str, is_required=False, default_value="bar")) def solid_requires_config_val(context): assert context.solid_config == "bar" return context.solid_config assert solid_requires_config_val(None) == "bar" @solid(config_schema={ "foo": Field(str, is_required=False, default_value="bar"), "baz": str, }) def solid_requires_config_partial(context): assert context.solid_config["foo"] == "bar" assert context.solid_config["baz"] == "bar" return context.solid_config["foo"] + context.solid_config["baz"] assert (solid_requires_config_partial( build_solid_context(solid_config={"baz": "bar"})) == "barbar")
def test_solid_invocation_with_resources(): @solid(required_resource_keys={"foo"}) def solid_requires_resources(context): assert context.resources.foo == "bar" return context.resources.foo # Ensure that a check invariant is raise when we attempt to invoke without context with pytest.raises( DagsterInvalidInvocationError, match= "Compute function of solid 'solid_requires_resources' has context argument, but no " "context was provided when invoking.", ): solid_requires_resources() # Ensure that alias is accounted for in error message with pytest.raises( DagsterInvalidInvocationError, match= "Compute function of solid 'aliased_solid_requires_resources' has context argument, but no " "context was provided when invoking.", ): solid_requires_resources.alias("aliased_solid_requires_resources")() # Ensure that error is raised when we attempt to invoke with a None context with pytest.raises( DagsterInvalidInvocationError, match= 'solid "solid_requires_resources" has required resources, but no context was ' "provided.", ): solid_requires_resources(None) # Ensure that error is raised when we attempt to invoke with a context without the required # resource. context = build_solid_context() with pytest.raises( DagsterInvalidInvocationError, match= 'solid "solid_requires_resources" requires resource "foo", but no resource ' "with that key was found on the context.", ): solid_requires_resources(context) context = build_solid_context(resources={"foo": "bar"}) assert solid_requires_resources(context) == "bar"
def test_unconfigured(dbt_seed, conn_string, test_project_dir, dbt_config_dir): # pylint: disable=unused-argument @solid(required_resource_keys={"dbt"}) def my_dbt_solid(context): return context.resources.dbt.run(project_dir=test_project_dir, profiles_dir=dbt_config_dir) context = build_solid_context(resources={"dbt": dbt_cli_resource}) dbt_result = my_dbt_solid(context) assert len(dbt_result.result["results"]) == 4
def test_solid_invocation_no_arg(): @solid def basic_solid(): return 5 result = basic_solid() assert result == 5 with pytest.raises( DagsterInvalidInvocationError, match="Compute function of solid 'basic_solid' has no context " "argument, but context was provided when invoking.", ): basic_solid(build_solid_context()) # Ensure alias is accounted for in error message with pytest.raises( DagsterInvalidInvocationError, match= "Compute function of solid 'aliased_basic_solid' has no context " "argument, but context was provided when invoking.", ): basic_solid.alias("aliased_basic_solid")(build_solid_context()) with pytest.raises( DagsterInvalidInvocationError, match= "Too many input arguments were provided for solid 'basic_solid'. This may be " "because an argument was provided for the context parameter, but no context parameter was " "defined for the solid.", ): basic_solid(None) # Ensure alias is accounted for in error message with pytest.raises( DagsterInvalidInvocationError, match= "Too many input arguments were provided for solid 'aliased_basic_solid'. This may be " "because an argument was provided for the context parameter, but no context parameter was " "defined for the solid.", ): basic_solid.alias("aliased_basic_solid")(None)
def test_async_gen_invocation(): @solid async def aio_gen(_): await asyncio.sleep(0.01) yield Output("done") context = build_solid_context() async def get_results(): res = [] async for output in aio_gen(context): res.append(output) return res loop = asyncio.get_event_loop() output = loop.run_until_complete(get_results())[0] assert output.value == "done"
def test_solid_invocation_with_bad_resources(capsys): @resource def bad_resource(_): if 1 == 1: raise Exception("oopsy daisy") yield "foo" @solid(required_resource_keys={"my_resource"}) def solid_requires_resource(context): return context.resources.my_resource with pytest.raises( DagsterResourceFunctionError, match="Error executing resource_fn on ResourceDefinition my_resource", ): with build_solid_context(resources={"my_resource": bad_resource}) as context: assert solid_requires_resource(context) == "foo" captured = capsys.readouterr() # make sure there are no exceptions in the context destructor (__del__) assert "Exception ignored in" not in captured.err
def test_download_items(): context = build_solid_context(resources={"hn_client": hn_snapshot_client}) id_range = (SNAPSHOT_START_ID, SNAPSHOT_START_ID + 2) table = download_items(context, id_range=id_range).value assert table.shape[0] == 2
def context() -> SolidExecutionContext: return build_solid_context()
def test_database_resource(): class BasicDatabase: def execute_query(self, query): pass solid_requires_resources(build_solid_context(resources={"database": BasicDatabase()}))
def test_solid_resource_def(): context = build_solid_context(resources={"foo": my_foo_resource.configured({"my_str": "bar"})}) assert solid_requires_foo(context) == "found bar"
def get_dbt_solid_context(project_dir, profiles_dir, **kwargs): return build_solid_context( resources={ "dbt": get_dbt_resource(project_dir, profiles_dir, **kwargs) })
def test_solid_with_context(): context = build_solid_context(resources={"foo": "bar"}) assert solid_requires_foo(context) == "found bar"