test(api): fix tests broken by refactor

api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py

@@ -3,6 +3,7 @@ import secrets
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
+from sqlalchemy import Engine
 from sqlalchemy.orm import Session
 
 from core.variables import StringSegment
@@ -17,6 +18,18 @@ from services.workflow_draft_variable_service import (
 )
 
 
+@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
+
+
 class TestDraftVariableSaver:
     def _get_test_app_id(self):
         suffix = secrets.token_hex(6)
@@ -104,11 +117,9 @@ class TestWorkflowDraftVariableService:
             conversation_variables=[],
         )
 
-    def test_reset_conversation_variable(self):
+    def test_reset_conversation_variable(self, mock_session):
         """Test resetting a conversation variable"""
-        mock_session = Mock(spec=Session)
-        mock_session_maker = Mock()
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        service = WorkflowDraftVariableService(mock_session)
 
         test_app_id = self._get_test_app_id()
         workflow = self._create_test_workflow(test_app_id)
@@ -131,11 +142,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):
+    def test_reset_node_variable_with_no_execution_id(self, mock_session):
         """Test resetting a node variable with no execution ID - should delete variable"""
-        mock_session = Mock(spec=Session)
-        mock_session_maker = Mock()
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        service = WorkflowDraftVariableService(mock_session)
 
         test_app_id = self._get_test_app_id()
         workflow = self._create_test_workflow(test_app_id)
@@ -159,14 +168,21 @@ class TestWorkflowDraftVariableService:
         mock_session.flush.assert_called_once()
         assert result is None
 
-    def test_reset_node_variable_with_missing_execution_record(self):
+    def test_reset_node_variable_with_missing_execution_record(
+        self,
+        mock_engine,
+        mock_session,
+        monkeypatch,
+    ):
         """Test resetting a node variable when execution record doesn't exist"""
-        mock_session = Mock(spec=Session)
+        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_session
+        mock_session_maker.return_value.__enter__.return_value = mock_repo_session
         mock_session_maker.return_value.__exit__.return_value = None
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        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()
@@ -184,19 +200,28 @@ class TestWorkflowDraftVariableService:
 
         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 result is None
 
-    def test_reset_node_variable_with_valid_execution_record(self):
+    def test_reset_node_variable_with_valid_execution_record(
+        self,
+        mock_session,
+        monkeypatch,
+    ):
         """Test resetting a node variable with valid execution record - should restore from execution"""
-        mock_session = Mock(spec=Session)
+        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_session
+        mock_session_maker.return_value.__enter__.return_value = mock_repo_session
         mock_session_maker.return_value.__exit__.return_value = None
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        mock_session_maker = monkeypatch.setattr(
+            "services.workflow_draft_variable_service.sessionmaker", mock_session_maker
+        )
+        service = WorkflowDraftVariableService(mock_session)
 
         # Create mock execution record
         mock_execution = Mock(spec=WorkflowNodeExecutionModel)
@@ -232,11 +257,9 @@ class TestWorkflowDraftVariableService:
             # Should return the updated variable
             assert result == variable
 
-    def test_reset_non_editable_system_variable_raises_error(self):
+    def test_reset_non_editable_system_variable_raises_error(self, mock_session):
         """Test that resetting a non-editable system variable raises an error"""
-        mock_session = Mock(spec=Session)
-        mock_session_maker = Mock()
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        service = WorkflowDraftVariableService(mock_session)
 
         test_app_id = self._get_test_app_id()
         workflow = self._create_test_workflow(test_app_id)
@@ -256,11 +279,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):
+    def test_reset_editable_system_variable_succeeds(self, mock_session):
         """Test that resetting an editable system variable succeeds"""
-        mock_session = Mock(spec=Session)
-        mock_session_maker = Mock()
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        service = WorkflowDraftVariableService(mock_session)
 
         test_app_id = self._get_test_app_id()
         workflow = self._create_test_workflow(test_app_id)
@@ -290,11 +311,9 @@ class TestWorkflowDraftVariableService:
         assert variable.last_edited_at is None
         mock_session.flush.assert_called()
 
-    def test_reset_query_system_variable_succeeds(self):
+    def test_reset_query_system_variable_succeeds(self, mock_session):
         """Test that resetting query system variable (another editable one) succeeds"""
-        mock_session = Mock(spec=Session)
-        mock_session_maker = Mock()
-        service = WorkflowDraftVariableService(mock_session, mock_session_maker)
+        service = WorkflowDraftVariableService(mock_session)
 
         test_app_id = self._get_test_app_id()
         workflow = self._create_test_workflow(test_app_id)