diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py
index 3a5591613b1..a92953313f2 100644
--- a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py
+++ b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py
@@ -6,6 +6,7 @@ import pytest
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.entities.queue_entities import QueueErrorEvent
+from models import Tenant
 
 
 class DummyQueueManager(AppQueueManager):
@@ -56,7 +57,7 @@ class TestBaseAppQueueManager:
             mock_redis.setex.return_value = True
             manager = DummyQueueManager(task_id="t1", user_id="u1", invoke_from=InvokeFrom.SERVICE_API)
 
-        bad = SimpleNamespace(_sa_instance_state=True)
+        bad = Tenant(name="Queued ORM model")
         with pytest.raises(TypeError):
             manager._check_for_sqlalchemy_models(bad)
 
diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry.py b/api/tests/unit_tests/core/workflow/test_workflow_entry.py
index 3a84255063c..23577105945 100644
--- a/api/tests/unit_tests/core/workflow/test_workflow_entry.py
+++ b/api/tests/unit_tests/core/workflow/test_workflow_entry.py
@@ -1,3 +1,4 @@
+import json
 from types import SimpleNamespace
 
 import pytest
@@ -10,12 +11,12 @@ from core.workflow.variable_prefixes import (
     ENVIRONMENT_VARIABLE_NODE_ID,
 )
 from core.workflow.workflow_entry import WorkflowEntry
-from graphon.entities.graph_config import NodeConfigDictAdapter
 from graphon.file import File, FileTransferMethod, FileType
 from graphon.nodes.code.code_node import CodeNode
 from graphon.nodes.code.limits import CodeNodeLimits
 from graphon.runtime import VariablePool
 from graphon.variables.variables import StringVariable
+from models.workflow import Workflow, WorkflowType
 
 
 @pytest.fixture(autouse=True)
@@ -116,18 +117,19 @@ class TestWorkflowEntry:
         }
         node_config = {"id": node_id, "data": node_data}
 
-        class StubWorkflow:
-            def __init__(self):
-                self.tenant_id = "tenant"
-                self.app_id = "app"
-                self.id = "workflow"
-                self.graph_dict = {"nodes": [node_config], "edges": []}
-
-            def get_node_config_by_id(self, target_id: str):
-                assert target_id == node_id
-                return NodeConfigDictAdapter.validate_python(node_config)
-
-        workflow = StubWorkflow()
+        workflow = Workflow.new(
+            tenant_id="tenant",
+            app_id="app",
+            type=WorkflowType.WORKFLOW,
+            version=Workflow.VERSION_DRAFT,
+            graph=json.dumps({"nodes": [node_config], "edges": []}),
+            features="{}",
+            created_by="account",
+            environment_variables=[],
+            conversation_variables=[],
+            rag_pipeline_variables=[],
+        )
+        workflow.id = "workflow"
         variable_pool = VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={})
         expected_limits = CodeNodeLimits(
             max_string_length=dify_config.CODE_MAX_STRING_LENGTH,
diff --git a/api/tests/unit_tests/services/retention/workflow_run/test_restore_archived_workflow_run.py b/api/tests/unit_tests/services/retention/workflow_run/test_restore_archived_workflow_run.py
index 62c042ff2cf..8d30ca231c7 100644
--- a/api/tests/unit_tests/services/retention/workflow_run/test_restore_archived_workflow_run.py
+++ b/api/tests/unit_tests/services/retention/workflow_run/test_restore_archived_workflow_run.py
@@ -19,7 +19,7 @@ import pytest
 from pydantic import ValidationError
 from sqlalchemy import Column, Engine, Integer, MetaData, String, Table, delete, event, func, select
 from sqlalchemy.dialects.sqlite import insert as sqlite_insert
-from sqlalchemy.orm import Session, sessionmaker
+from sqlalchemy.orm import Session, registry, sessionmaker
 
 from libs.archive_storage import ArchiveStorageNotConfiguredError
 from models.enums import CreatorUserRole
@@ -535,10 +535,13 @@ class TestGetModelColumnInfo:
             Column("defaulted_field", String(255), nullable=False, default="x"),
         )
 
-        class MockModel:
+        mapper_registry = registry()
+
+        @mapper_registry.mapped
+        class _AutoincrementModel:
             __table__ = test_table
 
-        _, required_columns, non_nullable_with_default = restore._get_model_column_info(MockModel)
+        _, required_columns, non_nullable_with_default = restore._get_model_column_info(_AutoincrementModel)
 
         assert required_columns == {"required_field"}
         assert "id" in non_nullable_with_default