Merge branch 'refs/heads/main' into feat/workflow-parallel-support

# Conflicts:
#	api/core/workflow/entities/variable_pool.py
#	api/core/workflow/nodes/iteration/iteration_node.py
#	api/core/workflow/workflow_engine_manager.py
This commit is contained in:
takatost
2024-07-31 02:25:31 +08:00
184 changed files with 3427 additions and 930 deletions

View File

@ -0,0 +1,307 @@
from uuid import uuid4
import pytest
from core.app.segments import (
ArrayFileVariable,
ArrayNumberVariable,
ArrayObjectVariable,
ArrayStringVariable,
FileVariable,
FloatVariable,
IntegerVariable,
NoneSegment,
ObjectSegment,
SecretVariable,
StringVariable,
factory,
)
def test_string_variable():
test_data = {'value_type': 'string', 'name': 'test_text', 'value': 'Hello, World!'}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, StringVariable)
def test_integer_variable():
test_data = {'value_type': 'number', 'name': 'test_int', 'value': 42}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, IntegerVariable)
def test_float_variable():
test_data = {'value_type': 'number', 'name': 'test_float', 'value': 3.14}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, FloatVariable)
def test_secret_variable():
test_data = {'value_type': 'secret', 'name': 'test_secret', 'value': 'secret_value'}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, SecretVariable)
def test_invalid_value_type():
test_data = {'value_type': 'unknown', 'name': 'test_invalid', 'value': 'value'}
with pytest.raises(ValueError):
factory.build_variable_from_mapping(test_data)
def test_build_a_blank_string():
result = factory.build_variable_from_mapping(
{
'value_type': 'string',
'name': 'blank',
'value': '',
}
)
assert isinstance(result, StringVariable)
assert result.value == ''
def test_build_a_object_variable_with_none_value():
var = factory.build_segment(
{
'key1': None,
}
)
assert isinstance(var, ObjectSegment)
assert isinstance(var.value['key1'], NoneSegment)
def test_object_variable():
mapping = {
'id': str(uuid4()),
'value_type': 'object',
'name': 'test_object',
'description': 'Description of the variable.',
'value': {
'key1': {
'id': str(uuid4()),
'value_type': 'string',
'name': 'text',
'value': 'text',
'description': 'Description of the variable.',
},
'key2': {
'id': str(uuid4()),
'value_type': 'number',
'name': 'number',
'value': 1,
'description': 'Description of the variable.',
},
},
}
variable = factory.build_variable_from_mapping(mapping)
assert isinstance(variable, ObjectSegment)
assert isinstance(variable.value['key1'], StringVariable)
assert isinstance(variable.value['key2'], IntegerVariable)
def test_array_string_variable():
mapping = {
'id': str(uuid4()),
'value_type': 'array[string]',
'name': 'test_array',
'description': 'Description of the variable.',
'value': [
{
'id': str(uuid4()),
'value_type': 'string',
'name': 'text',
'value': 'text',
'description': 'Description of the variable.',
},
{
'id': str(uuid4()),
'value_type': 'string',
'name': 'text',
'value': 'text',
'description': 'Description of the variable.',
},
],
}
variable = factory.build_variable_from_mapping(mapping)
assert isinstance(variable, ArrayStringVariable)
assert isinstance(variable.value[0], StringVariable)
assert isinstance(variable.value[1], StringVariable)
def test_array_number_variable():
mapping = {
'id': str(uuid4()),
'value_type': 'array[number]',
'name': 'test_array',
'description': 'Description of the variable.',
'value': [
{
'id': str(uuid4()),
'value_type': 'number',
'name': 'number',
'value': 1,
'description': 'Description of the variable.',
},
{
'id': str(uuid4()),
'value_type': 'number',
'name': 'number',
'value': 2.0,
'description': 'Description of the variable.',
},
],
}
variable = factory.build_variable_from_mapping(mapping)
assert isinstance(variable, ArrayNumberVariable)
assert isinstance(variable.value[0], IntegerVariable)
assert isinstance(variable.value[1], FloatVariable)
def test_array_object_variable():
mapping = {
'id': str(uuid4()),
'value_type': 'array[object]',
'name': 'test_array',
'description': 'Description of the variable.',
'value': [
{
'id': str(uuid4()),
'value_type': 'object',
'name': 'object',
'description': 'Description of the variable.',
'value': {
'key1': {
'id': str(uuid4()),
'value_type': 'string',
'name': 'text',
'value': 'text',
'description': 'Description of the variable.',
},
'key2': {
'id': str(uuid4()),
'value_type': 'number',
'name': 'number',
'value': 1,
'description': 'Description of the variable.',
},
},
},
{
'id': str(uuid4()),
'value_type': 'object',
'name': 'object',
'description': 'Description of the variable.',
'value': {
'key1': {
'id': str(uuid4()),
'value_type': 'string',
'name': 'text',
'value': 'text',
'description': 'Description of the variable.',
},
'key2': {
'id': str(uuid4()),
'value_type': 'number',
'name': 'number',
'value': 1,
'description': 'Description of the variable.',
},
},
},
],
}
variable = factory.build_variable_from_mapping(mapping)
assert isinstance(variable, ArrayObjectVariable)
assert isinstance(variable.value[0], ObjectSegment)
assert isinstance(variable.value[1], ObjectSegment)
assert isinstance(variable.value[0].value['key1'], StringVariable)
assert isinstance(variable.value[0].value['key2'], IntegerVariable)
assert isinstance(variable.value[1].value['key1'], StringVariable)
assert isinstance(variable.value[1].value['key2'], IntegerVariable)
def test_file_variable():
mapping = {
'id': str(uuid4()),
'value_type': 'file',
'name': 'test_file',
'description': 'Description of the variable.',
'value': {
'id': str(uuid4()),
'tenant_id': 'tenant_id',
'type': 'image',
'transfer_method': 'local_file',
'url': 'url',
'related_id': 'related_id',
'extra_config': {
'image_config': {
'width': 100,
'height': 100,
},
},
'filename': 'filename',
'extension': 'extension',
'mime_type': 'mime_type',
},
}
variable = factory.build_variable_from_mapping(mapping)
assert isinstance(variable, FileVariable)
def test_array_file_variable():
mapping = {
'id': str(uuid4()),
'value_type': 'array[file]',
'name': 'test_array_file',
'description': 'Description of the variable.',
'value': [
{
'id': str(uuid4()),
'name': 'file',
'value_type': 'file',
'value': {
'id': str(uuid4()),
'tenant_id': 'tenant_id',
'type': 'image',
'transfer_method': 'local_file',
'url': 'url',
'related_id': 'related_id',
'extra_config': {
'image_config': {
'width': 100,
'height': 100,
},
},
'filename': 'filename',
'extension': 'extension',
'mime_type': 'mime_type',
},
},
{
'id': str(uuid4()),
'name': 'file',
'value_type': 'file',
'value': {
'id': str(uuid4()),
'tenant_id': 'tenant_id',
'type': 'image',
'transfer_method': 'local_file',
'url': 'url',
'related_id': 'related_id',
'extra_config': {
'image_config': {
'width': 100,
'height': 100,
},
},
'filename': 'filename',
'extension': 'extension',
'mime_type': 'mime_type',
},
},
],
}
variable = factory.build_variable_from_mapping(mapping)
assert isinstance(variable, ArrayFileVariable)
assert isinstance(variable.value[0], FileVariable)
assert isinstance(variable.value[1], FileVariable)

