From d7860571aad97a20b070ae4e57f1ee0e4f3f12b1 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 21 Jul 2026 16:24:52 +0900 Subject: [PATCH] test: use SQLite sessions in services workflow (#39115) --- .../test_workflow_draft_variable_service.py | 188 +++++++----------- 1 file changed, 76 insertions(+), 112 deletions(-) diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py index c041a6203e9..9053f7c1b73 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py @@ -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"""