def test_execution_context(default_model_storage: ModelStorage): context = ExecutionContext(GraphSchema({}), "some_id") node = GraphNode( node_name="execution_context_aware", component_class=ExecutionContextAware, constructor_name="create", component_config={}, fn_name="get_execution_context", inputs={}, eager=False, model_storage=default_model_storage, resource=None, execution_context=context, ) context.model_id = "a_new_id" result = node()[1] assert result.model_id == "some_id" assert result.node_name == "execution_context_aware"
def test_execution_context(default_model_storage: ModelStorage): graph_schema = GraphSchema({ "execution_context_aware": SchemaNode( needs={}, uses=ExecutionContextAware, fn="get_execution_context", constructor_name="create", config={}, is_target=True, ) }) context = ExecutionContext(graph_schema=graph_schema, model_id="some_id") runner = DaskGraphRunner( graph_schema=graph_schema, model_storage=default_model_storage, execution_context=context, ) context.model_id = "a_new_id" result = runner.run()["execution_context_aware"] assert result.model_id == "some_id" assert result.node_name == "execution_context_aware"