2025-06-24 09:05:29 +08:00
|
|
|
import dataclasses
|
2025-05-21 18:38:16 +08:00
|
|
|
import json
|
2024-07-22 15:29:39 +08:00
|
|
|
from unittest import mock
|
|
|
|
from uuid import uuid4
|
|
|
|
|
|
|
|
from constants import HIDDEN_VALUE
|
2025-06-24 09:05:29 +08:00
|
|
|
from core.file.enums import FileTransferMethod, FileType
|
|
|
|
from core.file.models import File
|
2024-10-21 10:43:49 +08:00
|
|
|
from core.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable
|
2025-06-24 09:05:29 +08:00
|
|
|
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
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
|
|
|
|
def test_environment_variables():
|
2025-05-27 10:56:23 +08:00
|
|
|
# tenant_id context variable removed - using current_user.current_tenant_id directly
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
# Create a Workflow instance
|
2024-08-20 17:52:21 +08:00
|
|
|
workflow = Workflow(
|
2024-08-23 23:52:25 +08:00
|
|
|
tenant_id="tenant_id",
|
|
|
|
app_id="app_id",
|
|
|
|
type="workflow",
|
|
|
|
version="draft",
|
|
|
|
graph="{}",
|
|
|
|
features="{}",
|
|
|
|
created_by="account_id",
|
2024-08-20 17:52:21 +08:00
|
|
|
environment_variables=[],
|
|
|
|
conversation_variables=[],
|
|
|
|
)
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
# Create some EnvironmentVariable instances
|
2024-12-03 13:56:40 +08:00
|
|
|
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"]}
|
|
|
|
)
|
2024-07-22 15:29:39 +08:00
|
|
|
|
2025-05-27 10:56:23 +08:00
|
|
|
# Mock current_user as an EndUser
|
|
|
|
mock_user = mock.Mock()
|
|
|
|
mock_user.tenant_id = "tenant_id"
|
|
|
|
|
2024-07-22 15:29:39 +08:00
|
|
|
with (
|
2024-08-23 23:52:25 +08:00
|
|
|
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
|
|
|
|
mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"),
|
2025-05-27 10:56:23 +08:00
|
|
|
mock.patch("models.workflow.current_user", mock_user),
|
2024-07-22 15:29:39 +08:00
|
|
|
):
|
|
|
|
# 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():
|
2025-05-27 10:56:23 +08:00
|
|
|
# tenant_id context variable removed - using current_user.current_tenant_id directly
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
# Create a Workflow instance
|
2024-08-20 17:52:21 +08:00
|
|
|
workflow = Workflow(
|
2024-08-23 23:52:25 +08:00
|
|
|
tenant_id="tenant_id",
|
|
|
|
app_id="app_id",
|
|
|
|
type="workflow",
|
|
|
|
version="draft",
|
|
|
|
graph="{}",
|
|
|
|
features="{}",
|
|
|
|
created_by="account_id",
|
2024-08-20 17:52:21 +08:00
|
|
|
environment_variables=[],
|
|
|
|
conversation_variables=[],
|
|
|
|
)
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
# Create some EnvironmentVariable instances
|
2024-12-03 13:56:40 +08:00
|
|
|
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"]}
|
|
|
|
)
|
2024-07-22 15:29:39 +08:00
|
|
|
|
2025-05-27 10:56:23 +08:00
|
|
|
# Mock current_user as an EndUser
|
|
|
|
mock_user = mock.Mock()
|
|
|
|
mock_user.tenant_id = "tenant_id"
|
|
|
|
|
2024-07-22 15:29:39 +08:00
|
|
|
with (
|
2024-08-23 23:52:25 +08:00
|
|
|
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
|
|
|
|
mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"),
|
2025-05-27 10:56:23 +08:00
|
|
|
mock.patch("models.workflow.current_user", mock_user),
|
2024-07-22 15:29:39 +08:00
|
|
|
):
|
|
|
|
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={
|
2024-08-23 23:52:25 +08:00
|
|
|
"name": "new name",
|
|
|
|
"value": HIDDEN_VALUE,
|
2024-07-22 15:29:39 +08:00
|
|
|
}
|
|
|
|
)
|
|
|
|
|
|
|
|
workflow.environment_variables = variables
|
2024-08-23 23:52:25 +08:00
|
|
|
assert workflow.environment_variables[2].name == "new name"
|
2024-07-22 15:29:39 +08:00
|
|
|
assert workflow.environment_variables[2].value == variable3.value
|
|
|
|
|
|
|
|
|
|
|
|
def test_to_dict():
|
2025-05-27 10:56:23 +08:00
|
|
|
# tenant_id context variable removed - using current_user.current_tenant_id directly
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
# Create a Workflow instance
|
2024-08-20 17:52:21 +08:00
|
|
|
workflow = Workflow(
|
2024-08-23 23:52:25 +08:00
|
|
|
tenant_id="tenant_id",
|
|
|
|
app_id="app_id",
|
|
|
|
type="workflow",
|
|
|
|
version="draft",
|
|
|
|
graph="{}",
|
|
|
|
features="{}",
|
|
|
|
created_by="account_id",
|
2024-08-20 17:52:21 +08:00
|
|
|
environment_variables=[],
|
|
|
|
conversation_variables=[],
|
|
|
|
)
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
# Create some EnvironmentVariable instances
|
|
|
|
|
2025-05-27 10:56:23 +08:00
|
|
|
# Mock current_user as an EndUser
|
|
|
|
mock_user = mock.Mock()
|
|
|
|
mock_user.tenant_id = "tenant_id"
|
|
|
|
|
2024-07-22 15:29:39 +08:00
|
|
|
with (
|
2024-08-23 23:52:25 +08:00
|
|
|
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
|
|
|
|
mock.patch("core.helper.encrypter.decrypt_token", return_value="secret"),
|
2025-05-27 10:56:23 +08:00
|
|
|
mock.patch("models.workflow.current_user", mock_user),
|
2024-07-22 15:29:39 +08:00
|
|
|
):
|
|
|
|
# Set the environment_variables property of the Workflow instance
|
|
|
|
workflow.environment_variables = [
|
2024-08-23 23:52:25 +08:00
|
|
|
SecretVariable.model_validate({"name": "secret", "value": "secret", "id": str(uuid4())}),
|
|
|
|
StringVariable.model_validate({"name": "text", "value": "text", "id": str(uuid4())}),
|
2024-07-22 15:29:39 +08:00
|
|
|
]
|
|
|
|
|
|
|
|
workflow_dict = workflow.to_dict()
|
2024-08-23 23:52:25 +08:00
|
|
|
assert workflow_dict["environment_variables"][0]["value"] == ""
|
|
|
|
assert workflow_dict["environment_variables"][1]["value"] == "text"
|
2024-07-22 15:29:39 +08:00
|
|
|
|
|
|
|
workflow_dict = workflow.to_dict(include_secret=True)
|
2024-08-23 23:52:25 +08:00
|
|
|
assert workflow_dict["environment_variables"][0]["value"] == "secret"
|
|
|
|
assert workflow_dict["environment_variables"][1]["value"] == "text"
|
2025-05-21 18:38:16 +08:00
|
|
|
|
|
|
|
|
|
|
|
class TestWorkflowNodeExecution:
|
|
|
|
def test_execution_metadata_dict(self):
|
2025-05-30 04:56:37 +08:00
|
|
|
node_exec = WorkflowNodeExecutionModel()
|
2025-05-21 18:38:16 +08:00
|
|
|
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
|
2025-06-24 09:05:29 +08:00
|
|
|
|
|
|
|
|
|
|
|
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
|