Spaces:
Running
Running
| from os import getenv | |
| import pytest | |
| from core.app.entities.app_invoke_entities import InvokeFrom | |
| from core.workflow.entities.variable_pool import VariablePool | |
| from core.workflow.nodes.code.code_node import CodeNode | |
| from models.workflow import WorkflowNodeExecutionStatus | |
| from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock | |
| CODE_MAX_STRING_LENGTH = int(getenv('CODE_MAX_STRING_LENGTH', '10000')) | |
| def test_execute_code(setup_code_executor_mock): | |
| code = ''' | |
| def main(args1: int, args2: int) -> dict: | |
| return { | |
| "result": args1 + args2, | |
| } | |
| ''' | |
| # trim first 4 spaces at the beginning of each line | |
| code = '\n'.join([line[4:] for line in code.split('\n')]) | |
| node = CodeNode( | |
| tenant_id='1', | |
| app_id='1', | |
| workflow_id='1', | |
| user_id='1', | |
| user_from=InvokeFrom.WEB_APP, | |
| config={ | |
| 'id': '1', | |
| 'data': { | |
| 'outputs': { | |
| 'result': { | |
| 'type': 'number', | |
| }, | |
| }, | |
| 'title': '123', | |
| 'variables': [ | |
| { | |
| 'variable': 'args1', | |
| 'value_selector': ['1', '123', 'args1'], | |
| }, | |
| { | |
| 'variable': 'args2', | |
| 'value_selector': ['1', '123', 'args2'] | |
| } | |
| ], | |
| 'answer': '123', | |
| 'code_language': 'python3', | |
| 'code': 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=2) | |
| # execute node | |
| result = node.run(pool) | |
| assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED | |
| assert result.outputs['result'] == 3 | |
| assert result.error is None | |
| def test_execute_code_output_validator(setup_code_executor_mock): | |
| code = ''' | |
| def main(args1: int, args2: int) -> dict: | |
| return { | |
| "result": args1 + args2, | |
| } | |
| ''' | |
| # trim first 4 spaces at the beginning of each line | |
| code = '\n'.join([line[4:] for line in code.split('\n')]) | |
| node = CodeNode( | |
| tenant_id='1', | |
| app_id='1', | |
| workflow_id='1', | |
| user_id='1', | |
| user_from=InvokeFrom.WEB_APP, | |
| config={ | |
| 'id': '1', | |
| 'data': { | |
| "outputs": { | |
| "result": { | |
| "type": "string", | |
| }, | |
| }, | |
| 'title': '123', | |
| 'variables': [ | |
| { | |
| 'variable': 'args1', | |
| 'value_selector': ['1', '123', 'args1'], | |
| }, | |
| { | |
| 'variable': 'args2', | |
| 'value_selector': ['1', '123', 'args2'] | |
| } | |
| ], | |
| 'answer': '123', | |
| 'code_language': 'python3', | |
| 'code': 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=2) | |
| # execute node | |
| result = node.run(pool) | |
| assert result.status == WorkflowNodeExecutionStatus.FAILED | |
| assert result.error == 'Output variable `result` must be a string' | |
| def test_execute_code_output_validator_depth(): | |
| code = ''' | |
| def main(args1: int, args2: int) -> dict: | |
| return { | |
| "result": { | |
| "result": args1 + args2, | |
| } | |
| } | |
| ''' | |
| # trim first 4 spaces at the beginning of each line | |
| code = '\n'.join([line[4:] for line in code.split('\n')]) | |
| node = CodeNode( | |
| tenant_id='1', | |
| app_id='1', | |
| workflow_id='1', | |
| user_id='1', | |
| user_from=InvokeFrom.WEB_APP, | |
| config={ | |
| 'id': '1', | |
| 'data': { | |
| "outputs": { | |
| "string_validator": { | |
| "type": "string", | |
| }, | |
| "number_validator": { | |
| "type": "number", | |
| }, | |
| "number_array_validator": { | |
| "type": "array[number]", | |
| }, | |
| "string_array_validator": { | |
| "type": "array[string]", | |
| }, | |
| "object_validator": { | |
| "type": "object", | |
| "children": { | |
| "result": { | |
| "type": "number", | |
| }, | |
| "depth": { | |
| "type": "object", | |
| "children": { | |
| "depth": { | |
| "type": "object", | |
| "children": { | |
| "depth": { | |
| "type": "number", | |
| } | |
| } | |
| } | |
| } | |
| } | |
| } | |
| }, | |
| }, | |
| 'title': '123', | |
| 'variables': [ | |
| { | |
| 'variable': 'args1', | |
| 'value_selector': ['1', '123', 'args1'], | |
| }, | |
| { | |
| 'variable': 'args2', | |
| 'value_selector': ['1', '123', 'args2'] | |
| } | |
| ], | |
| 'answer': '123', | |
| 'code_language': 'python3', | |
| 'code': code | |
| } | |
| } | |
| ) | |
| # construct result | |
| result = { | |
| "number_validator": 1, | |
| "string_validator": "1", | |
| "number_array_validator": [1, 2, 3, 3.333], | |
| "string_array_validator": ["1", "2", "3"], | |
| "object_validator": { | |
| "result": 1, | |
| "depth": { | |
| "depth": { | |
| "depth": 1 | |
| } | |
| } | |
| } | |
| } | |
| # validate | |
| node._transform_result(result, node.node_data.outputs) | |
| # construct result | |
| result = { | |
| "number_validator": "1", | |
| "string_validator": 1, | |
| "number_array_validator": ["1", "2", "3", "3.333"], | |
| "string_array_validator": [1, 2, 3], | |
| "object_validator": { | |
| "result": "1", | |
| "depth": { | |
| "depth": { | |
| "depth": "1" | |
| } | |
| } | |
| } | |
| } | |
| # validate | |
| with pytest.raises(ValueError): | |
| node._transform_result(result, node.node_data.outputs) | |
| # construct result | |
| result = { | |
| "number_validator": 1, | |
| "string_validator": (CODE_MAX_STRING_LENGTH + 1) * "1", | |
| "number_array_validator": [1, 2, 3, 3.333], | |
| "string_array_validator": ["1", "2", "3"], | |
| "object_validator": { | |
| "result": 1, | |
| "depth": { | |
| "depth": { | |
| "depth": 1 | |
| } | |
| } | |
| } | |
| } | |
| # validate | |
| with pytest.raises(ValueError): | |
| node._transform_result(result, node.node_data.outputs) | |
| # construct result | |
| result = { | |
| "number_validator": 1, | |
| "string_validator": "1", | |
| "number_array_validator": [1, 2, 3, 3.333] * 2000, | |
| "string_array_validator": ["1", "2", "3"], | |
| "object_validator": { | |
| "result": 1, | |
| "depth": { | |
| "depth": { | |
| "depth": 1 | |
| } | |
| } | |
| } | |
| } | |
| # validate | |
| with pytest.raises(ValueError): | |
| node._transform_result(result, node.node_data.outputs) | |
| def test_execute_code_output_object_list(): | |
| code = ''' | |
| def main(args1: int, args2: int) -> dict: | |
| return { | |
| "result": { | |
| "result": args1 + args2, | |
| } | |
| } | |
| ''' | |
| # trim first 4 spaces at the beginning of each line | |
| code = '\n'.join([line[4:] for line in code.split('\n')]) | |
| node = CodeNode( | |
| tenant_id='1', | |
| app_id='1', | |
| workflow_id='1', | |
| user_id='1', | |
| user_from=InvokeFrom.WEB_APP, | |
| config={ | |
| 'id': '1', | |
| 'data': { | |
| "outputs": { | |
| "object_list": { | |
| "type": "array[object]", | |
| }, | |
| }, | |
| 'title': '123', | |
| 'variables': [ | |
| { | |
| 'variable': 'args1', | |
| 'value_selector': ['1', '123', 'args1'], | |
| }, | |
| { | |
| 'variable': 'args2', | |
| 'value_selector': ['1', '123', 'args2'] | |
| } | |
| ], | |
| 'answer': '123', | |
| 'code_language': 'python3', | |
| 'code': code | |
| } | |
| } | |
| ) | |
| # construct result | |
| result = { | |
| "object_list": [{ | |
| "result": 1, | |
| }, { | |
| "result": 2, | |
| }, { | |
| "result": [1, 2, 3], | |
| }] | |
| } | |
| # validate | |
| node._transform_result(result, node.node_data.outputs) | |
| # construct result | |
| result = { | |
| "object_list": [{ | |
| "result": 1, | |
| }, { | |
| "result": 2, | |
| }, { | |
| "result": [1, 2, 3], | |
| }, 1] | |
| } | |
| # validate | |
| with pytest.raises(ValueError): | |
| node._transform_result(result, node.node_data.outputs) | |