Underground-Digital/Workflow-Engine

1import time2import uuid3from os import getenv4from typing import cast5 6import pytest7 8from core.app.entities.app_invoke_entities import InvokeFrom9from core.workflow.entities.node_entities import NodeRunResult10from core.workflow.entities.variable_pool import VariablePool11from core.workflow.enums import SystemVariableKey12from core.workflow.graph_engine.entities.graph import Graph13from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams14from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState15from core.workflow.nodes.code.code_node import CodeNode16from core.workflow.nodes.code.entities import CodeNodeData17from models.enums import UserFrom18from models.workflow import WorkflowNodeExecutionStatus, WorkflowType19from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock20 21CODE_MAX_STRING_LENGTH = int(getenv("CODE_MAX_STRING_LENGTH", "10000"))22 23 24def init_code_node(code_config: dict):25    graph_config = {26        "edges": [27            {28                "id": "start-source-code-target",29                "source": "start",30                "target": "code",31            },32        ],33        "nodes": [{"data": {"type": "start"}, "id": "start"}, code_config],34    }35 36    graph = Graph.init(graph_config=graph_config)37 38    init_params = GraphInitParams(39        tenant_id="1",40        app_id="1",41        workflow_type=WorkflowType.WORKFLOW,42        workflow_id="1",43        graph_config=graph_config,44        user_id="1",45        user_from=UserFrom.ACCOUNT,46        invoke_from=InvokeFrom.DEBUGGER,47        call_depth=0,48    )49 50    # construct variable pool51    variable_pool = VariablePool(52        system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"},53        user_inputs={},54        environment_variables=[],55        conversation_variables=[],56    )57    variable_pool.add(["code", "123", "args1"], 1)58    variable_pool.add(["code", "123", "args2"], 2)59 60    node = CodeNode(61        id=str(uuid.uuid4()),62        graph_init_params=init_params,63        graph=graph,64        graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()),65        config=code_config,66    )67 68    return node69 70 71@pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True)72def test_execute_code(setup_code_executor_mock):73    code = """74    def main(args1: int, args2: int) -> dict:75        return {76            "result": args1 + args2,77        }78    """79    # trim first 4 spaces at the beginning of each line80    code = "\n".join([line[4:] for line in code.split("\n")])81 82    code_config = {83        "id": "code",84        "data": {85            "outputs": {86                "result": {87                    "type": "number",88                },89            },90            "title": "123",91            "variables": [92                {93                    "variable": "args1",94                    "value_selector": ["1", "123", "args1"],95                },96                {"variable": "args2", "value_selector": ["1", "123", "args2"]},97            ],98            "answer": "123",99            "code_language": "python3",100            "code": code,101        },102    }103 104    node = init_code_node(code_config)105    node.graph_runtime_state.variable_pool.add(["1", "123", "args1"], 1)106    node.graph_runtime_state.variable_pool.add(["1", "123", "args2"], 2)107 108    # execute node109    result = node._run()110    assert isinstance(result, NodeRunResult)111    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED112    assert result.outputs is not None113    assert result.outputs["result"] == 3114    assert result.error is None115 116 117@pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True)118def test_execute_code_output_validator(setup_code_executor_mock):119    code = """120    def main(args1: int, args2: int) -> dict:121        return {122            "result": args1 + args2,123        }124    """125    # trim first 4 spaces at the beginning of each line126    code = "\n".join([line[4:] for line in code.split("\n")])127 128    code_config = {129        "id": "code",130        "data": {131            "outputs": {132                "result": {133                    "type": "string",134                },135            },136            "title": "123",137            "variables": [138                {139                    "variable": "args1",140                    "value_selector": ["1", "123", "args1"],141                },142                {"variable": "args2", "value_selector": ["1", "123", "args2"]},143            ],144            "answer": "123",145            "code_language": "python3",146            "code": code,147        },148    }149 150    node = init_code_node(code_config)151    node.graph_runtime_state.variable_pool.add(["1", "123", "args1"], 1)152    node.graph_runtime_state.variable_pool.add(["1", "123", "args2"], 2)153 154    # execute node155    result = node._run()156    assert isinstance(result, NodeRunResult)157    assert result.status == WorkflowNodeExecutionStatus.FAILED158    assert result.error == "Output variable `result` must be a string"159 160 161def test_execute_code_output_validator_depth():162    code = """163    def main(args1: int, args2: int) -> dict:164        return {165            "result": {166                "result": args1 + args2,167            }168        }169    """170    # trim first 4 spaces at the beginning of each line171    code = "\n".join([line[4:] for line in code.split("\n")])172 173    code_config = {174        "id": "code",175        "data": {176            "outputs": {177                "string_validator": {178                    "type": "string",179                },180                "number_validator": {181                    "type": "number",182                },183                "number_array_validator": {184                    "type": "array[number]",185                },186                "string_array_validator": {187                    "type": "array[string]",188                },189                "object_validator": {190                    "type": "object",191                    "children": {192                        "result": {193                            "type": "number",194                        },195                        "depth": {196                            "type": "object",197                            "children": {198                                "depth": {199                                    "type": "object",200                                    "children": {201                                        "depth": {202                                            "type": "number",203                                        }204                                    },205                                }206                            },207                        },208                    },209                },210            },211            "title": "123",212            "variables": [213                {214                    "variable": "args1",215                    "value_selector": ["1", "123", "args1"],216                },217                {"variable": "args2", "value_selector": ["1", "123", "args2"]},218            ],219            "answer": "123",220            "code_language": "python3",221            "code": code,222        },223    }224 225    node = init_code_node(code_config)226 227    # construct result228    result = {229        "number_validator": 1,230        "string_validator": "1",231        "number_array_validator": [1, 2, 3, 3.333],232        "string_array_validator": ["1", "2", "3"],233        "object_validator": {"result": 1, "depth": {"depth": {"depth": 1}}},234    }235 236    node.node_data = cast(CodeNodeData, node.node_data)237 238    # validate239    node._transform_result(result, node.node_data.outputs)240 241    # construct result242    result = {243        "number_validator": "1",244        "string_validator": 1,245        "number_array_validator": ["1", "2", "3", "3.333"],246        "string_array_validator": [1, 2, 3],247        "object_validator": {"result": "1", "depth": {"depth": {"depth": "1"}}},248    }249 250    # validate251    with pytest.raises(ValueError):252        node._transform_result(result, node.node_data.outputs)253 254    # construct result255    result = {256        "number_validator": 1,257        "string_validator": (CODE_MAX_STRING_LENGTH + 1) * "1",258        "number_array_validator": [1, 2, 3, 3.333],259        "string_array_validator": ["1", "2", "3"],260        "object_validator": {"result": 1, "depth": {"depth": {"depth": 1}}},261    }262 263    # validate264    with pytest.raises(ValueError):265        node._transform_result(result, node.node_data.outputs)266 267    # construct result268    result = {269        "number_validator": 1,270        "string_validator": "1",271        "number_array_validator": [1, 2, 3, 3.333] * 2000,272        "string_array_validator": ["1", "2", "3"],273        "object_validator": {"result": 1, "depth": {"depth": {"depth": 1}}},274    }275 276    # validate277    with pytest.raises(ValueError):278        node._transform_result(result, node.node_data.outputs)279 280 281def test_execute_code_output_object_list():282    code = """283    def main(args1: int, args2: int) -> dict:284        return {285            "result": {286                "result": args1 + args2,287            }288        }289    """290    # trim first 4 spaces at the beginning of each line291    code = "\n".join([line[4:] for line in code.split("\n")])292 293    code_config = {294        "id": "code",295        "data": {296            "outputs": {297                "object_list": {298                    "type": "array[object]",299                },300            },301            "title": "123",302            "variables": [303                {304                    "variable": "args1",305                    "value_selector": ["1", "123", "args1"],306                },307                {"variable": "args2", "value_selector": ["1", "123", "args2"]},308            ],309            "answer": "123",310            "code_language": "python3",311            "code": code,312        },313    }314 315    node = init_code_node(code_config)316 317    # construct result318    result = {319        "object_list": [320            {321                "result": 1,322            },323            {324                "result": 2,325            },326            {327                "result": [1, 2, 3],328            },329        ]330    }331 332    node.node_data = cast(CodeNodeData, node.node_data)333 334    # validate335    node._transform_result(result, node.node_data.outputs)336 337    # construct result338    result = {339        "object_list": [340            {341                "result": 1,342            },343            {344                "result": 2,345            },346            {347                "result": [1, 2, 3],348            },349            1,350        ]351    }352 353    # validate354    with pytest.raises(ValueError):355        node._transform_result(result, node.node_data.outputs)356