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"