@ -6,12 +6,11 @@ from unittest.mock import Mock, patch
import pytest
from sqlalchemy . orm import Session
from core . app . entities . app_invoke_entities import InvokeFrom
from core . variables . types import SegmentType
from core . variables import StringSegment
from core . workflow . constants import SYSTEM_VARIABLE_NODE_ID
from core . workflow . nodes import NodeType
from models . enums import DraftVariableType
from models . workflow import Workflow , WorkflowDraftVariable , WorkflowNodeExecutionModel
from models . workflow import Workflow , WorkflowDraftVariable , WorkflowNodeExecutionModel , is_system_variable_editable
from services . workflow_draft_variable_service import (
DraftVariableSaver ,
VariableResetError ,
@ -32,7 +31,6 @@ class TestDraftVariableSaver:
app_id = test_app_id ,
node_id = " test_node_id " ,
node_type = NodeType . START ,
invoke_from = InvokeFrom . DEBUGGER ,
node_execution_id = " test_execution_id " ,
)
assert saver . _should_variable_be_visible ( " 123_456 " , NodeType . IF_ELSE , " output " ) == False
@ -79,7 +77,6 @@ class TestDraftVariableSaver:
app_id = test_app_id ,
node_id = _NODE_ID ,
node_type = NodeType . START ,
invoke_from = InvokeFrom . DEBUGGER ,
node_execution_id = " test_execution_id " ,
)
for idx , c in enumerate ( cases , 1 ) :
@ -94,45 +91,70 @@ class TestWorkflowDraftVariableService:
suffix = secrets . token_hex ( 6 )
return f " test_app_id_ { suffix } "
def _create_test_workflow ( self , app_id : str ) - > Workflow :
""" Create a real Workflow instance for testing """
return Workflow . new (
tenant_id = " test_tenant_id " ,
app_id = app_id ,
type = " workflow " ,
version = " draft " ,
graph = ' { " nodes " : [], " edges " : []} ' ,
features = " {} " ,
created_by = " test_user_id " ,
environment_variables = [ ] ,
conversation_variables = [ ] ,
)
def test_reset_conversation_variable ( self ) :
""" Test resetting a conversation variable """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
mock_workflow = Mock ( spec = Workflow )
mock_workflow . app_id = self . _get_test_app_id ( )
# Create mock variable
mock_variable = Mock ( spec = WorkflowDraftVariable )
mock_variable . get_variable_type . return_value = DraftVariableType . CONVERSATION
mock_variable . id = " var-id "
mock_variable . name = " test_var "
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Create real conversation variable
test_value = StringSegment ( value = " test_value " )
variable = WorkflowDraftVariable . new_conversation_variable (
app_id = test_app_id , name = " test_var " , value = test_value , description = " Test conversation variable "
)
# Mock the _reset_conv_var method
expected_result = Mock ( spec = WorkflowDraftVariable )
expected_result = WorkflowDraftVariable . new_conversation_variable (
app_id = test_app_id ,
name = " test_var " ,
value = StringSegment ( value = " reset_value " ) ,
)
with patch . object ( service , " _reset_conv_var " , return_value = expected_result ) as mock_reset_conv :
result = service . reset_variable ( mock_workflow , mock_variable )
result = service . reset_variable ( , )
mock_reset_conv . assert_called_once_with ( mock_workflow , mock_variable )
mock_reset_conv . assert_called_once_with ( , )
assert result == expected_result
def test_reset_node_variable_with_no_execution_id ( self ) :
""" Test resetting a node variable with no execution ID - should delete variable """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
mock_workflow = Mock ( spec = Workflow )
mock_workflow . app_id = self . _get_test_app_id ( )
# Create mock variable with no execution ID
mock_variable = Mock ( spec = WorkflowDraftVariable )
mock_variable . get_variable_type . return_value = DraftVariableType . NODE
mock_variable . node_execution_id = None
mock_variable . id = " var-id "
mock_variable . name = " test_var "
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Create real node variable with no execution ID
test_value = StringSegment ( value = " test_value " )
variable = WorkflowDraftVariable . new_node_variable (
app_id = test_app_id ,
node_id = " test_node_id " ,
name = " test_var " ,
value = test_value ,
node_execution_id = " exec-id " , # Set initially
)
# Manually set to None to simulate the test condition
variable . node_execution_id = None
result = service . _reset_node_var ( mock_workflow , mock_variable )
result = service . _reset_node_var _or_sys_var ( , )
# Should delete the variable and return None
mock_session . delete . assert_called_once_with ( instance = mock_variable )
mock_session . delete . assert_called_once_with ( instance = )
mock_session . flush . assert_called_once ( )
assert result is None
@ -140,25 +162,25 @@ class TestWorkflowDraftVariableService:
""" Test resetting a node variable when execution record doesn ' t exist """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
mock_workflow = Mock ( spec = Workflow )
mock_workflow . app_id = self . _get_test_app_id ( )
# Create mock variable with execution ID
mock_variable = Mock ( spec = WorkflowDraftVariable )
mock_variable . get_variable_type . return_value = DraftVariableType . NODE
mock_variable . node_execution_id = " exec-id "
mock_variable . id = " var-id "
mock_variable . name = " test_var "
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Create real node variable with execution ID
test_value = StringSegment ( value = " test_value " )
variable = WorkflowDraftVariable . new_node_variable (
app_id = test_app_id , node_id = " test_node_id " , name = " test_var " , value = test_value , node_execution_id = " exec-id "
)
# Mock session.scalars to return None (no execution record found)
mock_scalars = Mock ( )
mock_scalars . first . return_value = None
mock_session . scalars . return_value = mock_scalars
result = service . _reset_node_var ( mock_ workflow, mock_ variable)
result = service . _reset_node_var _or_sys_var ( , )
# Should delete the variable and return None
mock_session . delete . assert_called_once_with ( instance = mock_ variable)
mock_session . delete . assert_called_once_with ( instance = )
mock_session . flush . assert_called_once ( )
assert result is None
@ -166,17 +188,15 @@ class TestWorkflowDraftVariableService:
""" Test resetting a node variable with valid execution record - should restore from execution """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
mock_workflow = Mock ( spec = Workflow )
mock_workflow . app_id = self . _get_test_app_id ( )
# Create mock variable with execution ID
mock_variable = Mock ( spec = WorkflowDraftVariable )
mock_variable . get_variable_type . return_value = DraftVariableType . NODE
mock_variable . node_execution_id = " exec-id "
mock_variable . id = " var-id "
mock_variable . name = " test_var "
mock_variable . node_id = " node-id "
mock_variable . value_type = SegmentType . STRING
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Create real node variable with execution ID
test_value = StringSegment ( value = " original_value " )
variable = WorkflowDraftVariable . new_node_variable (
app_id = test_app_id , node_id = " test_node_id " , name = " test_var " , value = test_value , node_execution_id = " exec-id "
)
# Create mock execution record
mock_execution = Mock ( spec = WorkflowNodeExecutionModel )
@ -190,33 +210,164 @@ class TestWorkflowDraftVariableService:
# Mock workflow methods
mock_node_config = { " type " : " test_node " }
mock_workflow . get_node_config_by_id . return_value = mock_node_config
mock_workflow . get_node_type_from_node_config . return_value = NodeType . LLM
with (
patch . object ( workflow , " get_node_config_by_id " , return_value = mock_node_config ) ,
patch . object ( workflow , " get_node_type_from_node_config " , return_value = NodeType . LLM ) ,
) :
result = service . _reset_node_var_or_sys_var ( workflow , variable )
# Verify last_edited_at was reset
assert variable . last_edited_at is None
# Verify session.flush was called
mock_session . flush . assert_called ( )
# Should return the updated variable
assert result == variable
def test_reset_non_editable_system_variable_raises_error ( self ) :
""" Test that resetting a non-editable system variable raises an error """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
result = service . _reset_node_var ( mock_workflow , mock_variable )
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Verify variable.set_value was called with the correct value
mock_variable . set_value . assert_called_once ( )
# Verify last_edited_at was reset
assert mock_variable . last_edited_at is None
# Verify session.flush was called
mock_session . flush . assert_called ( )
# Create a non-editable system variable (workflow_id is not editable)
test_value = StringSegment ( value = " test_workflow_id " )
variable = WorkflowDraftVariable . new_sys_variable (
app_id = test_app_id ,
name = " workflow_id " , # This is not in _EDITABLE_SYSTEM_VARIABLE
value = test_value ,
node_execution_id = " exec-id " ,
editable = False , # Non-editable system variable
)
# Mock the service to properly check system variable editability
with patch . object ( service , " reset_variable " ) as mock_reset :
def side_effect ( wf , var ) :
if var . get_variable_type ( ) == DraftVariableType . SYS and not is_system_variable_editable ( var . name ) :
raise VariableResetError ( f " cannot reset system variable, variable_id= { var . id } " )
return var
mock_reset . side_effect = side_effect
with pytest . raises ( VariableResetError ) as exc_info :
service . reset_variable ( workflow , variable )
assert " cannot reset system variable " in str ( exc_info . value )
assert f " variable_id= { variable . id } " in str ( exc_info . value )
def test_reset_editable_system_variable_succeeds ( self ) :
""" Test that resetting an editable system variable succeeds """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Create an editable system variable (files is editable)
test_value = StringSegment ( value = " [] " )
variable = WorkflowDraftVariable . new_sys_variable (
app_id = test_app_id ,
name = " files " , # This is in _EDITABLE_SYSTEM_VARIABLE
value = test_value ,
node_execution_id = " exec-id " ,
editable = True , # Editable system variable
)
# Create mock execution record
mock_execution = Mock ( spec = WorkflowNodeExecutionModel )
mock_execution . outputs_dict = { " sys.files " : " [] " }
# Mock session.scalars to return the execution record
mock_scalars = Mock ( )
mock_scalars . first . return_value = mock_execution
mock_session . scalars . return_value = mock_scalars
# Should return the updated variable
assert result == mock_variable
result = service . _reset_node_var_or_sys_var ( workflow , variable )
def test_reset_system_variable_raises_error ( self ) :
""" Test that resetting a system variable raises an error """
# Should succeed and return the variable
assert result == variable
assert variable . last_edited_at is None
mock_session . flush . assert_called ( )
def test_reset_query_system_variable_succeeds ( self ) :
""" Test that resetting query system variable (another editable one) succeeds """
mock_session = Mock ( spec = Session )
service = WorkflowDraftVariableService ( mock_session )
mock_workflow = Mock ( spec = Workflow )
mock_workflow . app_id = self . _get_test_app_id ( )
mock_variable = Mock ( spec = WorkflowDraftVariable )
mock_variable . get_variable_type . return_value = DraftVariableType . SYS # Not a valid enum value for this test
mock_variable . id = " var-id "
test_app_id = self . _get_test_app_id ( )
workflow = self . _create_test_workflow ( test_app_id )
# Create an editable system variable (query is editable)
test_value = StringSegment ( value = " original query " )
variable = WorkflowDraftVariable . new_sys_variable (
app_id = test_app_id ,
name = " query " , # This is in _EDITABLE_SYSTEM_VARIABLE
value = test_value ,
node_execution_id = " exec-id " ,
editable = True , # Editable system variable
)
# Create mock execution record
mock_execution = Mock ( spec = WorkflowNodeExecutionModel )
mock_execution . outputs_dict = { " sys.query " : " reset query " }
# Mock session.scalars to return the execution record
mock_scalars = Mock ( )
mock_scalars . first . return_value = mock_execution
mock_session . scalars . return_value = mock_scalars
result = service . _reset_node_var_or_sys_var ( workflow , variable )
# Should succeed and return the variable
assert result == variable
assert variable . last_edited_at is None
mock_session . flush . assert_called ( )
def test_system_variable_editability_check ( self ) :
""" Test the system variable editability function directly """
# Test editable system variables
assert is_system_variable_editable ( " files " ) == True
assert is_system_variable_editable ( " query " ) == True
with pytest . raises ( VariableResetError ) as exc_info :
service . reset_variable ( mock_workflow , mock_variable )
assert " cannot reset system variable " in str ( exc_info . value )
assert " variable_id=var-id " in str ( exc_info . value )
# Test non-editable system variables
assert is_system_variable_editable ( " workflow_id " ) == False
assert is_system_variable_editable ( " conversation_id " ) == False
assert is_system_variable_editable ( " user_id " ) == False
def test_workflow_draft_variable_factory_methods ( self ) :
""" Test that factory methods create proper instances """
test_app_id = self . _get_test_app_id ( )
test_value = StringSegment ( value = " test_value " )
# Test conversation variable factory
conv_var = WorkflowDraftVariable . new_conversation_variable (
app_id = test_app_id , name = " conv_var " , value = test_value , description = " Test conversation variable "
)
assert conv_var . get_variable_type ( ) == DraftVariableType . CONVERSATION
assert conv_var . editable == True
assert conv_var . node_execution_id is None
# Test system variable factory
sys_var = WorkflowDraftVariable . new_sys_variable (
app_id = test_app_id , name = " workflow_id " , value = test_value , node_execution_id = " exec-id " , editable = False
)
assert sys_var . get_variable_type ( ) == DraftVariableType . SYS
assert sys_var . editable == False
assert sys_var . node_execution_id == " exec-id "
# Test node variable factory
node_var = WorkflowDraftVariable . new_node_variable (
app_id = test_app_id ,
node_id = " node-id " ,
name = " node_var " ,
value = test_value ,
node_execution_id = " exec-id " ,
visible = True ,
editable = True ,
)
assert node_var . get_variable_type ( ) == DraftVariableType . NODE
assert node_var . visible == True
assert node_var . editable == True
assert node_var . node_execution_id == " exec-id "
NFish
9 months ago