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 9c73bd3796..8b1348b75b 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
@@ -119,9 +119,7 @@ class TestWorkflowDraftVariableService:
 
     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)
@@ -146,9 +144,7 @@ class TestWorkflowDraftVariableService:
 
     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)
@@ -179,12 +175,14 @@ class TestWorkflowDraftVariableService:
         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()
@@ -214,12 +212,16 @@ class TestWorkflowDraftVariableService:
         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)
@@ -257,9 +259,7 @@ class TestWorkflowDraftVariableService:
 
     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)
@@ -279,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)
@@ -315,9 +313,7 @@ class TestWorkflowDraftVariableService:
 
     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)