mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 18:58:35 +08:00
test: use SQLite sessions in services workflow (#39115)
This commit is contained in:
Asuka Minato
GitHub
parent
3683831738
commit
d7860571aa
@ -5,7 +5,6 @@ from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import Engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.workflow.system_variables import SystemVariableKey
|
||||
@ -35,17 +34,11 @@ from services.workflow_draft_variable_service import (
|
||||
_model_to_insertion_dict,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_engine() -> Engine:
|
||||
return Mock(spec=Engine)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_session(mock_engine) -> Session:
|
||||
mock_session = Mock(spec=Session)
|
||||
mock_session.get_bind.return_value = mock_engine
|
||||
return mock_session
|
||||
SQLITE_MODELS = (Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, WorkflowNodeExecutionModel)
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures("sqlite_session"),
|
||||
pytest.mark.parametrize("sqlite_session", [SQLITE_MODELS], indirect=True),
|
||||
]
|
||||
|
||||
|
||||
class TestDraftVariableSaver:
|
||||
@ -53,13 +46,12 @@ class TestDraftVariableSaver:
|
||||
suffix = secrets.token_hex(6)
|
||||
return f"test_app_id_{suffix}"
|
||||
|
||||
def test__should_variable_be_visible(self):
|
||||
mock_session = MagicMock(spec=Session)
|
||||
def test__should_variable_be_visible(self, sqlite_session: Session):
|
||||
mock_user = Account(name="test", email="test@example.com")
|
||||
mock_user.id = str(uuid.uuid4())
|
||||
test_app_id = self._get_test_app_id()
|
||||
saver = DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="test-tenant-id",
|
||||
app_id=test_app_id,
|
||||
node_id="test_node_id",
|
||||
@ -70,7 +62,7 @@ class TestDraftVariableSaver:
|
||||
assert saver._should_variable_be_visible("123_456", BuiltinNodeTypes.IF_ELSE, "output") == False
|
||||
assert saver._should_variable_be_visible("123", BuiltinNodeTypes.START, "output") == True
|
||||
|
||||
def test__normalize_variable_for_start_node(self):
|
||||
def test__normalize_variable_for_start_node(self, sqlite_session: Session):
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class TestCase:
|
||||
name: str
|
||||
@ -118,11 +110,10 @@ class TestDraftVariableSaver:
|
||||
),
|
||||
]
|
||||
|
||||
mock_session = MagicMock(spec=Session)
|
||||
mock_user = MagicMock()
|
||||
test_app_id = self._get_test_app_id()
|
||||
saver = DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="test-tenant-id",
|
||||
app_id=test_app_id,
|
||||
node_id=_NODE_ID,
|
||||
@ -136,12 +127,11 @@ class TestDraftVariableSaver:
|
||||
assert node_id == c.expected_node_id, fail_msg
|
||||
assert name == c.expected_name, fail_msg
|
||||
|
||||
def test_build_variables_from_start_mapping_rebuilds_system_files(self):
|
||||
mock_session = MagicMock(spec=Session)
|
||||
def test_build_variables_from_start_mapping_rebuilds_system_files(self, sqlite_session: Session):
|
||||
mock_user = MagicMock(spec=Account)
|
||||
mock_user.id = str(uuid.uuid4())
|
||||
saver = DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="tenant-1",
|
||||
app_id=self._get_test_app_id(),
|
||||
node_id="start",
|
||||
@ -176,17 +166,7 @@ class TestDraftVariableSaver:
|
||||
rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1")
|
||||
|
||||
@pytest.fixture
|
||||
def mock_session(self):
|
||||
"""Mock SQLAlchemy session."""
|
||||
from sqlalchemy import Engine
|
||||
|
||||
mock_session = MagicMock(spec=Session)
|
||||
mock_engine = MagicMock(spec=Engine)
|
||||
mock_session.get_bind.return_value = mock_engine
|
||||
return mock_session
|
||||
|
||||
@pytest.fixture
|
||||
def draft_saver(self, mock_session):
|
||||
def draft_saver(self, sqlite_session: Session):
|
||||
"""Create DraftVariableSaver instance with user context."""
|
||||
# Create a mock user
|
||||
mock_user = MagicMock(spec=Account)
|
||||
@ -194,7 +174,7 @@ class TestDraftVariableSaver:
|
||||
mock_user.tenant_id = "test-tenant-id"
|
||||
|
||||
return DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="test-tenant-id",
|
||||
app_id="test-app-id",
|
||||
node_id="test-node-id",
|
||||
@ -203,7 +183,7 @@ class TestDraftVariableSaver:
|
||||
user=mock_user,
|
||||
)
|
||||
|
||||
def test_draft_saver_with_small_variables(self, draft_saver: DraftVariableSaver, mock_session):
|
||||
def test_draft_saver_with_small_variables(self, draft_saver: DraftVariableSaver):
|
||||
with patch(
|
||||
"services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True
|
||||
) as _mock_try_offload:
|
||||
@ -215,7 +195,7 @@ class TestDraftVariableSaver:
|
||||
assert draft_var.file_id is None
|
||||
_mock_try_offload.return_value = None
|
||||
|
||||
def test_draft_saver_with_large_variables(self, draft_saver: DraftVariableSaver, mock_session):
|
||||
def test_draft_saver_with_large_variables(self, draft_saver: DraftVariableSaver):
|
||||
with patch(
|
||||
"services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True
|
||||
) as _mock_try_offload:
|
||||
@ -237,12 +217,12 @@ class TestDraftVariableSaver:
|
||||
# Should not have large variable metadata
|
||||
assert draft_var.file_id == mock_draft_var_file.id
|
||||
|
||||
def test_try_offload_large_variable_uses_resource_tenant(self, mock_session):
|
||||
def test_try_offload_large_variable_uses_resource_tenant(self, sqlite_session: Session):
|
||||
mock_user = MagicMock(spec=Account)
|
||||
mock_user.id = "test-user-id"
|
||||
mock_user.current_tenant_id = ""
|
||||
saver = DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="app-tenant-id",
|
||||
app_id="test-app-id",
|
||||
node_id="test-node-id",
|
||||
@ -284,15 +264,16 @@ class TestDraftVariableSaver:
|
||||
assert len(draft_vars) == 2
|
||||
|
||||
@patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True)
|
||||
def test_start_node_save_persists_sys_timestamp_and_workflow_run_id(self, mock_batch_upsert):
|
||||
def test_start_node_save_persists_sys_timestamp_and_workflow_run_id(
|
||||
self, mock_batch_upsert, sqlite_session: Session
|
||||
):
|
||||
"""Start node should persist common `sys.*` variables, not only `sys.files`."""
|
||||
mock_session = MagicMock(spec=Session)
|
||||
mock_user = MagicMock(spec=Account)
|
||||
mock_user.id = "test-user-id"
|
||||
mock_user.tenant_id = "test-tenant-id"
|
||||
|
||||
saver = DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="test-tenant-id",
|
||||
app_id="test-app-id",
|
||||
node_id="start-node-id",
|
||||
@ -321,14 +302,13 @@ class TestDraftVariableSaver:
|
||||
}
|
||||
|
||||
@patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True)
|
||||
def test_start_node_save_normalizes_reserved_prefix_outputs(self, mock_batch_upsert):
|
||||
mock_session = MagicMock(spec=Session)
|
||||
def test_start_node_save_normalizes_reserved_prefix_outputs(self, mock_batch_upsert, sqlite_session: Session):
|
||||
mock_user = MagicMock(spec=Account)
|
||||
mock_user.id = "test-user-id"
|
||||
mock_user.tenant_id = "test-tenant-id"
|
||||
|
||||
saver = DraftVariableSaver(
|
||||
session=mock_session,
|
||||
session=sqlite_session,
|
||||
tenant_id="test-tenant-id",
|
||||
app_id="test-app-id",
|
||||
node_id="start-node-id",
|
||||
@ -382,8 +362,8 @@ class TestWorkflowDraftVariableService:
|
||||
rag_pipeline_variables=[],
|
||||
)
|
||||
|
||||
def test_list_variables_without_values_excludes_node_ids(self, mock_session):
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
def test_list_variables_without_values_excludes_node_ids(self, sqlite_session: Session):
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
variable = WorkflowDraftVariable.new_node_variable(
|
||||
app_id="app-1",
|
||||
node_id="node-1",
|
||||
@ -391,8 +371,21 @@ class TestWorkflowDraftVariableService:
|
||||
value=StringSegment(value="value"),
|
||||
node_execution_id="execution-1",
|
||||
)
|
||||
mock_session.scalar.return_value = 1
|
||||
mock_session.scalars.return_value = [variable]
|
||||
variable.user_id = "user-1"
|
||||
excluded_system = WorkflowDraftVariable.new_sys_variable(
|
||||
app_id="app-1", name="query", value=StringSegment(value="hidden"), node_execution_id="execution-1"
|
||||
)
|
||||
excluded_system.user_id = "user-1"
|
||||
other_user = WorkflowDraftVariable.new_node_variable(
|
||||
app_id="app-1",
|
||||
node_id="node-2",
|
||||
name="output",
|
||||
value=StringSegment(value="other"),
|
||||
node_execution_id="execution-2",
|
||||
)
|
||||
other_user.user_id = "user-2"
|
||||
sqlite_session.add_all([variable, excluded_system, other_user])
|
||||
sqlite_session.commit()
|
||||
|
||||
result = service.list_variables_without_values(
|
||||
app_id="app-1",
|
||||
@ -403,16 +396,11 @@ class TestWorkflowDraftVariableService:
|
||||
)
|
||||
|
||||
assert result.total == 1
|
||||
assert result.variables == [variable]
|
||||
assert [item.id for item in result.variables] == [variable.id]
|
||||
|
||||
stmt = mock_session.scalars.call_args.args[0]
|
||||
compiled = stmt.compile()
|
||||
excluded_node_ids = next(value for value in compiled.params.values() if isinstance(value, (list, tuple)))
|
||||
assert set(excluded_node_ids) == {SYSTEM_VARIABLE_NODE_ID, CONVERSATION_VARIABLE_NODE_ID}
|
||||
|
||||
def test_reset_conversation_variable(self, mock_session):
|
||||
def test_reset_conversation_variable(self, sqlite_session: Session):
|
||||
"""Test resetting a conversation variable"""
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
test_app_id = self._get_test_app_id()
|
||||
workflow = self._create_test_workflow(test_app_id)
|
||||
@ -435,9 +423,9 @@ class TestWorkflowDraftVariableService:
|
||||
mock_reset_conv.assert_called_once_with(workflow, variable)
|
||||
assert result == expected_result
|
||||
|
||||
def test_reset_node_variable_with_no_execution_id(self, mock_session):
|
||||
def test_reset_node_variable_with_no_execution_id(self, sqlite_session: Session):
|
||||
"""Test resetting a node variable with no execution ID - should delete variable"""
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
test_app_id = self._get_test_app_id()
|
||||
workflow = self._create_test_workflow(test_app_id)
|
||||
@ -453,33 +441,18 @@ class TestWorkflowDraftVariableService:
|
||||
)
|
||||
# Manually set to None to simulate the test condition
|
||||
variable.node_execution_id = None
|
||||
sqlite_session.add(variable)
|
||||
sqlite_session.commit()
|
||||
|
||||
result = service._reset_node_var_or_sys_var(workflow, variable)
|
||||
|
||||
# Should delete the variable and return None
|
||||
mock_session.delete.assert_called_once_with(instance=variable)
|
||||
mock_session.flush.assert_called_once()
|
||||
assert sqlite_session.get(WorkflowDraftVariable, variable.id) is None
|
||||
assert result is None
|
||||
|
||||
def test_reset_node_variable_with_missing_execution_record(
|
||||
self,
|
||||
mock_engine,
|
||||
mock_session,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
def test_reset_node_variable_with_missing_execution_record(self, sqlite_session: Session):
|
||||
"""Test resetting a node variable when execution record doesn't exist"""
|
||||
mock_repo_session = Mock(spec=Session)
|
||||
|
||||
mock_session_maker = MagicMock()
|
||||
# Mock the context manager protocol for sessionmaker
|
||||
mock_session_maker.return_value.__enter__.return_value = mock_repo_session
|
||||
mock_session_maker.return_value.__exit__.return_value = None
|
||||
monkeypatch.setattr("services.workflow_draft_variable_service.sessionmaker", mock_session_maker)
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
|
||||
# Mock the repository to return None (no execution record found)
|
||||
service._api_node_execution_repo = Mock()
|
||||
service._api_node_execution_repo.get_execution_by_id.return_value = None
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
test_app_id = self._get_test_app_id()
|
||||
workflow = self._create_test_workflow(test_app_id)
|
||||
@ -489,32 +462,17 @@ class TestWorkflowDraftVariableService:
|
||||
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"
|
||||
)
|
||||
# Variable is editable by default from factory method
|
||||
sqlite_session.add(variable)
|
||||
sqlite_session.commit()
|
||||
|
||||
result = service._reset_node_var_or_sys_var(workflow, variable)
|
||||
|
||||
mock_session_maker.assert_called_once_with(bind=mock_engine, expire_on_commit=False)
|
||||
# Should delete the variable and return None
|
||||
mock_session.delete.assert_called_once_with(instance=variable)
|
||||
mock_session.flush.assert_called_once()
|
||||
assert sqlite_session.get(WorkflowDraftVariable, variable.id) is None
|
||||
assert result is None
|
||||
|
||||
def test_reset_node_variable_with_valid_execution_record(
|
||||
self,
|
||||
mock_session,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Test resetting a node variable with valid execution record - should restore from execution"""
|
||||
mock_repo_session = Mock(spec=Session)
|
||||
|
||||
mock_session_maker = MagicMock()
|
||||
# Mock the context manager protocol for sessionmaker
|
||||
mock_session_maker.return_value.__enter__.return_value = mock_repo_session
|
||||
mock_session_maker.return_value.__exit__.return_value = None
|
||||
mock_session_maker = monkeypatch.setattr(
|
||||
"services.workflow_draft_variable_service.sessionmaker", mock_session_maker
|
||||
)
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
def test_reset_node_variable_with_valid_execution_record(self, sqlite_session: Session):
|
||||
"""Reset a node variable from its execution output and flush the restored value."""
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
# Create mock execution record
|
||||
mock_execution = Mock(spec=WorkflowNodeExecutionModel)
|
||||
@ -532,11 +490,13 @@ class TestWorkflowDraftVariableService:
|
||||
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"
|
||||
)
|
||||
# Variable is editable by default from factory method
|
||||
sqlite_session.add(variable)
|
||||
sqlite_session.commit()
|
||||
|
||||
# Mock workflow methods
|
||||
mock_node_config = {"type": "test_node"}
|
||||
with (
|
||||
patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush,
|
||||
patch.object(workflow, "get_node_config_by_id", return_value=mock_node_config, autospec=True),
|
||||
patch.object(workflow, "get_node_type_from_node_config", return_value=BuiltinNodeTypes.LLM, autospec=True),
|
||||
):
|
||||
@ -544,15 +504,13 @@ class TestWorkflowDraftVariableService:
|
||||
|
||||
# Verify last_edited_at was reset
|
||||
assert variable.last_edited_at is None
|
||||
# Verify session.flush was called
|
||||
mock_session.flush.assert_called()
|
||||
|
||||
flush.assert_called()
|
||||
# Should return the updated variable
|
||||
assert result == variable
|
||||
|
||||
def test_reset_non_editable_system_variable_raises_error(self, mock_session):
|
||||
def test_reset_non_editable_system_variable_raises_error(self, sqlite_session: Session):
|
||||
"""Test that resetting a non-editable system variable raises an error"""
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
test_app_id = self._get_test_app_id()
|
||||
workflow = self._create_test_workflow(test_app_id)
|
||||
@ -572,9 +530,9 @@ class TestWorkflowDraftVariableService:
|
||||
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, mock_session):
|
||||
def test_reset_editable_system_variable_succeeds(self, sqlite_session: Session):
|
||||
"""Test that resetting an editable system variable succeeds"""
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
test_app_id = self._get_test_app_id()
|
||||
workflow = self._create_test_workflow(test_app_id)
|
||||
@ -588,6 +546,8 @@ class TestWorkflowDraftVariableService:
|
||||
node_execution_id="exec-id",
|
||||
editable=True, # Editable system variable
|
||||
)
|
||||
sqlite_session.add(variable)
|
||||
sqlite_session.commit()
|
||||
|
||||
# Create mock execution record
|
||||
mock_execution = Mock(spec=WorkflowNodeExecutionModel)
|
||||
@ -597,16 +557,17 @@ class TestWorkflowDraftVariableService:
|
||||
service._api_node_execution_repo = Mock()
|
||||
service._api_node_execution_repo.get_execution_by_id.return_value = mock_execution
|
||||
|
||||
result = service._reset_node_var_or_sys_var(workflow, variable)
|
||||
with patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush:
|
||||
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()
|
||||
flush.assert_called()
|
||||
|
||||
def test_reset_query_system_variable_succeeds(self, mock_session):
|
||||
def test_reset_query_system_variable_succeeds(self, sqlite_session: Session):
|
||||
"""Test that resetting query system variable (another editable one) succeeds"""
|
||||
service = WorkflowDraftVariableService(mock_session)
|
||||
service = WorkflowDraftVariableService(sqlite_session)
|
||||
|
||||
test_app_id = self._get_test_app_id()
|
||||
workflow = self._create_test_workflow(test_app_id)
|
||||
@ -620,6 +581,8 @@ class TestWorkflowDraftVariableService:
|
||||
node_execution_id="exec-id",
|
||||
editable=True, # Editable system variable
|
||||
)
|
||||
sqlite_session.add(variable)
|
||||
sqlite_session.commit()
|
||||
|
||||
# Create mock execution record
|
||||
mock_execution = Mock(spec=WorkflowNodeExecutionModel)
|
||||
@ -629,12 +592,13 @@ class TestWorkflowDraftVariableService:
|
||||
service._api_node_execution_repo = Mock()
|
||||
service._api_node_execution_repo.get_execution_by_id.return_value = mock_execution
|
||||
|
||||
result = service._reset_node_var_or_sys_var(workflow, variable)
|
||||
with patch.object(sqlite_session, "flush", wraps=sqlite_session.flush) as flush:
|
||||
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()
|
||||
flush.assert_called()
|
||||
|
||||
def test_system_variable_editability_check(self):
|
||||
"""Test the system variable editability function directly"""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user