|
1 | | -from types import SimpleNamespace |
| 1 | +from contextlib import AbstractContextManager |
| 2 | +from types import SimpleNamespace, TracebackType |
2 | 3 |
|
3 | 4 | import pytest |
4 | 5 |
|
5 | 6 | from configs import dify_config |
6 | 7 | from core.helper.code_executor.code_executor import CodeLanguage |
| 8 | +from core.workflow import workflow_entry |
7 | 9 | from core.workflow.system_variables import build_system_variables, default_system_variables |
8 | 10 | from core.workflow.variable_prefixes import ( |
9 | 11 | CONVERSATION_VARIABLE_NODE_ID, |
|
14 | 16 | from graphon.file import File, FileTransferMethod, FileType |
15 | 17 | from graphon.nodes.code.code_node import CodeNode |
16 | 18 | from graphon.nodes.code.limits import CodeNodeLimits |
17 | | -from graphon.runtime import VariablePool |
| 19 | +from graphon.runtime import GraphRuntimeState, VariablePool |
18 | 20 | from graphon.variables.variables import StringVariable |
| 21 | +from tests.workflow_test_utils import build_test_graph_init_params |
19 | 22 |
|
20 | 23 |
|
21 | 24 | @pytest.fixture(autouse=True) |
@@ -52,6 +55,96 @@ def fake_head(method, url, *args, **kwargs): |
52 | 55 | class TestWorkflowEntry: |
53 | 56 | """Test WorkflowEntry class methods.""" |
54 | 57 |
|
| 58 | + def test_child_engine_enters_execution_context_while_initializing_graph(self, monkeypatch: pytest.MonkeyPatch): |
| 59 | + """Child graph node factories should run inside the parent execution context.""" |
| 60 | + |
| 61 | + class RecordingExecutionContext(AbstractContextManager[None]): |
| 62 | + entered: bool |
| 63 | + was_entered_during_graph_init: bool |
| 64 | + |
| 65 | + def __init__(self) -> None: |
| 66 | + self.entered = False |
| 67 | + self.was_entered_during_graph_init = False |
| 68 | + |
| 69 | + def __enter__(self) -> None: |
| 70 | + self.entered = True |
| 71 | + |
| 72 | + def __exit__( |
| 73 | + self, |
| 74 | + exc_type: type[BaseException] | None, |
| 75 | + exc_value: BaseException | None, |
| 76 | + traceback: TracebackType | None, |
| 77 | + ) -> bool: |
| 78 | + self.entered = False |
| 79 | + return False |
| 80 | + |
| 81 | + class StubDifyNodeFactory: |
| 82 | + graph_runtime_state: GraphRuntimeState |
| 83 | + |
| 84 | + def __init__(self, *, graph_init_params: object, graph_runtime_state: GraphRuntimeState) -> None: |
| 85 | + self.graph_init_params = graph_init_params |
| 86 | + self.graph_runtime_state = graph_runtime_state |
| 87 | + created_runtime_states.append(graph_runtime_state) |
| 88 | + |
| 89 | + class StubGraphEngine: |
| 90 | + graph_runtime_state: GraphRuntimeState |
| 91 | + layers: list[object] |
| 92 | + |
| 93 | + def __init__( |
| 94 | + self, |
| 95 | + *, |
| 96 | + workflow_id: str, |
| 97 | + graph: object, |
| 98 | + graph_runtime_state: GraphRuntimeState, |
| 99 | + command_channel: object, |
| 100 | + config: object, |
| 101 | + child_engine_builder: object, |
| 102 | + ) -> None: |
| 103 | + self.workflow_id = workflow_id |
| 104 | + self.graph = graph |
| 105 | + self.graph_runtime_state = graph_runtime_state |
| 106 | + self.command_channel = command_channel |
| 107 | + self.config = config |
| 108 | + self.child_engine_builder = child_engine_builder |
| 109 | + self.layers = [] |
| 110 | + |
| 111 | + def layer(self, layer: object) -> None: |
| 112 | + self.layers.append(layer) |
| 113 | + |
| 114 | + created_runtime_states: list[GraphRuntimeState] = [] |
| 115 | + execution_context = RecordingExecutionContext() |
| 116 | + parent_runtime_state = GraphRuntimeState( |
| 117 | + variable_pool=VariablePool(), |
| 118 | + start_at=0.0, |
| 119 | + execution_context=execution_context, |
| 120 | + ) |
| 121 | + graph_init_params = build_test_graph_init_params( |
| 122 | + graph_config={"nodes": [{"id": "root"}], "edges": []}, |
| 123 | + ) |
| 124 | + |
| 125 | + def init_graph(*, graph_config: object, node_factory: object, root_node_id: str) -> object: |
| 126 | + execution_context.was_entered_during_graph_init = execution_context.entered |
| 127 | + return {"graph_config": graph_config, "node_factory": node_factory, "root_node_id": root_node_id} |
| 128 | + |
| 129 | + monkeypatch.setattr(workflow_entry, "DifyNodeFactory", StubDifyNodeFactory) |
| 130 | + monkeypatch.setattr(workflow_entry.Graph, "init", staticmethod(init_graph)) |
| 131 | + monkeypatch.setattr(workflow_entry, "GraphEngine", StubGraphEngine) |
| 132 | + monkeypatch.setattr(workflow_entry, "LLMQuotaLayer", lambda tenant_id: ("quota", tenant_id)) |
| 133 | + |
| 134 | + engine = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-1").build_child_engine( |
| 135 | + workflow_id="workflow-1", |
| 136 | + graph_init_params=graph_init_params, |
| 137 | + parent_graph_runtime_state=parent_runtime_state, |
| 138 | + root_node_id="root", |
| 139 | + ) |
| 140 | + |
| 141 | + assert isinstance(engine, StubGraphEngine) |
| 142 | + assert execution_context.was_entered_during_graph_init is True |
| 143 | + assert execution_context.entered is False |
| 144 | + assert created_runtime_states[0].execution_context is execution_context |
| 145 | + assert engine.graph_runtime_state.execution_context is execution_context |
| 146 | + assert engine.layers == [("quota", "tenant-1")] |
| 147 | + |
55 | 148 | def test_mapping_user_inputs_to_variable_pool_with_system_variables(self): |
56 | 149 | """Test mapping system variables from user inputs to variable pool.""" |
57 | 150 | # Initialize variable pool with system variables |
|
0 commit comments