dify
This commit is contained in:
299
dify/api/tests/unit_tests/models/test_workflow.py
Normal file
299
dify/api/tests/unit_tests/models/test_workflow.py
Normal file
@@ -0,0 +1,299 @@
|
||||
import dataclasses
|
||||
import json
|
||||
from unittest import mock
|
||||
from uuid import uuid4
|
||||
|
||||
from constants import HIDDEN_VALUE
|
||||
from core.file.enums import FileTransferMethod, FileType
|
||||
from core.file.models import File
|
||||
from core.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable
|
||||
from core.variables.segments import IntegerSegment, Segment
|
||||
from factories.variable_factory import build_segment
|
||||
from models.workflow import Workflow, WorkflowDraftVariable, WorkflowNodeExecutionModel, is_system_variable_editable
|
||||
|
||||
|
||||
def test_environment_variables():
|
||||
# tenant_id context variable removed - using current_user.current_tenant_id directly
|
||||
|
||||
# Create a Workflow instance
|
||||
workflow = Workflow(
|
||||
tenant_id="tenant_id",
|
||||
app_id="app_id",
|
||||
type="workflow",
|
||||
version="draft",
|
||||
graph="{}",
|
||||
features="{}",
|
||||
created_by="account_id",
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
)
|
||||
|
||||
# Create some EnvironmentVariable instances
|
||||
variable1 = StringVariable.model_validate(
|
||||
{"name": "var1", "value": "value1", "id": str(uuid4()), "selector": ["env", "var1"]}
|
||||
)
|
||||
variable2 = IntegerVariable.model_validate(
|
||||
{"name": "var2", "value": 123, "id": str(uuid4()), "selector": ["env", "var2"]}
|
||||
)
|
||||
variable3 = SecretVariable.model_validate(
|
||||
{"name": "var3", "value": "secret", "id": str(uuid4()), "selector": ["env", "var3"]}
|
||||
)
|
||||
variable4 = FloatVariable.model_validate(
|
||||
{"name": "var4", "value": 3.14, "id": str(uuid4()), "selector": ["env", "var4"]}
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
|
||||
mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"),
|
||||
):
|
||||
# Set the environment_variables property of the Workflow instance
|
||||
variables = [variable1, variable2, variable3, variable4]
|
||||
workflow.environment_variables = variables
|
||||
|
||||
# Get the environment_variables property and assert its value
|
||||
assert workflow.environment_variables == variables
|
||||
|
||||
|
||||
def test_update_environment_variables():
|
||||
# tenant_id context variable removed - using current_user.current_tenant_id directly
|
||||
|
||||
# Create a Workflow instance
|
||||
workflow = Workflow(
|
||||
tenant_id="tenant_id",
|
||||
app_id="app_id",
|
||||
type="workflow",
|
||||
version="draft",
|
||||
graph="{}",
|
||||
features="{}",
|
||||
created_by="account_id",
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
)
|
||||
|
||||
# Create some EnvironmentVariable instances
|
||||
variable1 = StringVariable.model_validate(
|
||||
{"name": "var1", "value": "value1", "id": str(uuid4()), "selector": ["env", "var1"]}
|
||||
)
|
||||
variable2 = IntegerVariable.model_validate(
|
||||
{"name": "var2", "value": 123, "id": str(uuid4()), "selector": ["env", "var2"]}
|
||||
)
|
||||
variable3 = SecretVariable.model_validate(
|
||||
{"name": "var3", "value": "secret", "id": str(uuid4()), "selector": ["env", "var3"]}
|
||||
)
|
||||
variable4 = FloatVariable.model_validate(
|
||||
{"name": "var4", "value": 3.14, "id": str(uuid4()), "selector": ["env", "var4"]}
|
||||
)
|
||||
|
||||
with (
|
||||
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
|
||||
mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"),
|
||||
):
|
||||
variables = [variable1, variable2, variable3, variable4]
|
||||
|
||||
# Set the environment_variables property of the Workflow instance
|
||||
workflow.environment_variables = variables
|
||||
assert workflow.environment_variables == [variable1, variable2, variable3, variable4]
|
||||
|
||||
# Update the name of variable3 and keep the value as it is
|
||||
variables[2] = variable3.model_copy(
|
||||
update={
|
||||
"name": "new name",
|
||||
"value": HIDDEN_VALUE,
|
||||
}
|
||||
)
|
||||
|
||||
workflow.environment_variables = variables
|
||||
assert workflow.environment_variables[2].name == "new name"
|
||||
assert workflow.environment_variables[2].value == variable3.value
|
||||
|
||||
|
||||
def test_to_dict():
|
||||
# tenant_id context variable removed - using current_user.current_tenant_id directly
|
||||
|
||||
# Create a Workflow instance
|
||||
workflow = Workflow(
|
||||
tenant_id="tenant_id",
|
||||
app_id="app_id",
|
||||
type="workflow",
|
||||
version="draft",
|
||||
graph="{}",
|
||||
features="{}",
|
||||
created_by="account_id",
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
)
|
||||
|
||||
# Create some EnvironmentVariable instances
|
||||
|
||||
with (
|
||||
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
|
||||
mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"),
|
||||
):
|
||||
# Set the environment_variables property of the Workflow instance
|
||||
workflow.environment_variables = [
|
||||
SecretVariable.model_validate({"name": "secret", "value": "secret", "id": str(uuid4())}),
|
||||
StringVariable.model_validate({"name": "text", "value": "text", "id": str(uuid4())}),
|
||||
]
|
||||
|
||||
workflow_dict = workflow.to_dict()
|
||||
assert workflow_dict["environment_variables"][0]["value"] == ""
|
||||
assert workflow_dict["environment_variables"][1]["value"] == "text"
|
||||
|
||||
workflow_dict = workflow.to_dict(include_secret=True)
|
||||
assert workflow_dict["environment_variables"][0]["value"] == "secret"
|
||||
assert workflow_dict["environment_variables"][1]["value"] == "text"
|
||||
|
||||
|
||||
class TestWorkflowNodeExecution:
|
||||
def test_execution_metadata_dict(self):
|
||||
node_exec = WorkflowNodeExecutionModel()
|
||||
node_exec.execution_metadata = None
|
||||
assert node_exec.execution_metadata_dict == {}
|
||||
|
||||
original = {"a": 1, "b": ["2"]}
|
||||
node_exec.execution_metadata = json.dumps(original)
|
||||
assert node_exec.execution_metadata_dict == original
|
||||
|
||||
|
||||
class TestIsSystemVariableEditable:
|
||||
def test_is_system_variable(self):
|
||||
cases = [
|
||||
("query", True),
|
||||
("files", True),
|
||||
("dialogue_count", False),
|
||||
("conversation_id", False),
|
||||
("user_id", False),
|
||||
("app_id", False),
|
||||
("workflow_id", False),
|
||||
("workflow_run_id", False),
|
||||
]
|
||||
for name, editable in cases:
|
||||
assert editable == is_system_variable_editable(name)
|
||||
|
||||
assert is_system_variable_editable("invalid_or_new_system_variable") == False
|
||||
|
||||
|
||||
class TestWorkflowDraftVariableGetValue:
|
||||
def test_get_value_by_case(self):
|
||||
@dataclasses.dataclass
|
||||
class TestCase:
|
||||
name: str
|
||||
value: Segment
|
||||
|
||||
tenant_id = "test_tenant_id"
|
||||
|
||||
test_file = File(
|
||||
tenant_id=tenant_id,
|
||||
type=FileType.IMAGE,
|
||||
transfer_method=FileTransferMethod.REMOTE_URL,
|
||||
remote_url="https://example.com/example.jpg",
|
||||
filename="example.jpg",
|
||||
extension=".jpg",
|
||||
mime_type="image/jpeg",
|
||||
size=100,
|
||||
)
|
||||
cases: list[TestCase] = [
|
||||
TestCase(
|
||||
name="number/int",
|
||||
value=build_segment(1),
|
||||
),
|
||||
TestCase(
|
||||
name="number/float",
|
||||
value=build_segment(1.0),
|
||||
),
|
||||
TestCase(
|
||||
name="string",
|
||||
value=build_segment("a"),
|
||||
),
|
||||
TestCase(
|
||||
name="object",
|
||||
value=build_segment({}),
|
||||
),
|
||||
TestCase(
|
||||
name="file",
|
||||
value=build_segment(test_file),
|
||||
),
|
||||
TestCase(
|
||||
name="array[any]",
|
||||
value=build_segment([1, "a"]),
|
||||
),
|
||||
TestCase(
|
||||
name="array[string]",
|
||||
value=build_segment(["a", "b"]),
|
||||
),
|
||||
TestCase(
|
||||
name="array[number]/int",
|
||||
value=build_segment([1, 2]),
|
||||
),
|
||||
TestCase(
|
||||
name="array[number]/float",
|
||||
value=build_segment([1.0, 2.0]),
|
||||
),
|
||||
TestCase(
|
||||
name="array[number]/mixed",
|
||||
value=build_segment([1, 2.0]),
|
||||
),
|
||||
TestCase(
|
||||
name="array[object]",
|
||||
value=build_segment([{}, {"a": 1}]),
|
||||
),
|
||||
TestCase(
|
||||
name="none",
|
||||
value=build_segment(None),
|
||||
),
|
||||
]
|
||||
|
||||
for idx, c in enumerate(cases, 1):
|
||||
fail_msg = f"test case {c.name} failed, index={idx}"
|
||||
draft_var = WorkflowDraftVariable()
|
||||
draft_var.set_value(c.value)
|
||||
assert c.value == draft_var.get_value(), fail_msg
|
||||
|
||||
def test_file_variable_preserves_all_fields(self):
|
||||
"""Test that File type variables preserve all fields during encoding/decoding."""
|
||||
tenant_id = "test_tenant_id"
|
||||
|
||||
# Create a File with specific field values
|
||||
test_file = File(
|
||||
id="test_file_id",
|
||||
tenant_id=tenant_id,
|
||||
type=FileType.IMAGE,
|
||||
transfer_method=FileTransferMethod.REMOTE_URL,
|
||||
remote_url="https://example.com/test.jpg",
|
||||
filename="test.jpg",
|
||||
extension=".jpg",
|
||||
mime_type="image/jpeg",
|
||||
size=12345, # Specific size to test preservation
|
||||
storage_key="test_storage_key",
|
||||
)
|
||||
|
||||
# Create a FileSegment and WorkflowDraftVariable
|
||||
file_segment = build_segment(test_file)
|
||||
draft_var = WorkflowDraftVariable()
|
||||
draft_var.set_value(file_segment)
|
||||
|
||||
# Retrieve the value and verify all fields are preserved
|
||||
retrieved_segment = draft_var.get_value()
|
||||
retrieved_file = retrieved_segment.value
|
||||
|
||||
# Verify all important fields are preserved
|
||||
assert retrieved_file.id == test_file.id
|
||||
assert retrieved_file.tenant_id == test_file.tenant_id
|
||||
assert retrieved_file.type == test_file.type
|
||||
assert retrieved_file.transfer_method == test_file.transfer_method
|
||||
assert retrieved_file.remote_url == test_file.remote_url
|
||||
assert retrieved_file.filename == test_file.filename
|
||||
assert retrieved_file.extension == test_file.extension
|
||||
assert retrieved_file.mime_type == test_file.mime_type
|
||||
assert retrieved_file.size == test_file.size # This was the main issue being fixed
|
||||
# Note: storage_key is not serialized in model_dump() so it won't be preserved
|
||||
|
||||
# Verify the segments have the same type and the important fields match
|
||||
assert file_segment.value_type == retrieved_segment.value_type
|
||||
|
||||
def test_get_and_set_value(self):
|
||||
draft_var = WorkflowDraftVariable()
|
||||
int_var = IntegerSegment(value=1)
|
||||
draft_var.set_value(int_var)
|
||||
value = draft_var.get_value()
|
||||
assert value == int_var
|
||||
Reference in New Issue
Block a user