From a5394fa2ce45e1c9dad1a3077bd9f09226afdd26 Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Mon, 11 Mar 2024 21:52:49 +0800 Subject: [PATCH] test: template transform --- .../template_transform_node.py | 9 +++- .../workflow/nodes/__mock/code_executor.py | 4 ++ .../workflow/nodes/test_template_transform.py | 46 +++++++++++++++++++ .../workflow/nodes/test_tool.py | 25 ---------- 4 files changed, 58 insertions(+), 26 deletions(-) create mode 100644 api/tests/integration_tests/workflow/nodes/test_template_transform.py diff --git a/api/core/workflow/nodes/template_transform/template_transform_node.py b/api/core/workflow/nodes/template_transform/template_transform_node.py index c41f5d1030..15d4b2a6e7 100644 --- a/api/core/workflow/nodes/template_transform/template_transform_node.py +++ b/api/core/workflow/nodes/template_transform/template_transform_node.py @@ -7,6 +7,7 @@ from core.workflow.nodes.base_node import BaseNode from core.workflow.nodes.template_transform.entities import TemplateTransformNodeData from models.workflow import WorkflowNodeExecutionStatus +MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH = 1000 class TemplateTransformNode(BaseNode): _node_data_cls = TemplateTransformNodeData @@ -48,7 +49,6 @@ class TemplateTransformNode(BaseNode): ) variables[variable] = value - # Run code try: result = CodeExecutor.execute_code( @@ -62,6 +62,13 @@ class TemplateTransformNode(BaseNode): status=WorkflowNodeExecutionStatus.FAILED, error=str(e) ) + + if len(result['result']) > MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH: + return NodeRunResult( + inputs=variables, + status=WorkflowNodeExecutionStatus.FAILED, + error=f"Output length exceeds {MAX_TEMPLATE_TRANSFORM_OUTPUT_LENGTH} characters" + ) return NodeRunResult( status=WorkflowNodeExecutionStatus.SUCCEEDED, diff --git a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py b/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py index a1c8eb71dc..2eb987181f 100644 --- a/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py +++ b/api/tests/integration_tests/workflow/nodes/__mock/code_executor.py @@ -15,6 +15,10 @@ class MockedCodeExecutor: return { "result": 3 } + elif language == 'jinja2': + return { + "result": "3" + } @pytest.fixture def setup_code_executor_mock(request, monkeypatch: MonkeyPatch): diff --git a/api/tests/integration_tests/workflow/nodes/test_template_transform.py b/api/tests/integration_tests/workflow/nodes/test_template_transform.py new file mode 100644 index 0000000000..4348995a05 --- /dev/null +++ b/api/tests/integration_tests/workflow/nodes/test_template_transform.py @@ -0,0 +1,46 @@ +import pytest +from core.app.entities.app_invoke_entities import InvokeFrom + +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.nodes.template_transform.template_transform_node import TemplateTransformNode +from models.workflow import WorkflowNodeExecutionStatus +from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock + +@pytest.mark.parametrize('setup_code_executor_mock', [['none']], indirect=True) +def test_execute_code(setup_code_executor_mock): + code = '''{{args2}}''' + node = TemplateTransformNode( + tenant_id='1', + app_id='1', + workflow_id='1', + user_id='1', + user_from=InvokeFrom.WEB_APP, + config={ + 'id': '1', + 'data': { + 'title': '123', + 'variables': [ + { + 'variable': 'args1', + 'value_selector': ['1', '123', 'args1'], + }, + { + 'variable': 'args2', + 'value_selector': ['1', '123', 'args2'] + } + ], + 'template': code, + } + } + ) + + # construct variable pool + pool = VariablePool(system_variables={}, user_inputs={}) + pool.append_variable(node_id='1', variable_key_list=['123', 'args1'], value=1) + pool.append_variable(node_id='1', variable_key_list=['123', 'args2'], value=3) + + # execute node + result = node.run(pool) + + assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs['output'] == '3' diff --git a/api/tests/integration_tests/workflow/nodes/test_tool.py b/api/tests/integration_tests/workflow/nodes/test_tool.py index 72e0d6f853..66139563e2 100644 --- a/api/tests/integration_tests/workflow/nodes/test_tool.py +++ b/api/tests/integration_tests/workflow/nodes/test_tool.py @@ -1,34 +1,9 @@ -import pytest from core.app.entities.app_invoke_entities import InvokeFrom from core.workflow.entities.variable_pool import VariablePool from core.workflow.nodes.tool.tool_node import ToolNode from models.workflow import WorkflowNodeExecutionStatus -""" -class ToolEntity(BaseModel): - provider_id: str - provider_type: Literal['builtin', 'api'] - provider_name: str # redundancy - tool_name: str - tool_label: str # redundancy - tool_configurations: dict[str, ToolParameterValue] - -class ToolNodeData(BaseNodeData, ToolEntity): - class ToolInput(VariableSelector): - variable_type: Literal['selector', 'static'] - value: Optional[str] - - @validator('value') - def check_value(cls, value, values, **kwargs): - if values['variable_type'] == 'static' and value is None: - raise ValueError('value is required for static variable') - return value - - tool_parameters: list[ToolInput] - -""" - def test_tool_invoke(): pool = VariablePool(system_variables={}, user_inputs={}) pool.append_variable(node_id='1', variable_key_list=['123', 'args1'], value='1+1')