dify/api/tests/unit_tests/models/test_workflow.py

import dataclasses
import json
from unittest import mock
from uuid import uuid4

import pytest
from sqlalchemy.orm import Session

from constants import HIDDEN_VALUE
from core.helper import encrypter
from core.workflow.file_reference import build_file_reference
from core.workflow.llm_environment_variable import LLMEnvironmentVariable
from factories.variable_factory import build_segment
from graphon.file import File, FileTransferMethod, FileType
from graphon.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable
from graphon.variables.segments import IntegerSegment, Segment
from models.account import Account
from models.tools import WorkflowToolProvider
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(name="var1", value="value1", id=str(uuid4()), selector=["env", "var1"])
    variable2 = IntegerVariable(name="var2", value=123, id=str(uuid4()), selector=["env", "var2"])
    variable3 = SecretVariable(name="var3", value="secret", id=str(uuid4()), selector=["env", "var3"])
    variable4 = FloatVariable(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_llm_environment_variable_round_trip():
    workflow = Workflow(
        tenant_id="tenant_id",
        app_id="app_id",
        type="workflow",
        version="draft",
        graph="{}",
        features="{}",
        created_by="account_id",
        environment_variables=[],
        conversation_variables=[],
    )
    variable = LLMEnvironmentVariable(
        name="for_research",
        value={"provider": "langgenius/anthropic/anthropic", "name": "claude-sonnet", "mode": "chat"},
        id=str(uuid4()),
        selector=["env", "for_research"],
    )

    workflow.environment_variables = [variable]

    assert workflow.environment_variables == [variable]
    assert json.loads(workflow._environment_variables)["for_research"]["value_type"] == "llm"
    assert workflow.to_dict()["environment_variables"][0]["value_type"] == "llm"


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(name="var1", value="value1", id=str(uuid4()), selector=["env", "var1"])
    variable2 = IntegerVariable(name="var2", value=123, id=str(uuid4()), selector=["env", "var2"])
    variable3 = SecretVariable(name="var3", value="secret", id=str(uuid4()), selector=["env", "var3"])
    variable4 = FloatVariable(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(name="secret", value="secret", id=str(uuid4())),
            StringVariable(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"


@pytest.mark.parametrize("sqlite_session", [(Workflow, Account)], indirect=True)
def test_workflow_account_getters_use_caller_session(sqlite_session: Session):
    created_account = Account(name="Created Account", email="created@example.com")
    created_account.id = "created-account-id"
    updated_account = Account(name="Updated Account", email="updated@example.com")
    updated_account.id = "updated-account-id"
    decoy_account = Account(name="Decoy Account", email="decoy@example.com")
    decoy_account.id = "decoy-account-id"
    workflow = Workflow(
        tenant_id="tenant_id",
        app_id="app_id",
        type="workflow",
        version="draft",
        graph="{}",
        features="{}",
        created_by="created-account-id",
        environment_variables=[],
        conversation_variables=[],
        updated_by="updated-account-id",
    )
    sqlite_session.add_all([decoy_account, updated_account, workflow, created_account])
    sqlite_session.flush()

    assert workflow.get_created_by_account(session=sqlite_session) is created_account
    assert workflow.get_updated_by_account(session=sqlite_session) is updated_account


@pytest.mark.parametrize("sqlite_session", [(Workflow, WorkflowToolProvider)], indirect=True)
def test_workflow_tool_published_getter_uses_caller_session(sqlite_session: Session):
    workflow = Workflow(
        tenant_id="tenant_id",
        app_id="app_id",
        type="workflow",
        version="draft",
        graph="{}",
        features="{}",
        created_by="account_id",
        environment_variables=[],
        conversation_variables=[],
    )
    matching_provider = WorkflowToolProvider(
        name="matching-provider",
        label="Matching provider",
        icon="tool",
        app_id=workflow.app_id,
        version="1",
        user_id="account-id",
        tenant_id=workflow.tenant_id,
        description="Matching workflow tool",
    )
    decoy_provider = WorkflowToolProvider(
        name="decoy-provider",
        label="Decoy provider",
        icon="tool",
        app_id="other-app",
        version="1",
        user_id="account-id",
        tenant_id=workflow.tenant_id,
        description="Different app",
    )
    sqlite_session.add_all([decoy_provider, workflow, matching_provider])
    sqlite_session.flush()

    assert workflow.get_tool_published(session=sqlite_session) is True


def test_normalize_environment_variable_mappings_converts_full_mask_to_hidden_value():
    normalized = Workflow.normalize_environment_variable_mappings(
        [
            {
                "id": str(uuid4()),
                "name": "secret",
                "value": encrypter.full_mask_token(),
                "value_type": "secret",
            }
        ]
    )

    assert normalized[0]["value"] == HIDDEN_VALUE


def test_normalize_environment_variable_mappings_keeps_hidden_value():
    normalized = Workflow.normalize_environment_variable_mappings(
        [
            {
                "id": str(uuid4()),
                "name": "secret",
                "value": HIDDEN_VALUE,
                "value_type": "secret",
            }
        ]
    )

    assert normalized[0]["value"] == HIDDEN_VALUE


class TestWorkflowNodeExecution:
    def test_execution_metadata_dict(self):
        node_exec = WorkflowNodeExecutionModel(
            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(
            file_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(
            file_id="test_file_id",
            file_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.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_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self):
        persisted_file = File(
            file_id="test_file_id",
            file_type=FileType.DOCUMENT,
            transfer_method=FileTransferMethod.LOCAL_FILE,
            reference=build_file_reference(record_id="upload-1", storage_key="legacy-storage-key"),
            filename="test.txt",
            extension=".txt",
            mime_type="text/plain",
            size=12,
        )
        rebuilt_file = File(
            file_id="test_file_id",
            file_type=FileType.DOCUMENT,
            transfer_method=FileTransferMethod.LOCAL_FILE,
            reference=build_file_reference(record_id="upload-1"),
            filename="test.txt",
            extension=".txt",
            mime_type="text/plain",
            size=12,
            storage_key="canonical-storage-key",
        )
        draft_var = WorkflowDraftVariable(
            app_id="app-1",
        )
        draft_var.set_value(build_segment(persisted_file))
        draft_var._WorkflowDraftVariable__value = None

        with (
            mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"),
            mock.patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file,
        ):
            retrieved_segment = draft_var.get_value()

        assert retrieved_segment.value == rebuilt_file
        rebuild_file.assert_called_once()
        assert rebuild_file.call_args.kwargs["tenant_id"] == "tenant-1"

    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