View File

@ -1,4 +1,4 @@
from core.app.segments import SecretVariable, parser
from core.app.segments import SecretVariable, StringSegment, parser
from core.helper import encrypter
from core.workflow.entities.node_entities import SystemVariable
from core.workflow.entities.variable_pool import VariablePool
@ -51,3 +51,4 @@ def test_convert_variable_to_segment_group():
segments_group = parser.convert_template(template=template, variable_pool=variable_pool)
assert segments_group.text == 'fake-user-id'
assert segments_group.log == 'fake-user-id'
assert segments_group.value == [StringSegment(value='fake-user-id')]

View File

@ -2,48 +2,16 @@ import pytest
from pydantic import ValidationError
from core.app.segments import (
ArrayVariable,
ArrayAnyVariable,
FloatVariable,
IntegerVariable,
NoneVariable,
ObjectVariable,
SecretVariable,
SegmentType,
StringVariable,
factory,
)
def test_string_variable():
test_data = {'value_type': 'string', 'name': 'test_text', 'value': 'Hello, World!'}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, StringVariable)
def test_integer_variable():
test_data = {'value_type': 'number', 'name': 'test_int', 'value': 42}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, IntegerVariable)
def test_float_variable():
test_data = {'value_type': 'number', 'name': 'test_float', 'value': 3.14}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, FloatVariable)
def test_secret_variable():
test_data = {'value_type': 'secret', 'name': 'test_secret', 'value': 'secret_value'}
result = factory.build_variable_from_mapping(test_data)
assert isinstance(result, SecretVariable)
def test_invalid_value_type():
test_data = {'value_type': 'unknown', 'name': 'test_invalid', 'value': 'value'}
with pytest.raises(ValueError):
factory.build_variable_from_mapping(test_data)
def test_frozen_variables():
var = StringVariable(name='text', value='text')
with pytest.raises(ValidationError):
@ -64,34 +32,22 @@ def test_frozen_variables():
def test_variable_value_type_immutable():
with pytest.raises(ValidationError):
StringVariable(value_type=SegmentType.ARRAY, name='text', value='text')
StringVariable(value_type=SegmentType.ARRAY_ANY, name='text', value='text')
with pytest.raises(ValidationError):
StringVariable.model_validate({'value_type': 'not text', 'name': 'text', 'value': 'text'})
var = IntegerVariable(name='integer', value=42)
with pytest.raises(ValidationError):
IntegerVariable(value_type=SegmentType.ARRAY, name=var.name, value=var.value)
IntegerVariable(value_type=SegmentType.ARRAY_ANY, name=var.name, value=var.value)
var = FloatVariable(name='float', value=3.14)
with pytest.raises(ValidationError):
FloatVariable(value_type=SegmentType.ARRAY, name=var.name, value=var.value)
FloatVariable(value_type=SegmentType.ARRAY_ANY, name=var.name, value=var.value)
var = SecretVariable(name='secret', value='secret_value')
with pytest.raises(ValidationError):
SecretVariable(value_type=SegmentType.ARRAY, name=var.name, value=var.value)
def test_build_a_blank_string():
result = factory.build_variable_from_mapping(
{
'value_type': 'string',
'name': 'blank',
'value': '',
}
)
assert isinstance(result, StringVariable)
assert result.value == ''
SecretVariable(value_type=SegmentType.ARRAY_ANY, name=var.name, value=var.value)
def test_object_variable_to_object():
@ -104,7 +60,7 @@ def test_object_variable_to_object():
'key2': StringVariable(name='key2', value='value2'),
},
),
'key2': ArrayVariable(
'key2': ArrayAnyVariable(
name='array',
value=[
StringVariable(name='key5_1', value='value5_1'),
@ -136,13 +92,3 @@ def test_variable_to_object():
assert var.to_object() == 3.14
var = SecretVariable(name='secret', value='secret_value')
assert var.to_object() == 'secret_value'
def test_build_a_object_variable_with_none_value():
var = factory.build_anonymous_variable(
{
'key1': None,
}
)
assert isinstance(var, ObjectVariable)
assert isinstance(var.value['key1'], NoneVariable)

View File

@ -0,0 +1,52 @@
import random
from unittest.mock import MagicMock, patch
from core.helper.ssrf_proxy import SSRF_DEFAULT_MAX_RETRIES, STATUS_FORCELIST, make_request
@patch('httpx.request')
def test_successful_request(mock_request):
mock_response = MagicMock()
mock_response.status_code = 200
mock_request.return_value = mock_response
response = make_request('GET', 'http://example.com')
assert response.status_code == 200
@patch('httpx.request')
def test_retry_exceed_max_retries(mock_request):
mock_response = MagicMock()
mock_response.status_code = 500
side_effects = [mock_response] * SSRF_DEFAULT_MAX_RETRIES
mock_request.side_effect = side_effects
try:
make_request('GET', 'http://example.com', max_retries=SSRF_DEFAULT_MAX_RETRIES - 1)
raise AssertionError("Expected Exception not raised")
except Exception as e:
assert str(e) == f"Reached maximum retries ({SSRF_DEFAULT_MAX_RETRIES - 1}) for URL http://example.com"
@patch('httpx.request')
def test_retry_logic_success(mock_request):
side_effects = []
for _ in range(SSRF_DEFAULT_MAX_RETRIES):
status_code = random.choice(STATUS_FORCELIST)
mock_response = MagicMock()
mock_response.status_code = status_code
side_effects.append(mock_response)
mock_response_200 = MagicMock()
mock_response_200.status_code = 200
side_effects.append(mock_response_200)
mock_request.side_effect = side_effects
response = make_request('GET', 'http://example.com', max_retries=SSRF_DEFAULT_MAX_RETRIES)
assert response.status_code == 200
assert mock_request.call_count == SSRF_DEFAULT_MAX_RETRIES + 1
assert mock_request.call_args_list[0][1].get('method') == 'GET'

View File

@ -21,6 +21,20 @@ def prepare_example_positions_yaml(tmp_path, monkeypatch) -> str:
return str(tmp_path)
@pytest.fixture
def prepare_empty_commented_positions_yaml(tmp_path, monkeypatch) -> str:
monkeypatch.chdir(tmp_path)
tmp_path.joinpath("example_positions_all_commented.yaml").write_text(dedent(
"""\
# - commented1
# - commented2
-
-
"""))
return str(tmp_path)
def test_position_helper(prepare_example_positions_yaml):
position_map = get_position_map(
folder_path=prepare_example_positions_yaml,
@ -32,3 +46,10 @@ def test_position_helper(prepare_example_positions_yaml):
'third': 2,
'forth': 3,
}
def test_position_helper_with_all_commented(prepare_empty_commented_positions_yaml):
position_map = get_position_map(
folder_path=prepare_empty_commented_positions_yaml,
file_name='example_positions_all_commented.yaml')
assert position_map == {}

View File

@ -53,6 +53,9 @@ def test_load_yaml_non_existing_file():
assert load_yaml_file(file_path=NON_EXISTING_YAML_FILE) == {}
assert load_yaml_file(file_path='') == {}
with pytest.raises(FileNotFoundError):
load_yaml_file(file_path=NON_EXISTING_YAML_FILE, ignore_error=False)
def test_load_valid_yaml_file(prepare_example_yaml_file):
yaml_data = load_yaml_file(file_path=prepare_example_yaml_file)
@ -68,7 +71,7 @@ def test_load_valid_yaml_file(prepare_example_yaml_file):
def test_load_invalid_yaml_file(prepare_invalid_yaml_file):
# yaml syntax error
with pytest.raises(YAMLError):
load_yaml_file(file_path=prepare_invalid_yaml_file)
load_yaml_file(file_path=prepare_invalid_yaml_file, ignore_error=False)
# ignore error
assert load_yaml_file(file_path=prepare_invalid_yaml_file, ignore_error=True) == {}
assert load_yaml_file(file_path=prepare_invalid_yaml_file) == {}