test: exercise human input submissions with SQLite repository

api/tests/unit_tests/controllers/common/test_schema.py

@@ -1,12 +1,13 @@
 import sys
 from datetime import datetime
 from enum import StrEnum
+from types import SimpleNamespace
 from typing import Literal
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import pytest
 from flask import Flask
-from flask_restx import Namespace
+from flask_restx import Namespace, fields
 from pydantic import BaseModel, ConfigDict, Field
 
 
@@ -70,14 +71,11 @@ class ResponseAliasModel(BaseModel):
 
 
 @pytest.fixture(autouse=True)
-def mock_console_ns():
-    """Mock the console_ns to avoid circular imports during test collection."""
-    mock_ns = MagicMock(spec=Namespace)
-    mock_ns.models = {}
-
-    # Inject mock before importing schema module
-    with patch.dict(sys.modules, {"controllers.console": MagicMock(console_ns=mock_ns)}):
-        yield mock_ns
+def console_ns():
+    """Isolate the console import while exercising a real namespace."""
+    namespace = Namespace("console")
+    with patch.dict(sys.modules, {"controllers.console": SimpleNamespace(console_ns=namespace)}):
+        yield namespace
 
 
 def test_default_ref_template_value():
@@ -89,15 +87,12 @@ def test_default_ref_template_value():
 def test_register_schema_model_calls_namespace_schema_model():
     from controllers.common.schema import register_schema_model
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_schema_model(namespace, UserModel)
 
-    namespace.schema_model.assert_called_once()
-
-    model_name, schema = namespace.schema_model.call_args.args
-
-    assert model_name == "UserModel"
+    assert list(namespace.models) == ["UserModel"]
+    schema = namespace.models["UserModel"].__schema__
     assert isinstance(schema, dict)
     assert "properties" in schema
 
@@ -105,11 +100,11 @@ def test_register_schema_model_calls_namespace_schema_model():
 def test_register_schema_model_passes_schema_from_pydantic():
     from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0, register_schema_model
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_schema_model(namespace, UserModel)
 
-    schema = namespace.schema_model.call_args.args[1]
+    schema = namespace.models["UserModel"].__schema__
 
     expected_schema = UserModel.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_OPENAPI_3_0)
 
@@ -119,11 +114,11 @@ def test_register_schema_model_passes_schema_from_pydantic():
 def test_register_schema_model_promotes_nested_pydantic_definitions():
     from controllers.common.schema import DEFAULT_REF_TEMPLATE_OPENAPI_3_0, register_schema_model
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_schema_model(namespace, ParentModel)
 
-    called_schemas = {call.args[0]: call.args[1] for call in namespace.schema_model.call_args_list}
+    called_schemas = {name: model.__schema__ for name, model in namespace.models.items()}
     parent_schema = ParentModel.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_OPENAPI_3_0)
 
     assert set(called_schemas) == {"ParentModel", "ChildModel"}
@@ -135,20 +130,18 @@ def test_register_schema_model_promotes_nested_pydantic_definitions():
 def test_register_schema_models_registers_multiple_models():
     from controllers.common.schema import register_schema_models
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_schema_models(namespace, UserModel, ProductModel)
 
-    assert namespace.schema_model.call_count == 2
-
-    called_names = [call.args[0] for call in namespace.schema_model.call_args_list]
+    called_names = list(namespace.models)
     assert called_names == ["UserModel", "ProductModel"]
 
 
 def test_register_schema_models_calls_register_schema_model(monkeypatch: pytest.MonkeyPatch):
     from controllers.common.schema import register_schema_models
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     calls = []
 
@@ -171,13 +164,12 @@ def test_register_schema_models_calls_register_schema_model(monkeypatch: pytest.
 def test_register_response_schema_model_uses_serialized_field_names():
     from controllers.common.schema import register_response_schema_model
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_response_schema_model(namespace, ResponseAliasModel)
 
-    model_name, schema = namespace.schema_model.call_args.args
-
-    assert model_name == "ResponseAliasModel"
+    assert list(namespace.models) == ["ResponseAliasModel"]
+    schema = namespace.models["ResponseAliasModel"].__schema__
     assert "public_name" in schema["properties"]
     assert "internal_name" not in schema["properties"]
 
@@ -185,11 +177,11 @@ def test_register_response_schema_model_uses_serialized_field_names():
 def test_register_schema_model_preserves_openapi_nullable_unions():
     from controllers.common.schema import register_schema_model
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_schema_model(namespace, NullableSchemaModel)
 
-    called_schemas = {call.args[0]: call.args[1] for call in namespace.schema_model.call_args_list}
+    called_schemas = {name: model.__schema__ for name, model in namespace.models.items()}
     properties = called_schemas["NullableSchemaModel"]["properties"]
 
     assert properties["name"]["anyOf"] == [{"type": "string"}, {"type": "null"}]
@@ -198,80 +190,70 @@ def test_register_schema_model_preserves_openapi_nullable_unions():
     assert "anyOf" in properties["ambiguous"]
 
 
-def test_get_or_create_model_returns_existing_model(mock_console_ns):
+def test_get_or_create_model_returns_existing_model(console_ns):
     from controllers.common.schema import get_or_create_model
 
-    existing_model = MagicMock()
-    mock_console_ns.models = {"TestModel": existing_model}
+    existing_model = console_ns.model("TestModel", {"name": fields.String})
 
     result = get_or_create_model("TestModel", {"key": "value"})
 
-    assert result == existing_model
-    mock_console_ns.model.assert_not_called()
+    assert result is existing_model
+    assert len(console_ns.models) == 1
 
 
-def test_get_or_create_model_creates_new_model_when_not_exists(mock_console_ns):
+def test_get_or_create_model_creates_new_model_when_not_exists(console_ns):
     from controllers.common.schema import get_or_create_model
 
-    mock_console_ns.models = {}
-    new_model = MagicMock()
-    mock_console_ns.model.return_value = new_model
-    field_def = {"name": {"type": "string"}}
+    field_def = {"name": fields.String}
 
     result = get_or_create_model("NewModel", field_def)
 
-    assert result == new_model
-    mock_console_ns.model.assert_called_once_with("NewModel", field_def)
+    assert result is console_ns.models["NewModel"]
+    assert result.__schema__["properties"] == {"name": {"type": "string"}}
 
 
-def test_get_or_create_model_does_not_call_model_if_exists(mock_console_ns):
+def test_get_or_create_model_does_not_call_model_if_exists(console_ns):
     from controllers.common.schema import get_or_create_model
 
-    existing_model = MagicMock()
-    mock_console_ns.models = {"ExistingModel": existing_model}
+    existing_model = console_ns.model("ExistingModel", {"name": fields.String})
 
     result = get_or_create_model("ExistingModel", {"key": "value"})
 
-    assert result == existing_model
-    mock_console_ns.model.assert_not_called()
+    assert result is existing_model
+    assert len(console_ns.models) == 1
 
 
 def test_register_enum_models_registers_single_enum():
     from controllers.common.schema import register_enum_models
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_enum_models(namespace, StatusEnum)
 
-    namespace.schema_model.assert_called_once()
-
-    model_name, schema = namespace.schema_model.call_args.args
-
-    assert model_name == "StatusEnum"
+    assert list(namespace.models) == ["StatusEnum"]
+    schema = namespace.models["StatusEnum"].__schema__
     assert isinstance(schema, dict)
 
 
 def test_register_enum_models_registers_multiple_enums():
     from controllers.common.schema import register_enum_models
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_enum_models(namespace, StatusEnum, PriorityEnum)
 
-    assert namespace.schema_model.call_count == 2
-
-    called_names = [call.args[0] for call in namespace.schema_model.call_args_list]
+    called_names = list(namespace.models)
     assert called_names == ["StatusEnum", "PriorityEnum"]
 
 
 def test_register_enum_models_uses_correct_ref_template():
     from controllers.common.schema import register_enum_models
 
-    namespace = MagicMock(spec=Namespace)
+    namespace = Namespace("test")
 
     register_enum_models(namespace, StatusEnum)
 
-    schema = namespace.schema_model.call_args.args[1]
+    schema = namespace.models["StatusEnum"].__schema__
 
     # Verify the schema contains enum values
     assert "enum" in schema or "anyOf" in schema

api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py

@@ -13,7 +13,7 @@ from core.app.app_config.entities import (
     PromptTemplateEntity,
 )
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
-from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom, ModelConfigWithCredentialsEntity
+from core.app.entities.app_invoke_entities import ChatAppGenerateEntity, InvokeFrom
 from core.app.entities.queue_entities import (
     AppQueueEvent,
     MessageQueueMessage,
@@ -37,20 +37,13 @@ from core.app.entities.task_entities import (
 )
 from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
 from core.base.tts import AppGeneratorTTSPublisher
-from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
-from core.entities.provider_entities import CustomConfiguration, SystemConfiguration
 from core.ops.ops_trace_manager import TraceQueueManager
-from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime
-from graphon.model_runtime.entities.common_entities import I18nObject
 from graphon.model_runtime.entities.llm_entities import LLMResult as RuntimeLLMResult
 from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage
 from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, TextPromptMessageContent
-from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
-from graphon.model_runtime.entities.provider_entities import ProviderEntity
-from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
 from models.enums import ConversationFromSource
 from models.model import AppMode, Conversation, Message
-from models.provider import ProviderType
+from tests.unit_tests.core.model_fixtures import make_model_config
 
 
 class _QueueManager(AppQueueManager):
@@ -86,42 +79,6 @@ def queue_message(event: AppQueueEvent) -> MessageQueueMessage:
     )
 
 
-def model_config() -> ModelConfigWithCredentialsEntity:
-    provider = ProviderEntity(
-        provider="test-provider",
-        label=I18nObject(en_US="Test"),
-        supported_model_types=[ModelType.LLM],
-        configurate_methods=[],
-    )
-    return ModelConfigWithCredentialsEntity(
-        provider=provider.provider,
-        model="test-model",
-        mode="chat",
-        model_schema=AIModelEntity(
-            model="test-model",
-            label=I18nObject(en_US="Test"),
-            model_type=ModelType.LLM,
-            fetch_from=FetchFrom.PREDEFINED_MODEL,
-            model_properties={},
-        ),
-        provider_model_bundle=ProviderModelBundle(
-            configuration=ProviderConfiguration(
-                tenant_id="test-tenant-id",
-                provider=provider,
-                preferred_provider_type=ProviderType.CUSTOM,
-                using_provider_type=ProviderType.CUSTOM,
-                system_configuration=SystemConfiguration(enabled=False),
-                custom_configuration=CustomConfiguration(provider=None),
-                model_settings=[],
-            ),
-            model_type_instance=LargeLanguageModel(
-                provider_schema=provider,
-                model_runtime=create_plugin_model_runtime(tenant_id="test-tenant-id"),
-            ),
-        ),
-    )
-
-
 def llm_chunk(content: str) -> LLMResultChunk:
     return LLMResultChunk(
         model="test-model",
@@ -165,7 +122,7 @@ class TestEasyUIBasedGenerateTaskPipelineProcessStreamResponse:
                 model=ModelConfigEntity(provider="test-provider", model="test-model"),
                 prompt_template=PromptTemplateEntity(prompt_type=PromptTemplateEntity.PromptType.SIMPLE),
             ),
-            model_conf=model_config(),
+            model_conf=make_model_config(provider="test-provider", model="test-model", mode="chat"),
             inputs={},
             files=[],
             user_id="test-user-id",

api/tests/unit_tests/core/helper/test_creators.py

@@ -1,6 +1,6 @@
 """Tests for the Creators Platform helper module."""
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import httpx
 import pytest
@@ -19,10 +19,11 @@ def _patch_creators_url(monkeypatch: pytest.MonkeyPatch):
 class TestUploadDSL:
     @patch("core.helper.creators.httpx.post")
     def test_returns_claim_code(self, mock_post):
-        mock_response = MagicMock(spec=httpx.Response)
-        mock_response.json.return_value = {"data": {"claim_code": "abc123"}}
-        mock_response.raise_for_status = MagicMock()
-        mock_post.return_value = mock_response
+        mock_post.return_value = httpx.Response(
+            200,
+            request=httpx.Request("POST", "https://creators.example.com/api/v1/templates/anonymous-upload"),
+            json={"data": {"claim_code": "abc123"}},
+        )
 
         from core.helper.creators import upload_dsl
 
@@ -46,10 +47,11 @@ class TestUploadDSL:
 
     @patch("core.helper.creators.httpx.post")
     def test_raises_on_missing_claim_code(self, mock_post):
-        mock_response = MagicMock(spec=httpx.Response)
-        mock_response.json.return_value = {"data": {}}
-        mock_response.raise_for_status = MagicMock()
-        mock_post.return_value = mock_response
+        mock_post.return_value = httpx.Response(
+            200,
+            request=httpx.Request("POST", "https://creators.example.com/api/v1/templates/anonymous-upload"),
+            json={"data": {}},
+        )
 
         from core.helper.creators import upload_dsl
 
@@ -58,10 +60,11 @@ class TestUploadDSL:
 
     @patch("core.helper.creators.httpx.post")
     def test_raises_on_non_string_claim_code(self, mock_post):
-        mock_response = MagicMock(spec=httpx.Response)
-        mock_response.json.return_value = {"data": {"claim_code": 123}}
-        mock_response.raise_for_status = MagicMock()
-        mock_post.return_value = mock_response
+        mock_post.return_value = httpx.Response(
+            200,
+            request=httpx.Request("POST", "https://creators.example.com/api/v1/templates/anonymous-upload"),
+            json={"data": {"claim_code": 123}},
+        )
 
         from core.helper.creators import upload_dsl
 
@@ -70,18 +73,16 @@ class TestUploadDSL:
 
     @patch("core.helper.creators.httpx.post")
     def test_raises_on_http_error(self, mock_post):
-        mock_response = MagicMock(spec=httpx.Response)
-        mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
-            "Server Error",
-            request=MagicMock(),
-            response=MagicMock(),
+        mock_post.return_value = httpx.Response(
+            500,
+            request=httpx.Request("POST", "https://creators.example.com/api/v1/templates/anonymous-upload"),
         )
-        mock_post.return_value = mock_response
 
         from core.helper.creators import upload_dsl
 
-        with pytest.raises(httpx.HTTPStatusError):
+        with pytest.raises(httpx.HTTPStatusError) as caught:
             upload_dsl(b"app: demo")
+        assert caught.value.response is mock_post.return_value
 
 
 class TestGetRedirectUrl:

api/tests/unit_tests/core/helper/test_http_client_pooling.py

@@ -1,16 +1,28 @@
 from __future__ import annotations
 
-from unittest.mock import MagicMock
+from collections.abc import Iterator
 
 import httpx
+import pytest
 
 from core.helper.http_client_pooling import HttpClientPoolFactory
 
 
-def test_get_or_create_reuses_client_for_same_key() -> None:
+@pytest.fixture
+def http_clients() -> Iterator[tuple[httpx.Client, httpx.Client]]:
+    def reject_request(request: httpx.Request) -> httpx.Response:
+        raise AssertionError(f"Pooling must not send requests: {request.url}")
+
+    with (
+        httpx.Client(transport=httpx.MockTransport(reject_request)) as first,
+        httpx.Client(transport=httpx.MockTransport(reject_request)) as second,
+    ):
+        yield first, second
+
+
+def test_get_or_create_reuses_client_for_same_key(http_clients: tuple[httpx.Client, httpx.Client]) -> None:
     factory = HttpClientPoolFactory()
-    first_client = MagicMock(spec=httpx.Client)
-    second_client = MagicMock(spec=httpx.Client)
+    first_client, second_client = http_clients
     clients = [first_client, second_client]
 
     def _builder() -> httpx.Client:
@@ -20,22 +32,24 @@ def test_get_or_create_reuses_client_for_same_key() -> None:
     assert factory.get_or_create("shared", _builder) is first_client
 
 
-def test_get_or_create_creates_distinct_clients_for_distinct_keys() -> None:
+def test_get_or_create_creates_distinct_clients_for_distinct_keys(
+    http_clients: tuple[httpx.Client, httpx.Client],
+) -> None:
     factory = HttpClientPoolFactory()
-    client_a = MagicMock(spec=httpx.Client)
-    client_b = MagicMock(spec=httpx.Client)
+    client_a, client_b = http_clients
 
     assert factory.get_or_create("a", lambda: client_a) is client_a
     assert factory.get_or_create("b", lambda: client_b) is client_b
 
 
-def test_close_all_closes_pooled_clients_and_allows_recreate() -> None:
+def test_close_all_closes_pooled_clients_and_allows_recreate(http_clients: tuple[httpx.Client, httpx.Client]) -> None:
     factory = HttpClientPoolFactory()
-    first_client = MagicMock(spec=httpx.Client)
-    replacement_client = MagicMock(spec=httpx.Client)
+    first_client, replacement_client = http_clients
 
     assert factory.get_or_create("x", lambda: first_client) is first_client
+    assert not first_client.is_closed
     factory.close_all()
 
-    first_client.close.assert_called_once()
+    assert first_client.is_closed
+    assert not replacement_client.is_closed
     assert factory.get_or_create("x", lambda: replacement_client) is replacement_client

api/tests/unit_tests/core/mcp/auth/test_auth_flow.py

@@ -43,6 +43,15 @@ from core.mcp.types import (
     OAuthTokens,
     ProtectedResourceMetadata,
 )
+from tests.unit_tests.core.mcp.fixtures import make_provider
+
+
+@pytest.fixture
+def provider(monkeypatch: pytest.MonkeyPatch) -> MCPProviderEntity:
+    """Use real credential parsing with plaintext values at the decryption boundary."""
+    monkeypatch.setattr("core.entities.mcp_provider.encrypter.decrypt_token", lambda _tenant_id, token: token)
+    monkeypatch.setattr(MCPProviderEntity, "_decrypt_dict", lambda _self, data: data.copy())
+    return make_provider(credentials={})
 
 
 class TestPKCEGeneration:
@@ -600,22 +609,6 @@ class TestCallbackHandling:
 class TestAuthOrchestration:
     """Test the main auth orchestration function."""
 
-    @pytest.fixture
-    def mock_provider(self):
-        """Create a mock provider entity."""
-        provider = Mock(spec=MCPProviderEntity)
-        provider.id = "provider-id"
-        provider.tenant_id = "tenant-id"
-        provider.decrypt_server_url.return_value = "https://api.example.com"
-        provider.client_metadata = OAuthClientMetadata(
-            client_name="Dify",
-            redirect_uris=["https://redirect.example.com"],
-        )
-        provider.redirect_url = "https://redirect.example.com"
-        provider.retrieve_client_information.return_value = None
-        provider.retrieve_tokens.return_value = None
-        return provider
-
     @pytest.fixture
     def mock_service(self):
         """Create a mock MCP service."""
@@ -624,7 +617,7 @@ class TestAuthOrchestration:
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
     @patch("core.mcp.auth.auth_flow.register_client")
     @patch("core.mcp.auth.auth_flow.start_authorization")
-    def test_auth_new_registration(self, mock_start_auth, mock_register, mock_discover, mock_provider, mock_service):
+    def test_auth_new_registration(self, mock_start_auth, mock_register, mock_discover, provider, mock_service):
         """Test auth flow for new client registration."""
         # Setup
         mock_discover.return_value = (
@@ -644,7 +637,7 @@ class TestAuthOrchestration:
         )
         mock_start_auth.return_value = ("https://auth.example.com/authorize?...", "code-verifier")
 
-        result = auth(mock_provider)
+        result = auth(provider)
 
         # auth() now returns AuthResult
         assert isinstance(result, AuthResult)
@@ -670,7 +663,7 @@ class TestAuthOrchestration:
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
     @patch("core.mcp.auth.auth_flow._retrieve_redis_state")
     @patch("core.mcp.auth.auth_flow.exchange_authorization")
-    def test_auth_exchange_code(self, mock_exchange, mock_retrieve_state, mock_discover, mock_provider, mock_service):
+    def test_auth_exchange_code(self, mock_exchange, mock_retrieve_state, mock_discover, provider, mock_service):
         """Test auth flow for exchanging authorization code."""
         # Setup metadata discovery
         mock_discover.return_value = (
@@ -685,7 +678,7 @@ class TestAuthOrchestration:
         )
 
         # Setup existing client
-        mock_provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="existing-client")
+        provider.credentials["client_information"] = {"client_id": "existing-client"}
 
         # Setup state retrieval
         state_data = OAuthCallbackState(
@@ -703,7 +696,7 @@ class TestAuthOrchestration:
         tokens = OAuthTokens(access_token="new-token", token_type="Bearer", expires_in=3600)
         mock_exchange.return_value = tokens
 
-        result = auth(mock_provider, authorization_code="auth-code", state_param="state-key")
+        result = auth(provider, authorization_code="auth-code", state_param="state-key")
 
         # auth() now returns AuthResult, not a dict
         assert isinstance(result, AuthResult)
@@ -717,7 +710,7 @@ class TestAuthOrchestration:
         assert result.actions[0].tenant_id == "tenant-id"
 
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_exchange_code_without_state(self, mock_discover, mock_provider, mock_service):
+    def test_auth_exchange_code_without_state(self, mock_discover, provider, mock_service):
         """Test auth flow fails when exchanging code without state."""
         # Setup metadata discovery
         mock_discover.return_value = (
@@ -731,19 +724,19 @@ class TestAuthOrchestration:
             None,
         )
 
-        mock_provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="existing-client")
+        provider.credentials["client_information"] = {"client_id": "existing-client"}
 
         with pytest.raises(ValueError) as exc_info:
-            auth(mock_provider, authorization_code="auth-code")
+            auth(provider, authorization_code="auth-code")
 
         assert "State parameter is required" in str(exc_info.value)
 
     @patch("core.mcp.auth.auth_flow.refresh_authorization")
-    def test_auth_refresh_token(self, mock_refresh, mock_provider, mock_service):
+    def test_auth_refresh_token(self, mock_refresh, provider, mock_service):
         """Test auth flow for refreshing tokens."""
         # Setup existing client and tokens
-        mock_provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="existing-client")
-        mock_provider.retrieve_tokens.return_value = OAuthTokens(
+        provider.credentials["client_information"] = {"client_id": "existing-client"}
+        provider.credentials.update(
             access_token="old-token",
             token_type="Bearer",
             expires_in=0,
@@ -771,7 +764,7 @@ class TestAuthOrchestration:
                 None,
             )
 
-            result = auth(mock_provider)
+            result = auth(provider)
 
             # auth() now returns AuthResult
             assert isinstance(result, AuthResult)
@@ -788,7 +781,7 @@ class TestAuthOrchestration:
             mock_refresh.assert_called_once()
 
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_registration_fails_with_code(self, mock_discover, mock_provider, mock_service):
+    def test_auth_registration_fails_with_code(self, mock_discover, provider, mock_service):
         """Test auth fails when no client info exists but code is provided."""
         # Setup metadata discovery
         mock_discover.return_value = (
@@ -802,10 +795,10 @@ class TestAuthOrchestration:
             None,
         )
 
-        mock_provider.retrieve_client_information.return_value = None
+        provider.credentials.pop("client_information", None)
 
         with pytest.raises(ValueError) as exc_info:
-            auth(mock_provider, authorization_code="auth-code")
+            auth(provider, authorization_code="auth-code")
 
         assert "Existing OAuth client information is required" in str(exc_info.value)
 
@@ -999,15 +992,17 @@ class TestAuthOrchestration:
     @patch("core.mcp.auth.auth_flow._retrieve_redis_state")
     @patch("core.mcp.auth.auth_flow.exchange_authorization")
     def test_handle_callback(self, mock_exchange, mock_retrieve):
-        state = Mock(spec=OAuthCallbackState)
-        state.server_url = "https://api"
-        state.metadata = None
-        state.client_information = Mock()
-        state.code_verifier = "cv"
-        state.redirect_uri = "https://re"
+        state = OAuthCallbackState(
+            provider_id="provider-id",
+            tenant_id="tenant-id",
+            server_url="https://api",
+            client_information=OAuthClientInformation(client_id="c1"),
+            code_verifier="cv",
+            redirect_uri="https://re",
+        )
         mock_retrieve.return_value = state
 
-        tokens = Mock(spec=OAuthTokens)
+        tokens = OAuthTokens(access_token="at", token_type="Bearer")
         mock_exchange.return_value = tokens
 
         s, t = handle_callback("key", "code")
@@ -1070,7 +1065,11 @@ class TestAuthOrchestration:
                 mock_prm.return_value = ProtectedResourceMetadata(
                     resource="https://api", authorization_servers=["https://auth"]
                 )
-                mock_asm.return_value = Mock(spec=OAuthMetadata)
+                mock_asm.return_value = OAuthMetadata(
+                    authorization_endpoint="https://auth/authorize",
+                    token_endpoint="https://auth/token",
+                    response_types_supported=["code"],
+                )
 
                 asm, prm, hint = discover_oauth_metadata("https://api")
                 assert asm == mock_asm.return_value
@@ -1285,11 +1284,8 @@ class TestAuthOrchestration:
             register_client("https://api", None, client_metadata)
 
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_orchestration_failures(self, mock_discover):
-        provider = Mock(spec=MCPProviderEntity)
-        provider.decrypt_server_url.return_value = "https://api"
-        provider.id = "p1"
-        provider.tenant_id = "t1"
+    def test_auth_orchestration_failures(self, mock_discover, provider):
+        provider.server_url = "https://api"
 
         # Case 1: No server metadata
         mock_discover.return_value = (None, None, None)
@@ -1303,7 +1299,7 @@ class TestAuthOrchestration:
             response_types_supported=["code"],
         )
         mock_discover.return_value = (asm, None, None)
-        provider.retrieve_client_information.return_value = None
+        provider.credentials.pop("client_information", None)
         with pytest.raises(ValueError, match="Existing OAuth client information is required"):
             auth(provider, authorization_code="code")
 
@@ -1320,13 +1316,9 @@ class TestAuthOrchestration:
                 auth(provider)
 
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_orchestration_client_credentials(self, mock_discover):
-        provider = Mock(spec=MCPProviderEntity)
-        provider.decrypt_server_url.return_value = "https://api"
-        provider.id = "p1"
-        provider.tenant_id = "t1"
-        provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="c1", client_secret="s1")
-        provider.decrypt_credentials.return_value = {"scope": "read"}
+    def test_auth_orchestration_client_credentials(self, mock_discover, provider):
+        provider.server_url = "https://api"
+        provider.credentials = {"client_information": {"client_id": "c1", "client_secret": "s1"}, "scope": "read"}
 
         asm = OAuthMetadata(
             authorization_endpoint="https://auth/auth",
@@ -1350,13 +1342,9 @@ class TestAuthOrchestration:
                 auth(provider)
 
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_orchestration_authorization_code(self, mock_discover):
-        provider = Mock(spec=MCPProviderEntity)
-        provider.decrypt_server_url.return_value = "https://api"
-        provider.id = "p1"
-        provider.tenant_id = "t1"
-        provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="c1")
-        provider.decrypt_credentials.return_value = {}
+    def test_auth_orchestration_authorization_code(self, mock_discover, provider):
+        provider.server_url = "https://api"
+        provider.credentials = {"client_information": {"client_id": "c1"}}
 
         asm = OAuthMetadata(
             authorization_endpoint="https://auth/auth",
@@ -1368,9 +1356,14 @@ class TestAuthOrchestration:
 
         # Case 1: Exchange code
         with patch("core.mcp.auth.auth_flow._retrieve_redis_state") as mock_retrieve:
-            state = Mock(spec=OAuthCallbackState)
-            state.code_verifier = "cv"
-            state.redirect_uri = "https://re"
+            state = OAuthCallbackState(
+                provider_id=provider.id,
+                tenant_id=provider.tenant_id,
+                server_url=provider.server_url,
+                client_information=OAuthClientInformation(client_id="c1"),
+                code_verifier="cv",
+                redirect_uri="https://re",
+            )
             mock_retrieve.return_value = state
 
             with patch("core.mcp.auth.auth_flow.exchange_authorization") as mock_exchange:
@@ -1385,7 +1378,7 @@ class TestAuthOrchestration:
                     auth(provider, authorization_code="code")
 
                 # Missing verifier in state
-                state.code_verifier = None
+                state.code_verifier = ""
                 with pytest.raises(ValueError, match="Missing code_verifier"):
                     auth(provider, authorization_code="code", state_param="sp")
 
@@ -1395,14 +1388,10 @@ class TestAuthOrchestration:
                     auth(provider, authorization_code="code", state_param="sp")
 
     @patch("core.mcp.auth.auth_flow.discover_oauth_metadata")
-    def test_auth_orchestration_refresh_failure(self, mock_discover):
-        provider = Mock(spec=MCPProviderEntity)
-        provider.decrypt_server_url.return_value = "https://api"
-        provider.id = "p1"
-        provider.tenant_id = "t1"
-        provider.retrieve_client_information.return_value = OAuthClientInformation(client_id="c1")
-        provider.decrypt_credentials.return_value = {}
-        provider.retrieve_tokens.return_value = OAuthTokens(access_token="at", token_type="Bearer", refresh_token="rt")
+    def test_auth_orchestration_refresh_failure(self, mock_discover, provider):
+        provider.server_url = "https://api"
+        provider.credentials = {"client_information": {"client_id": "c1"}}
+        provider.credentials.update(access_token="at", token_type="Bearer", refresh_token="rt")
 
         asm = OAuthMetadata(
             authorization_endpoint="https://auth/auth",

api/tests/unit_tests/core/mcp/fixtures.py

@@ -0,0 +1,25 @@
+"""Real MCP provider fixtures for authentication tests."""
+
+from datetime import datetime
+
+from core.entities.mcp_provider import MCPProviderEntity
+
+
+def make_provider(credentials: dict[str, object]) -> MCPProviderEntity:
+    return MCPProviderEntity(
+        id="provider-id",
+        tenant_id="tenant-id",
+        user_id="user-id",
+        server_identifier="test-server",
+        name="Test server",
+        server_url="https://api.example.com",
+        headers={},
+        timeout=30,
+        sse_read_timeout=300,
+        authed=False,
+        credentials=credentials,
+        tools=[],
+        icon="",
+        created_at=datetime(2024, 1, 1),
+        updated_at=datetime(2024, 1, 1),
+    )

api/tests/unit_tests/core/mcp/session/test_base_session.py

@@ -556,17 +556,17 @@ def test_session_exit_timeout(streams):
     read_stream, write_stream = streams
     session = MockSession(read_stream, write_stream, ReceiveRequest, ReceiveNotification)
 
-    mock_future = MagicMock(spec=Future)
-    mock_future.result.side_effect = TimeoutError()
-    mock_future.done.return_value = False
+    future = Future()
+    session._receiver_future = future
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        session._executor = executor
+        with patch.object(executor, "shutdown", wraps=executor.shutdown) as shutdown:
+            session.__exit__(None, None, None)
 
-    session._receiver_future = mock_future
-    session._executor = MagicMock(spec=ThreadPoolExecutor)
-
-    session.__exit__(None, None, None)
-
-    mock_future.cancel.assert_called_once()
-    session._executor.shutdown.assert_called_once_with(wait=False)
+            assert future.cancelled()
+            shutdown.assert_called_once_with(wait=False)
+            with pytest.raises(RuntimeError, match="cannot schedule new futures after shutdown"):
+                executor.submit(lambda: None)
 
 
 @pytest.mark.timeout(10)

api/tests/unit_tests/core/mcp/session/test_client_session.py

@@ -531,16 +531,26 @@ def test_received_notification_logging(streams):
     assert logging_cb.call_args[0][0].level == "info"
 
 
-def test_default_message_handler():
+def test_default_message_handler(streams):
     # Exception case
     with pytest.raises(ValueError, match="test error"):
         _default_message_handler(Exception("test error"))
 
     # Notification case - should do nothing
-    _default_message_handler(MagicMock(spec=types.ServerNotification))
+    _default_message_handler(types.ServerNotification(root=types.ToolListChangedNotification()))
 
     # RequestResponder case - should do nothing
-    _default_message_handler(MagicMock(spec=RequestResponder))
+    read_stream, write_stream = streams
+    session = ClientSession(read_stream, write_stream)
+    with RequestResponder(
+        request_id=1,
+        request_meta=None,
+        request=types.ServerRequest(root=types.PingRequest()),
+        session=session,
+        on_complete=lambda _responder: None,
+    ) as responder:
+        _default_message_handler(responder)
+        assert not responder.completed
 
 
 def test_default_sampling_callback():

api/tests/unit_tests/core/mcp/test_mcp_client.py

@@ -14,7 +14,8 @@ from core.entities.mcp_provider import MCPProviderEntity
 from core.mcp.auth_client import MCPClientWithAuthRetry
 from core.mcp.error import MCPAuthError, MCPConnectionError
 from core.mcp.mcp_client import MCPClient
-from core.mcp.types import CallToolResult, ListToolsResult, OAuthTokens, TextContent, Tool, ToolAnnotations
+from core.mcp.types import CallToolResult, ListToolsResult, TextContent, Tool, ToolAnnotations
+from tests.unit_tests.core.mcp.fixtures import make_provider
 
 
 class TestMCPClient:
@@ -428,23 +429,25 @@ class TestMCPClient:
     def test_cleanup(self):
         """Test cleanup method."""
         client = MCPClient(server_url="http://test.example.com")
-        mock_exit_stack = Mock(spec=ExitStack)
-        client._exit_stack = mock_exit_stack
+        cleaned = []
+        client._exit_stack.callback(cleaned.append, "closed")
         client._session = Mock()
         client._initialized = True
 
         client.cleanup()
 
-        mock_exit_stack.close.assert_called_once()
+        assert cleaned == ["closed"]
         assert client._session is None
         assert client._initialized is False
 
     def test_cleanup_with_error(self):
         """Test cleanup method with error."""
         client = MCPClient(server_url="http://test.example.com")
-        mock_exit_stack = Mock(spec=ExitStack)
-        mock_exit_stack.close.side_effect = Exception("Cleanup error")
-        client._exit_stack = mock_exit_stack
+
+        def fail_cleanup():
+            raise RuntimeError("Cleanup error")
+
+        client._exit_stack.callback(fail_cleanup)
         client._session = Mock()
         client._initialized = True
 
@@ -530,43 +533,41 @@ class TestMCPClientWithAuthRetry:
     """Test suite for MCPClientWithAuthRetry."""
 
     @pytest.fixture
-    def mock_provider(self):
-        provider = MagicMock(spec=MCPProviderEntity)
-        provider.id = "test-provider-id"
-        provider.server_identifier = "test-server-identifier"
-        provider.tenant_id = "test-tenant-id"
-        provider.retrieve_tokens.return_value = OAuthTokens(
-            access_token="new-token",
-            token_type="Bearer",
-            expires_in=3600,
-            refresh_token="refresh-token",
+    def provider(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setattr(MCPProviderEntity, "_decrypt_dict", lambda _self, data: data.copy())
+        return make_provider(
+            credentials={
+                "access_token": "new-token",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+                "refresh_token": "refresh-token",
+            }
         )
-        return provider
 
     @pytest.fixture
-    def auth_client(self, mock_provider):
+    def auth_client(self, provider):
         client = MCPClientWithAuthRetry(
             server_url="http://test.example.com",
             headers={"Authorization": "Bearer old-token"},
-            provider_entity=mock_provider,
+            provider_entity=provider,
             authorization_code="test-code",
         )
         return client
 
-    def test_init(self, mock_provider):
+    def test_init(self, provider):
         """Test initialization."""
         client = MCPClientWithAuthRetry(
             server_url="http://test.example.com",
             headers={"Authorization": "Bearer test"},
             timeout=30.0,
-            provider_entity=mock_provider,
+            provider_entity=provider,
             authorization_code="initial-code",
         )
 
         assert client.server_url == "http://test.example.com"
         assert client.headers == {"Authorization": "Bearer test"}
         assert client.timeout == 30.0
-        assert client.provider_entity == mock_provider
+        assert client.provider_entity == provider
         assert client.authorization_code == "initial-code"
         assert client._has_retried is False
 
@@ -575,19 +576,20 @@ class TestMCPClientWithAuthRetry:
         self,
         mock_service_class,
         auth_client,
-        mock_provider,
+        provider,
         monkeypatch: pytest.MonkeyPatch,
         sqlite_engine: Engine,
     ):
         monkeypatch.setattr("core.mcp.auth_client.db", SimpleNamespace(engine=sqlite_engine))
 
         mock_service = mock_service_class.return_value
-        new_provider = MagicMock(spec=MCPProviderEntity)
-        new_provider.retrieve_tokens.return_value = OAuthTokens(
-            access_token="new-access-token",
-            token_type="Bearer",
-            expires_in=3600,
-            refresh_token="new-refresh-token",
+        new_provider = make_provider(
+            credentials={
+                "access_token": "new-access-token",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+                "refresh_token": "new-refresh-token",
+            }
         )
         mock_service.get_provider_entity_by_server_identifier.return_value = new_provider
 
@@ -603,13 +605,13 @@ class TestMCPClientWithAuthRetry:
 
         # Verify service calls - error.resource_metadata_url and error.scope_hint are parsed from header
         mock_service.auth_with_actions.assert_called_once_with(
-            mock_provider,
+            provider,
             "test-code",
             resource_metadata_url="http://meta",
             scope_hint="read",
         )
         mock_service.get_provider_entity_by_server_identifier.assert_called_once_with(
-            server_identifier=mock_provider.server_identifier, tenant_id=mock_provider.tenant_id
+            server_identifier=provider.server_identifier, tenant_id=provider.tenant_id
         )
 
         # Verify client updates
@@ -643,7 +645,7 @@ class TestMCPClientWithAuthRetry:
         self,
         mock_service_class,
         auth_client,
-        mock_provider,
+        provider,
         monkeypatch: pytest.MonkeyPatch,
         sqlite_engine: Engine,
     ):
@@ -651,8 +653,7 @@ class TestMCPClientWithAuthRetry:
         monkeypatch.setattr("core.mcp.auth_client.db", SimpleNamespace(engine=sqlite_engine))
         mock_service = mock_service_class.return_value
 
-        new_provider = MagicMock(spec=MCPProviderEntity)
-        new_provider.retrieve_tokens.return_value = None
+        new_provider = make_provider(credentials={})
         mock_service.get_provider_entity_by_server_identifier.return_value = new_provider
 
         error = MCPAuthError("Auth failed")

api/tests/unit_tests/core/model_fixtures.py

@@ -0,0 +1,61 @@
+"""Build real model configuration and memory objects without resolving provider credentials."""
+
+from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
+from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
+from core.entities.provider_entities import CustomConfiguration, SystemConfiguration
+from core.memory.token_buffer_memory import TokenBufferMemory
+from core.model_manager import ModelInstance
+from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime
+from graphon.model_runtime.entities.common_entities import I18nObject
+from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
+from graphon.model_runtime.entities.provider_entities import ProviderEntity
+from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
+from models.model import Conversation
+from models.provider import ProviderType
+
+
+def make_model_config(*, provider: str, model: str, mode: str) -> ModelConfigWithCredentialsEntity:
+    provider_schema = ProviderEntity(
+        provider=provider,
+        label=I18nObject(en_US="Test"),
+        supported_model_types=[ModelType.LLM],
+        configurate_methods=[],
+    )
+    return ModelConfigWithCredentialsEntity(
+        provider=provider_schema.provider,
+        model=model,
+        mode=mode,
+        model_schema=AIModelEntity(
+            model=model,
+            label=I18nObject(en_US="Test"),
+            model_type=ModelType.LLM,
+            fetch_from=FetchFrom.PREDEFINED_MODEL,
+            model_properties={},
+        ),
+        provider_model_bundle=ProviderModelBundle(
+            configuration=ProviderConfiguration(
+                tenant_id="test-tenant-id",
+                provider=provider_schema,
+                preferred_provider_type=ProviderType.CUSTOM,
+                using_provider_type=ProviderType.CUSTOM,
+                system_configuration=SystemConfiguration(enabled=False),
+                custom_configuration=CustomConfiguration(provider=None),
+                model_settings=[],
+            ),
+            model_type_instance=LargeLanguageModel(
+                provider_schema=provider_schema,
+                model_runtime=create_plugin_model_runtime(tenant_id="test-tenant-id"),
+            ),
+        ),
+    )
+
+
+def make_token_buffer_memory(config: ModelConfigWithCredentialsEntity) -> TokenBufferMemory:
+    return TokenBufferMemory(
+        conversation=Conversation(),
+        model_instance=ModelInstance(
+            provider_model_bundle=config.provider_model_bundle,
+            model=config.model,
+            credentials=config.credentials,
+        ),
+    )

api/tests/unit_tests/core/moderation/test_input_moderation.py

@@ -1,6 +1,8 @@
+import queue
 from unittest.mock import MagicMock, patch
 
 import pytest
+from flask import Flask
 
 from core.app.app_config.entities import AppConfig, SensitiveWordAvoidanceEntity
 from core.moderation.base import ModerationAction, ModerationError, ModerationInputsResult
@@ -13,10 +15,20 @@ from models.model import AppMode
 class TestInputModeration:
     @pytest.fixture
     def app_config(self):
-        config = MagicMock(spec=AppConfig)
-        config.app_mode = AppMode.CHAT
-        config.sensitive_word_avoidance = None
-        return config
+        return AppConfig(tenant_id="test_tenant_id", app_id="test_app_id", app_mode=AppMode.CHAT)
+
+    @pytest.fixture
+    def trace_manager(self):
+        with (
+            Flask(__name__).app_context(),
+            patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None),
+            patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True),
+            patch("core.ops.ops_trace_manager.trace_manager_queue", queue.Queue()),
+            patch.object(TraceQueueManager, "start_timer"),
+        ):
+            manager = TraceQueueManager(app_id="test_app_id")
+            with patch.object(manager, "add_trace_task", wraps=manager.add_trace_task):
+                yield manager
 
     @pytest.fixture
     def input_moderation(self):
@@ -46,9 +58,7 @@ class TestInputModeration:
         message_id = "test_message_id"
 
         # Setup config
-        sensitive_word_config = MagicMock(spec=SensitiveWordAvoidanceEntity)
-        sensitive_word_config.type = "keywords"
-        sensitive_word_config.config = {"keywords": ["bad"]}
+        sensitive_word_config = SensitiveWordAvoidanceEntity(type="keywords", config={"keywords": ["bad"]})
         app_config.sensitive_word_avoidance = sensitive_word_config
 
         # Setup factory mock
@@ -70,18 +80,17 @@ class TestInputModeration:
 
     @patch("core.moderation.input_moderation.ModerationFactory")
     @patch("core.moderation.input_moderation.TraceTask")
-    def test_check_with_trace_manager(self, mock_trace_task, mock_factory_cls, app_config, input_moderation):
+    def test_check_with_trace_manager(
+        self, mock_trace_task, mock_factory_cls, app_config, input_moderation, trace_manager
+    ):
         app_id = "test_app_id"
         tenant_id = "test_tenant_id"
         inputs = {"input_key": "input_value"}
         query = "test query"
         message_id = "test_message_id"
-        trace_manager = MagicMock(spec=TraceQueueManager)
 
         # Setup config
-        sensitive_word_config = MagicMock(spec=SensitiveWordAvoidanceEntity)
-        sensitive_word_config.type = "keywords"
-        sensitive_word_config.config = {}
+        sensitive_word_config = SensitiveWordAvoidanceEntity(type="keywords", config={})
         app_config.sensitive_word_avoidance = sensitive_word_config
 
         # Setup factory mock
@@ -100,6 +109,7 @@ class TestInputModeration:
         )
 
         trace_manager.add_trace_task.assert_called_once_with(mock_trace_task.return_value)
+        assert trace_manager.collect_tasks() == [mock_trace_task.return_value]
         mock_trace_task.assert_called_once()
         call_kwargs = mock_trace_task.call_args.kwargs
         call_args = mock_trace_task.call_args.args
@@ -118,9 +128,7 @@ class TestInputModeration:
         message_id = "test_message_id"
 
         # Setup config
-        sensitive_word_config = MagicMock(spec=SensitiveWordAvoidanceEntity)
-        sensitive_word_config.type = "keywords"
-        sensitive_word_config.config = {}
+        sensitive_word_config = SensitiveWordAvoidanceEntity(type="keywords", config={})
         app_config.sensitive_word_avoidance = sensitive_word_config
 
         # Setup factory mock
@@ -151,9 +159,7 @@ class TestInputModeration:
         message_id = "test_message_id"
 
         # Setup config
-        sensitive_word_config = MagicMock(spec=SensitiveWordAvoidanceEntity)
-        sensitive_word_config.type = "keywords"
-        sensitive_word_config.config = {}
+        sensitive_word_config = SensitiveWordAvoidanceEntity(type="keywords", config={})
         app_config.sensitive_word_avoidance = sensitive_word_config
 
         # Setup factory mock
@@ -183,9 +189,7 @@ class TestInputModeration:
         message_id = "test_message_id"
 
         # Setup config
-        sensitive_word_config = MagicMock(spec=SensitiveWordAvoidanceEntity)
-        sensitive_word_config.type = "keywords"
-        sensitive_word_config.config = {}
+        sensitive_word_config = SensitiveWordAvoidanceEntity(type="keywords", config={})
         app_config.sensitive_word_avoidance = sensitive_word_config
 
         # Setup factory mock

api/tests/unit_tests/core/moderation/test_output_moderation.py

@@ -3,7 +3,9 @@ from unittest.mock import MagicMock, patch
 import pytest
 from flask import Flask
 
-from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
+from core.app.apps.base_app_queue_manager import PublishFrom
+from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
+from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.entities.queue_entities import QueueMessageReplaceEvent
 from core.moderation.base import ModerationAction, ModerationOutputsResult
 from core.moderation.output_moderation import ModerationRule, OutputModeration
@@ -11,17 +13,27 @@ from core.moderation.output_moderation import ModerationRule, OutputModeration
 
 class TestOutputModeration:
     @pytest.fixture
-    def mock_queue_manager(self):
-        return MagicMock(spec=AppQueueManager)
+    def queue_manager(self):
+        with patch("core.app.apps.base_app_queue_manager.redis_client.setex"):
+            manager = MessageBasedAppQueueManager(
+                task_id="test-task",
+                user_id="test-user",
+                invoke_from=InvokeFrom.SERVICE_API,
+                conversation_id="test-conversation",
+                app_mode="chat",
+                message_id="test-message",
+            )
+        manager.publish = MagicMock(wraps=manager.publish)
+        return manager
 
     @pytest.fixture
     def moderation_rule(self):
         return ModerationRule(type="keywords", config={"keywords": "badword"})
 
     @pytest.fixture
-    def output_moderation(self, mock_queue_manager, moderation_rule: ModerationRule):
+    def output_moderation(self, queue_manager, moderation_rule: ModerationRule):
         return OutputModeration(
-            tenant_id="test_tenant", app_id="test_app", rule=moderation_rule, queue_manager=mock_queue_manager
+            tenant_id="test_tenant", app_id="test_app", rule=moderation_rule, queue_manager=queue_manager
         )
 
     def test_should_direct_output(self, output_moderation: OutputModeration):
@@ -55,7 +67,7 @@ class TestOutputModeration:
             assert flagged is False
             assert output_moderation.is_final_chunk is True
 
-    def test_moderation_completion_flagged_direct_output(self, output_moderation: OutputModeration, mock_queue_manager):
+    def test_moderation_completion_flagged_direct_output(self, output_moderation: OutputModeration, queue_manager):
         with patch.object(OutputModeration, "moderation") as mock_moderation:
             mock_moderation.return_value = ModerationOutputsResult(
                 flagged=True, action=ModerationAction.DIRECT_OUTPUT, preset_response="preset"
@@ -65,13 +77,13 @@ class TestOutputModeration:
 
             assert output == "preset"
             assert flagged is True
-            mock_queue_manager.publish.assert_called_once()
-            args, _ = mock_queue_manager.publish.call_args
+            queue_manager.publish.assert_called_once()
+            args, _ = queue_manager.publish.call_args
             assert isinstance(args[0], QueueMessageReplaceEvent)
             assert args[0].text == "preset"
             assert args[1] == PublishFrom.TASK_PIPELINE
 
-    def test_moderation_completion_flagged_overridden(self, output_moderation: OutputModeration, mock_queue_manager):
+    def test_moderation_completion_flagged_overridden(self, output_moderation: OutputModeration, queue_manager):
         with patch.object(OutputModeration, "moderation") as mock_moderation:
             mock_moderation.return_value = ModerationOutputsResult(
                 flagged=True, action=ModerationAction.OVERRIDDEN, text="masked content"
@@ -81,14 +93,13 @@ class TestOutputModeration:
 
             assert output == "masked content"
             assert flagged is True
-            mock_queue_manager.publish.assert_called_once()
-            args, _ = mock_queue_manager.publish.call_args
+            queue_manager.publish.assert_called_once()
+            args, _ = queue_manager.publish.call_args
             assert args[0].text == "masked content"
 
     def test_start_thread(self, output_moderation: OutputModeration):
-        mock_app = MagicMock(spec=Flask)
-        with patch("core.moderation.output_moderation.current_app") as mock_current_app:
-            mock_current_app._get_current_object = MagicMock(return_value=mock_app)
+        app = Flask(__name__)
+        with app.app_context():
             with patch("threading.Thread") as mock_thread_class:
                 mock_thread_instance = MagicMock()
                 mock_thread_class.return_value = mock_thread_instance
@@ -132,16 +143,16 @@ class TestOutputModeration:
         result = output_moderation.moderation("tenant", "app", "buffer")
         assert result is None
 
-    def test_worker_loop_and_exit(self, output_moderation: OutputModeration, mock_queue_manager):
-        mock_app = MagicMock(spec=Flask)
+    def test_worker_loop_and_exit(self, output_moderation: OutputModeration, queue_manager):
+        app = Flask(__name__)
 
         # Test exit on thread_running=False
         output_moderation.thread_running = False
-        output_moderation.worker(mock_app, 10)
+        output_moderation.worker(app, 10)
         # Should exit immediately
 
     def test_worker_no_flag(self, output_moderation: OutputModeration):
-        mock_app = MagicMock(spec=Flask)
+        app = Flask(__name__)
 
         with patch.object(OutputModeration, "moderation") as mock_moderation:
             mock_moderation.return_value = ModerationOutputsResult(flagged=False, action=ModerationAction.DIRECT_OUTPUT)
@@ -156,12 +167,12 @@ class TestOutputModeration:
 
             mock_moderation.side_effect = side_effect
 
-            output_moderation.worker(mock_app, 10)
+            output_moderation.worker(app, 10)
 
             assert mock_moderation.called
 
-    def test_worker_flagged_direct_output(self, output_moderation: OutputModeration, mock_queue_manager):
-        mock_app = MagicMock(spec=Flask)
+    def test_worker_flagged_direct_output(self, output_moderation: OutputModeration, queue_manager):
+        app = Flask(__name__)
 
         with patch.object(OutputModeration, "moderation") as mock_moderation:
             mock_moderation.return_value = ModerationOutputsResult(
@@ -171,14 +182,14 @@ class TestOutputModeration:
             output_moderation.buffer = "badword"
             output_moderation.is_final_chunk = True
 
-            output_moderation.worker(mock_app, 10)
+            output_moderation.worker(app, 10)
 
             assert output_moderation.final_output == "preset"
-            mock_queue_manager.publish.assert_called_once()
+            queue_manager.publish.assert_called_once()
             # It breaks on DIRECT_OUTPUT
 
-    def test_worker_flagged_overridden(self, output_moderation: OutputModeration, mock_queue_manager):
-        mock_app = MagicMock(spec=Flask)
+    def test_worker_flagged_overridden(self, output_moderation: OutputModeration, queue_manager):
+        app = Flask(__name__)
 
         with patch.object(OutputModeration, "moderation") as mock_moderation:
             # Use side_effect to change thread_running on second call
@@ -193,14 +204,14 @@ class TestOutputModeration:
             output_moderation.buffer = "badword"
             output_moderation.is_final_chunk = True
 
-            output_moderation.worker(mock_app, 10)
+            output_moderation.worker(app, 10)
 
-            mock_queue_manager.publish.assert_called_once()
-            args, _ = mock_queue_manager.publish.call_args
+            queue_manager.publish.assert_called_once()
+            args, _ = queue_manager.publish.call_args
             assert args[0].text == "masked"
 
     def test_worker_chunk_too_small(self, output_moderation: OutputModeration):
-        mock_app = MagicMock(spec=Flask)
+        app = Flask(__name__)
         with patch("time.sleep") as mock_sleep:
             # chunk_length < buffer_size and not is_final_chunk
             output_moderation.buffer = "123"  # length 3
@@ -211,12 +222,12 @@ class TestOutputModeration:
 
             mock_sleep.side_effect = sleep_side_effect
 
-            output_moderation.worker(mock_app, 10)  # buffer_size 10
+            output_moderation.worker(app, 10)  # buffer_size 10
 
             mock_sleep.assert_called_once_with(1)
 
-    def test_worker_empty_not_flagged(self, output_moderation: OutputModeration, mock_queue_manager):
-        mock_app = MagicMock(spec=Flask)
+    def test_worker_empty_not_flagged(self, output_moderation: OutputModeration, queue_manager):
+        app = Flask(__name__)
         with patch.object(OutputModeration, "moderation") as mock_moderation:
             # Return None (exception or no rule)
             mock_moderation.return_value = None
@@ -229,6 +240,6 @@ class TestOutputModeration:
             output_moderation.buffer = "something"
             output_moderation.is_final_chunk = True
 
-            output_moderation.worker(mock_app, 10)
+            output_moderation.worker(app, 10)
 
-            mock_queue_manager.publish.assert_not_called()
+            queue_manager.publish.assert_not_called()

api/tests/unit_tests/core/ops/test_trace_queue_manager.py

@@ -11,58 +11,19 @@ When BOTH are false, tasks are silently dropped (correct behavior).
 """
 
 import queue
-import sys
-import types
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 
 @pytest.fixture
-def trace_queue_manager_and_task(monkeypatch: pytest.MonkeyPatch):
-    """Fixture to provide TraceQueueManager and TraceTask with delayed imports."""
-    module_name = "core.ops.ops_trace_manager"
-    if module_name not in sys.modules:
-        ops_stub = types.ModuleType(module_name)
-
-        class StubTraceTask:
-            def __init__(self, trace_type):
-                self.trace_type = trace_type
-                self.app_id = None
-
-        class StubTraceQueueManager:
-            def __init__(self, app_id=None):
-                self.app_id = app_id
-                from core.telemetry.gateway import is_enterprise_telemetry_enabled
-
-                self._enterprise_telemetry_enabled = is_enterprise_telemetry_enabled()
-                self.trace_instance = StubOpsTraceManager.get_ops_trace_instance(app_id)
-
-            def add_trace_task(self, trace_task):
-                if self._enterprise_telemetry_enabled or self.trace_instance:
-                    trace_task.app_id = self.app_id
-                    from core.ops.ops_trace_manager import trace_manager_queue
-
-                    trace_manager_queue.put(trace_task)
-
-        class StubOpsTraceManager:
-            @staticmethod
-            def get_ops_trace_instance(app_id):
-                return None
-
-        ops_stub.TraceQueueManager = StubTraceQueueManager
-        ops_stub.TraceTask = StubTraceTask
-        ops_stub.OpsTraceManager = StubOpsTraceManager
-        ops_stub.trace_manager_queue = MagicMock(spec=queue.Queue)
-        monkeypatch.setitem(sys.modules, module_name, ops_stub)
-
+def trace_queue_manager_and_task():
+    """Exercise the production guard without starting background dispatch."""
     from core.ops.entities.trace_entity import TraceTaskName
+    from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
 
-    ops_module = __import__(module_name, fromlist=["TraceQueueManager", "TraceTask"])
-    TraceQueueManager = ops_module.TraceQueueManager
-    TraceTask = ops_module.TraceTask
-
-    return TraceQueueManager, TraceTask, TraceTaskName
+    with patch.object(TraceQueueManager, "start_timer"):
+        yield TraceQueueManager, TraceTask, TraceTaskName
 
 
 class TestTraceQueueManagerTelemetryGuard:
@@ -76,19 +37,19 @@ class TestTraceQueueManagerTelemetryGuard:
         """
         TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
 
-        mock_queue = MagicMock(spec=queue.Queue)
+        trace_queue = queue.Queue()
 
         trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
 
         with (
             patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False),
             patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None),
-            patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue),
+            patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue),
         ):
             manager = TraceQueueManager(app_id="test-app-id")
             manager.add_trace_task(trace_task)
 
-            mock_queue.put.assert_not_called()
+            assert trace_queue.empty()
 
     def test_task_enqueued_when_telemetry_enabled(self, trace_queue_manager_and_task):
         """Verify task IS enqueued when enterprise telemetry is enabled.
@@ -98,20 +59,21 @@ class TestTraceQueueManagerTelemetryGuard:
         """
         TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
 
-        mock_queue = MagicMock(spec=queue.Queue)
+        trace_queue = queue.Queue()
 
         trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
 
         with (
             patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True),
             patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None),
-            patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue),
+            patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue),
         ):
             manager = TraceQueueManager(app_id="test-app-id")
             manager.add_trace_task(trace_task)
 
-            mock_queue.put.assert_called_once()
-            called_task = mock_queue.put.call_args[0][0]
+            called_task = trace_queue.get_nowait()
+            assert called_task is trace_task
+            assert trace_queue.empty()
             assert called_task.app_id == "test-app-id"
 
     def test_task_enqueued_when_trace_instance_configured(self, trace_queue_manager_and_task):
@@ -122,7 +84,7 @@ class TestTraceQueueManagerTelemetryGuard:
         """
         TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
 
-        mock_queue = MagicMock(spec=queue.Queue)
+        trace_queue = queue.Queue()
 
         mock_trace_instance = MagicMock()
 
@@ -133,13 +95,14 @@ class TestTraceQueueManagerTelemetryGuard:
             patch(
                 "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=mock_trace_instance
             ),
-            patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue),
+            patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue),
         ):
             manager = TraceQueueManager(app_id="test-app-id")
             manager.add_trace_task(trace_task)
 
-            mock_queue.put.assert_called_once()
-            called_task = mock_queue.put.call_args[0][0]
+            called_task = trace_queue.get_nowait()
+            assert called_task is trace_task
+            assert trace_queue.empty()
             assert called_task.app_id == "test-app-id"
 
     def test_task_enqueued_when_both_telemetry_and_trace_instance_enabled(self, trace_queue_manager_and_task):
@@ -150,7 +113,7 @@ class TestTraceQueueManagerTelemetryGuard:
         """
         TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
 
-        mock_queue = MagicMock(spec=queue.Queue)
+        trace_queue = queue.Queue()
 
         mock_trace_instance = MagicMock()
 
@@ -161,13 +124,14 @@ class TestTraceQueueManagerTelemetryGuard:
             patch(
                 "core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=mock_trace_instance
             ),
-            patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue),
+            patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue),
         ):
             manager = TraceQueueManager(app_id="test-app-id")
             manager.add_trace_task(trace_task)
 
-            mock_queue.put.assert_called_once()
-            called_task = mock_queue.put.call_args[0][0]
+            called_task = trace_queue.get_nowait()
+            assert called_task is trace_task
+            assert trace_queue.empty()
             assert called_task.app_id == "test-app-id"
 
     def test_app_id_set_before_enqueue(self, trace_queue_manager_and_task):
@@ -178,17 +142,19 @@ class TestTraceQueueManagerTelemetryGuard:
         """
         TraceQueueManager, TraceTask, TraceTaskName = trace_queue_manager_and_task
 
-        mock_queue = MagicMock(spec=queue.Queue)
+        trace_queue = queue.Queue()
 
         trace_task = TraceTask(trace_type=TraceTaskName.WORKFLOW_TRACE)
 
         with (
             patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True),
             patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None),
-            patch("core.ops.ops_trace_manager.trace_manager_queue", mock_queue),
+            patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue),
         ):
             manager = TraceQueueManager(app_id="expected-app-id")
             manager.add_trace_task(trace_task)
 
-            called_task = mock_queue.put.call_args[0][0]
+            called_task = trace_queue.get_nowait()
+            assert called_task is trace_task
+            assert trace_queue.empty()
             assert called_task.app_id == "expected-app-id"

api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py

@@ -3,8 +3,6 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from core.app.app_config.entities import ModelConfigEntity
-from core.memory.token_buffer_memory import TokenBufferMemory
 from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
 from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
@@ -17,14 +15,12 @@ from graphon.model_runtime.entities.message_entities import (
     TextPromptMessageContent,
     UserPromptMessage,
 )
-from models.model import Conversation
 from tests.unit_tests.config_override import apply_config_overrides
+from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory
 
 
 def test__get_completion_model_prompt_messages():
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
-    model_config_mock.provider = "openai"
-    model_config_mock.model = "gpt-3.5-turbo-instruct"
+    model_config = make_model_config(provider="openai", model="gpt-3.5-turbo-instruct", mode="completion")
 
     prompt_template = "Context:\n{{#context#}}\n\nHistories:\n{{#histories#}}\n\nyou are {{name}}."
     prompt_template_config = CompletionModelPromptTemplate(text=prompt_template)
@@ -38,7 +34,7 @@ def test__get_completion_model_prompt_messages():
     files = []
     context = "I am superman."
 
-    memory = TokenBufferMemory(conversation=Conversation(), model_instance=model_config_mock)
+    memory = make_token_buffer_memory(model_config)
 
     history_prompt_messages = [UserPromptMessage(content="Hi"), AssistantPromptMessage(content="Hello")]
     memory.get_history_prompt_messages = MagicMock(return_value=history_prompt_messages)
@@ -53,7 +49,7 @@ def test__get_completion_model_prompt_messages():
         context=context,
         memory_config=memory_config,
         memory=memory,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
 
     assert len(prompt_messages) == 1
@@ -72,12 +68,12 @@ def test__get_completion_model_prompt_messages():
 
 
 def test__get_chat_model_prompt_messages(get_chat_model_args):
-    model_config_mock, memory_config, messages, inputs, context = get_chat_model_args
+    model_config, memory_config, messages, inputs, context = get_chat_model_args
 
     files = []
     query = "Hi2."
 
-    memory = TokenBufferMemory(conversation=Conversation(), model_instance=model_config_mock)
+    memory = make_token_buffer_memory(model_config)
 
     history_prompt_messages = [UserPromptMessage(content="Hi1."), AssistantPromptMessage(content="Hello1!")]
     memory.get_history_prompt_messages = MagicMock(return_value=history_prompt_messages)
@@ -92,7 +88,7 @@ def test__get_chat_model_prompt_messages(get_chat_model_args):
         context=context,
         memory_config=memory_config,
         memory=memory,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
 
     assert len(prompt_messages) == 6
@@ -104,7 +100,7 @@ def test__get_chat_model_prompt_messages(get_chat_model_args):
 
 
 def test__get_chat_model_prompt_messages_no_memory(get_chat_model_args):
-    model_config_mock, _, messages, inputs, context = get_chat_model_args
+    model_config, _, messages, inputs, context = get_chat_model_args
 
     files = []
 
@@ -118,7 +114,7 @@ def test__get_chat_model_prompt_messages_no_memory(get_chat_model_args):
         context=context,
         memory_config=None,
         memory=None,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
 
     assert len(prompt_messages) == 3
@@ -129,7 +125,7 @@ def test__get_chat_model_prompt_messages_no_memory(get_chat_model_args):
 
 
 def test__get_chat_model_prompt_messages_with_files_no_memory(get_chat_model_args, monkeypatch: pytest.MonkeyPatch):
-    model_config_mock, _, messages, inputs, context = get_chat_model_args
+    model_config, _, messages, inputs, context = get_chat_model_args
     apply_config_overrides(monkeypatch, MULTIMODAL_SEND_FORMAT="url")
 
     files = [
@@ -157,7 +153,7 @@ def test__get_chat_model_prompt_messages_with_files_no_memory(get_chat_model_arg
             context=context,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
 
     assert len(prompt_messages) == 4
@@ -172,9 +168,7 @@ def test__get_chat_model_prompt_messages_with_files_no_memory(get_chat_model_arg
 
 @pytest.fixture
 def get_chat_model_args():
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
-    model_config_mock.provider = "openai"
-    model_config_mock.model = "gpt-4"
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
 
     memory_config = MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False))
 
@@ -190,12 +184,12 @@ def get_chat_model_args():
 
     context = "I am superman."
 
-    return model_config_mock, memory_config, prompt_messages, inputs, context
+    return model_config, memory_config, prompt_messages, inputs, context
 
 
 def test_get_prompt_dispatches_completion_and_chat_and_invalid():
     transform = AdvancedPromptTransform()
-    model_config = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     completion_template = CompletionModelPromptTemplate(text="Hello {{name}}", edition_type="basic")
     chat_template = [ChatModelMessage(text="Hello {{name}}", role=PromptMessageRole.USER, edition_type="basic")]
 
@@ -240,7 +234,7 @@ def test_get_prompt_dispatches_completion_and_chat_and_invalid():
 
 
 def test_completion_prompt_jinja2_with_files():
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     transform = AdvancedPromptTransform()
     completion_template = CompletionModelPromptTemplate(text="Hi {{name}}", edition_type="jinja2")
 
@@ -268,7 +262,7 @@ def test_completion_prompt_jinja2_with_files():
             context=None,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
 
     assert len(messages) == 1
@@ -279,7 +273,7 @@ def test_completion_prompt_jinja2_with_files():
 
 
 def test_completion_prompt_basic_sets_query_variable():
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     transform = AdvancedPromptTransform()
     template = CompletionModelPromptTemplate(text="Q={{#query#}}", edition_type="basic")
 
@@ -291,7 +285,7 @@ def test_completion_prompt_basic_sets_query_variable():
         context=None,
         memory_config=None,
         memory=None,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
 
     assert messages[0].content == "Q=what?"
@@ -299,7 +293,7 @@ def test_completion_prompt_basic_sets_query_variable():
 
 def test_chat_prompt_with_variable_template_and_context():
     transform = AdvancedPromptTransform(with_variable_tmpl=True)
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     prompt_template = [ChatModelMessage(text="sys={{#node.name#}} ctx={{#context#}}", role=PromptMessageRole.SYSTEM)]
 
     messages = transform._get_chat_model_prompt_messages(
@@ -310,7 +304,7 @@ def test_chat_prompt_with_variable_template_and_context():
         context="context-text",
         memory_config=None,
         memory=None,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
 
     assert len(messages) == 1
@@ -320,7 +314,7 @@ def test_chat_prompt_with_variable_template_and_context():
 
 def test_chat_prompt_jinja2_branch_and_invalid_edition():
     transform = AdvancedPromptTransform()
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     prompt_template = [ChatModelMessage(text="Hello {{name}}", role=PromptMessageRole.USER, edition_type="jinja2")]
 
     with patch("core.prompt.advanced_prompt_transform.Jinja2Formatter.format", return_value="Hello John"):
@@ -332,7 +326,7 @@ def test_chat_prompt_jinja2_branch_and_invalid_edition():
             context=None,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
     assert messages[0].content == "Hello John"
 
@@ -346,13 +340,13 @@ def test_chat_prompt_jinja2_branch_and_invalid_edition():
             context=None,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
 
 
 def test_chat_prompt_query_template_and_query_only_branch():
     transform = AdvancedPromptTransform()
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     memory_config = MemoryConfig(
         window=MemoryConfig.WindowConfig(enabled=False),
         query_prompt_template="query={{#sys.query#}} ctx={{#context#}}",
@@ -367,16 +361,16 @@ def test_chat_prompt_query_template_and_query_only_branch():
         context="ctx",
         memory_config=memory_config,
         memory=None,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
     assert messages[-1].content == "query={{#sys.query#}} ctx=ctx"
 
 
 def test_chat_prompt_memory_with_files_and_query():
     transform = AdvancedPromptTransform()
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     memory_config = MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False))
-    memory = MagicMock(spec=TokenBufferMemory)
+    memory = make_token_buffer_memory(model_config)
     prompt_template = [ChatModelMessage(text="sys", role=PromptMessageRole.SYSTEM)]
     file = File(
         file_id="file1",
@@ -402,7 +396,7 @@ def test_chat_prompt_memory_with_files_and_query():
             context=None,
             memory_config=memory_config,
             memory=memory,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
 
     assert isinstance(messages[-1].content, list)
@@ -411,7 +405,7 @@ def test_chat_prompt_memory_with_files_and_query():
 
 def test_chat_prompt_files_without_query_updates_last_user_or_appends_new():
     transform = AdvancedPromptTransform()
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     file = File(
         file_id="file1",
         tenant_id="tenant1",
@@ -434,7 +428,7 @@ def test_chat_prompt_files_without_query_updates_last_user_or_appends_new():
             context=None,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
     assert isinstance(messages[-1].content, list)
     assert messages[-1].content[1].data == "u"
@@ -452,7 +446,7 @@ def test_chat_prompt_files_without_query_updates_last_user_or_appends_new():
             context=None,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
     assert isinstance(messages[-1], UserPromptMessage)
     assert isinstance(messages[-1].content, list)
@@ -461,7 +455,7 @@ def test_chat_prompt_files_without_query_updates_last_user_or_appends_new():
 
 def test_chat_prompt_files_with_query_branch():
     transform = AdvancedPromptTransform()
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     file = File(
         file_id="file1",
         tenant_id="tenant1",
@@ -483,7 +477,7 @@ def test_chat_prompt_files_with_query_branch():
             context=None,
             memory_config=None,
             memory=None,
-            model_config=model_config_mock,
+            model_config=model_config,
         )
 
     assert isinstance(messages[-1].content, list)
@@ -495,7 +489,7 @@ def test_set_context_query_histories_variable_helpers():
     parser_context = PromptTemplateParser(template="{{#context#}}")
     parser_query = PromptTemplateParser(template="{{#query#}}")
     parser_hist = PromptTemplateParser(template="{{#histories#}}")
-    model_config_mock = MagicMock(spec=ModelConfigEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     memory_config = MemoryConfig(
         role_prefix=MemoryConfig.RolePrefix(user="Human", assistant="Assistant"),
         window=MemoryConfig.WindowConfig(enabled=False),
@@ -512,7 +506,7 @@ def test_set_context_query_histories_variable_helpers():
             role_prefix=memory_config.role_prefix,  # type: ignore[arg-type]
             parser=parser_hist,
             prompt_inputs={},
-            model_config=model_config_mock,
+            model_config=model_config,
         )["#histories#"]
         == ""
     )

api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py

@@ -3,8 +3,6 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
-from core.memory.token_buffer_memory import TokenBufferMemory
 from core.prompt.prompt_templates.advanced_prompt_templates import (
     CHAT_APP_CHAT_PROMPT_CONFIG,
     CHAT_APP_COMPLETION_PROMPT_CONFIG,
@@ -19,7 +17,8 @@ from graphon.model_runtime.entities.message_entities import (
     TextPromptMessageContent,
     UserPromptMessage,
 )
-from models.model import AppMode, Conversation
+from models.model import AppMode
+from tests.unit_tests.core.model_fixtures import make_model_config, make_token_buffer_memory
 
 
 def test_get_common_chat_app_prompt_template_with_pcqm():
@@ -108,13 +107,11 @@ def test_get_common_chat_app_prompt_template_with_p():
 
 
 def test__get_chat_model_prompt_messages():
-    model_config_mock = MagicMock(spec=ModelConfigWithCredentialsEntity)
-    model_config_mock.provider = "openai"
-    model_config_mock.model = "gpt-4"
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
 
-    memory_mock = MagicMock(spec=TokenBufferMemory)
+    memory = make_token_buffer_memory(model_config)
     history_prompt_messages = [UserPromptMessage(content="Hi"), AssistantPromptMessage(content="Hello")]
-    memory_mock.get_history_prompt_messages.return_value = history_prompt_messages
+    memory.get_history_prompt_messages = MagicMock(return_value=history_prompt_messages)
 
     prompt_transform = SimplePromptTransform()
     prompt_transform._calculate_rest_token = MagicMock(return_value=2000)
@@ -130,8 +127,8 @@ def test__get_chat_model_prompt_messages():
         query=query,
         files=[],
         context=context,
-        memory=memory_mock,
-        model_config=model_config_mock,
+        memory=memory,
+        model_config=model_config,
     )
 
     prompt_template = prompt_transform.get_prompt_template(
@@ -153,11 +150,9 @@ def test__get_chat_model_prompt_messages():
 
 
 def test__get_completion_model_prompt_messages():
-    model_config_mock = MagicMock(spec=ModelConfigWithCredentialsEntity)
-    model_config_mock.provider = "openai"
-    model_config_mock.model = "gpt-3.5-turbo-instruct"
+    model_config = make_model_config(provider="openai", model="gpt-3.5-turbo-instruct", mode="completion")
 
-    memory = TokenBufferMemory(conversation=Conversation(), model_instance=model_config_mock)
+    memory = make_token_buffer_memory(model_config)
 
     history_prompt_messages = [UserPromptMessage(content="Hi"), AssistantPromptMessage(content="Hello")]
     memory.get_history_prompt_messages = MagicMock(return_value=history_prompt_messages)
@@ -176,7 +171,7 @@ def test__get_completion_model_prompt_messages():
         files=[],
         context=context,
         memory=memory,
-        model_config=model_config_mock,
+        model_config=model_config,
     )
 
     prompt_template = prompt_transform.get_prompt_template(
@@ -207,10 +202,8 @@ def test__get_completion_model_prompt_messages():
 
 def test_get_prompt_dispatches_chat_and_completion():
     transform = SimplePromptTransform()
-    model_config_chat = MagicMock(spec=ModelConfigWithCredentialsEntity)
-    model_config_chat.mode = "chat"
-    model_config_completion = MagicMock(spec=ModelConfigWithCredentialsEntity)
-    model_config_completion.mode = "completion"
+    model_config_chat = make_model_config(provider="openai", model="gpt-4", mode="chat")
+    model_config_completion = make_model_config(provider="openai", model="gpt-3.5-turbo-instruct", mode="completion")
     prompt_entity = SimpleNamespace(simple_prompt_template="hello")
 
     transform._get_chat_model_prompt_messages = MagicMock(return_value=(["chat-msg"], None))
@@ -292,7 +285,7 @@ def test_get_prompt_str_and_rules_type_validation_errors():
 
 def test_chat_model_prompt_messages_uses_prompt_when_query_empty():
     transform = SimplePromptTransform()
-    model_config = MagicMock(spec=ModelConfigWithCredentialsEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     transform._get_prompt_str_and_rules = MagicMock(return_value=("prompt-text", {}))
     transform._get_last_user_message = MagicMock(return_value=UserPromptMessage(content="prompt-text"))
 
@@ -313,7 +306,7 @@ def test_chat_model_prompt_messages_uses_prompt_when_query_empty():
 
 def test_completion_model_prompt_messages_empty_stops_becomes_none():
     transform = SimplePromptTransform()
-    model_config = MagicMock(spec=ModelConfigWithCredentialsEntity)
+    model_config = make_model_config(provider="openai", model="gpt-4", mode="chat")
     transform._get_prompt_str_and_rules = MagicMock(return_value=("prompt", {"stops": []}))
 
     prompt_messages, stops = transform._get_completion_model_prompt_messages(

api/tests/unit_tests/core/schemas/test_schema_manager.py

@@ -1,80 +1,78 @@
-from unittest.mock import MagicMock, patch
+from pathlib import Path
+from unittest.mock import patch
 
 from core.schemas.registry import SchemaRegistry
 from core.schemas.schema_manager import SchemaManager
 
 
-def test_init_with_provided_registry():
-    mock_registry = MagicMock(spec=SchemaRegistry)
-    manager = SchemaManager(registry=mock_registry)
-    assert manager.registry == mock_registry
+def test_init_with_provided_registry(tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
+    manager = SchemaManager(registry=registry)
+    assert manager.registry == registry
 
 
 @patch("core.schemas.schema_manager.SchemaRegistry.default_registry")
-def test_init_with_default_registry(mock_default_registry):
-    mock_registry = MagicMock(spec=SchemaRegistry)
-    mock_default_registry.return_value = mock_registry
+def test_init_with_default_registry(mock_default_registry, tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
+    mock_default_registry.return_value = registry
 
     manager = SchemaManager()
 
     mock_default_registry.assert_called_once()
-    assert manager.registry == mock_registry
+    assert manager.registry == registry
 
 
-def test_get_all_schema_definitions():
-    mock_registry = MagicMock(spec=SchemaRegistry)
-    expected_definitions = [{"name": "schema1", "schema": {}}, {"name": "schema2", "schema": {}}]
-    mock_registry.get_all_schemas_for_version.return_value = expected_definitions
+def test_get_all_schema_definitions(tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
+    expected_definitions = [
+        {"name": "schema1", "label": "schema1", "schema": {}},
+        {"name": "schema2", "label": "schema2", "schema": {}},
+    ]
+    registry.versions = {"v1": {"other": {}}, "v2": {"schema1": {}, "schema2": {}}}
 
-    manager = SchemaManager(registry=mock_registry)
+    manager = SchemaManager(registry=registry)
     result = manager.get_all_schema_definitions(version="v2")
 
-    mock_registry.get_all_schemas_for_version.assert_called_once_with("v2")
     assert result == expected_definitions
 
 
-def test_get_schema_by_name_success():
-    mock_registry = MagicMock(spec=SchemaRegistry)
+def test_get_schema_by_name_success(tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
     mock_schema = {"type": "object"}
-    mock_registry.get_schema.return_value = mock_schema
+    registry.versions = {"v1": {"my_schema": mock_schema}, "v2": {"my_schema": {"type": "string"}}}
 
-    manager = SchemaManager(registry=mock_registry)
+    manager = SchemaManager(registry=registry)
     result = manager.get_schema_by_name("my_schema", version="v1")
 
-    expected_uri = "https://dify.ai/schemas/v1/my_schema.json"
-    mock_registry.get_schema.assert_called_once_with(expected_uri)
     assert result == {"name": "my_schema", "schema": mock_schema}
 
 
-def test_get_schema_by_name_not_found():
-    mock_registry = MagicMock(spec=SchemaRegistry)
-    mock_registry.get_schema.return_value = None
+def test_get_schema_by_name_not_found(tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
 
-    manager = SchemaManager(registry=mock_registry)
+    manager = SchemaManager(registry=registry)
     result = manager.get_schema_by_name("non_existent", version="v1")
 
     assert result is None
 
 
-def test_list_available_schemas():
-    mock_registry = MagicMock(spec=SchemaRegistry)
+def test_list_available_schemas(tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
     expected_schemas = ["schema1", "schema2"]
-    mock_registry.list_schemas.return_value = expected_schemas
+    registry.versions = {"v1": {"schema2": {}, "schema1": {}}, "v2": {"other": {}}}
 
-    manager = SchemaManager(registry=mock_registry)
+    manager = SchemaManager(registry=registry)
     result = manager.list_available_schemas(version="v1")
 
-    mock_registry.list_schemas.assert_called_once_with("v1")
     assert result == expected_schemas
 
 
-def test_list_available_versions():
-    mock_registry = MagicMock(spec=SchemaRegistry)
+def test_list_available_versions(tmp_path: Path):
+    registry = SchemaRegistry(str(tmp_path))
     expected_versions = ["v1", "v2"]
-    mock_registry.list_versions.return_value = expected_versions
+    registry.versions = {"v2": {}, "v1": {}}
 
-    manager = SchemaManager(registry=mock_registry)
+    manager = SchemaManager(registry=registry)
     result = manager.list_available_versions()
 
-    mock_registry.list_versions.assert_called_once()
     assert result == expected_versions

api/tests/unit_tests/core/telemetry/test_facade.py

@@ -3,59 +3,32 @@
 from __future__ import annotations
 
 import queue
-import sys
-import types
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 from core.ops.entities.trace_entity import TraceTaskName
+from core.ops.ops_trace_manager import TraceQueueManager, TraceTask
+from core.telemetry import emit
 from core.telemetry.events import DraftNodeExecutionTraceEvent, TelemetryContext
 
 
 @pytest.fixture
-def telemetry_test_setup(monkeypatch: pytest.MonkeyPatch):
-    module_name = "core.ops.ops_trace_manager"
-    ops_stub = types.ModuleType(module_name)
-
-    class StubTraceTask:
-        def __init__(self, trace_type, **kwargs):
-            self.trace_type = trace_type
-            self.app_id = None
-            self.kwargs = kwargs
-
-    class StubTraceQueueManager:
-        def __init__(self, app_id=None, user_id=None):
-            self.app_id = app_id
-            self.user_id = user_id
-            self.trace_instance = StubOpsTraceManager.get_ops_trace_instance(app_id)
-
-        def add_trace_task(self, trace_task):
-            trace_task.app_id = self.app_id
-            from core.ops.ops_trace_manager import trace_manager_queue
-
-            trace_manager_queue.put(trace_task)
-
-    class StubOpsTraceManager:
-        @staticmethod
-        def get_ops_trace_instance(app_id):
-            return None
-
-    ops_stub.TraceQueueManager = StubTraceQueueManager
-    ops_stub.TraceTask = StubTraceTask
-    ops_stub.OpsTraceManager = StubOpsTraceManager
-    ops_stub.trace_manager_queue = MagicMock(spec=queue.Queue)
-    monkeypatch.setitem(sys.modules, module_name, ops_stub)
-
-    from core.telemetry import emit
-
-    return emit, ops_stub.trace_manager_queue
+def telemetry_test_setup():
+    trace_queue: queue.Queue[TraceTask] = queue.Queue()
+    with (
+        patch("core.ops.ops_trace_manager.OpsTraceManager.get_ops_trace_instance", return_value=None),
+        patch("core.ops.ops_trace_manager.trace_manager_queue", trace_queue),
+        patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=False),
+        patch.object(TraceQueueManager, "start_timer"),
+    ):
+        yield emit, trace_queue
 
 
 class TestTelemetryEmit:
     @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True)
     def test_emit_enterprise_trace_creates_trace_task(self, mock_ee, telemetry_test_setup):
-        emit_fn, mock_queue = telemetry_test_setup
+        emit_fn, trace_queue = telemetry_test_setup
 
         event = DraftNodeExecutionTraceEvent(
             context=TelemetryContext(
@@ -68,12 +41,14 @@ class TestTelemetryEmit:
 
         emit_fn(event)
 
-        mock_queue.put.assert_called_once()
-        called_task = mock_queue.put.call_args[0][0]
+        called_task = trace_queue.get_nowait()
+        assert trace_queue.empty()
+        assert called_task.app_id == "test-app"
+        assert called_task.user_id == "test-user"
         assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE
 
     def test_emit_enterprise_only_trace_dropped_when_ee_disabled(self, telemetry_test_setup):
-        emit_fn, mock_queue = telemetry_test_setup
+        emit_fn, trace_queue = telemetry_test_setup
 
         event = DraftNodeExecutionTraceEvent(
             context=TelemetryContext(
@@ -86,11 +61,11 @@ class TestTelemetryEmit:
 
         emit_fn(event)
 
-        mock_queue.put.assert_not_called()
+        assert trace_queue.empty()
 
     @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True)
     def test_emit_passes_name_directly_to_trace_task(self, mock_ee, telemetry_test_setup):
-        emit_fn, mock_queue = telemetry_test_setup
+        emit_fn, trace_queue = telemetry_test_setup
 
         event = DraftNodeExecutionTraceEvent(
             context=TelemetryContext(
@@ -103,14 +78,16 @@ class TestTelemetryEmit:
 
         emit_fn(event)
 
-        mock_queue.put.assert_called_once()
-        called_task = mock_queue.put.call_args[0][0]
+        called_task = trace_queue.get_nowait()
+        assert trace_queue.empty()
+        assert called_task.app_id == "test-app"
+        assert called_task.user_id == "test-user"
         assert called_task.trace_type == TraceTaskName.DRAFT_NODE_EXECUTION_TRACE
         assert isinstance(called_task.trace_type, TraceTaskName)
 
     @patch("core.telemetry.gateway.is_enterprise_telemetry_enabled", return_value=True)
     def test_emit_with_provided_trace_manager(self, mock_ee, telemetry_test_setup):
-        emit_fn, mock_queue = telemetry_test_setup
+        emit_fn, trace_queue = telemetry_test_setup
 
         mock_trace_manager = MagicMock()
         mock_trace_manager.add_trace_task = MagicMock()

api/tests/unit_tests/services/test_human_input_service.py

@@ -1,6 +1,9 @@
 import dataclasses
+import json
 import logging
+from collections.abc import Callable
 from datetime import datetime, timedelta
+from typing import cast
 from unittest.mock import MagicMock
 
 import pytest
@@ -17,6 +20,7 @@ from core.repositories.human_input_repository import (
     HumanInputFormRecord,
     HumanInputFormSubmissionRepository,
 )
+from core.workflow.human_input_adapter import DeliveryMethodType
 from core.workflow.nodes.human_input.entities import (
     FileInputConfig,
     FileListInputConfig,
@@ -30,7 +34,7 @@ from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInput
 from graphon.file import File, FileTransferMethod, FileType
 from graphon.runtime import GraphRuntimeState, VariablePool
 from libs.datetime_utils import naive_utc_now
-from models.human_input import RecipientType
+from models.human_input import HumanInputDelivery, HumanInputForm, HumanInputFormRecipient, RecipientType
 from models.model import App, AppMode
 from models.workflow import WorkflowRun
 from services.human_input_service import (
@@ -93,6 +97,66 @@ def sample_form_record() -> HumanInputFormRecord:
     )
 
 
+@pytest.fixture
+def form_repository(
+    sqlite_session_factory: sessionmaker[Session],
+    monkeypatch: pytest.MonkeyPatch,
+    mocker: MockerFixture,
+) -> Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository]:
+    monkeypatch.setattr(
+        "core.repositories.human_input_repository.session_factory.create_session", sqlite_session_factory
+    )
+
+    def build(record: HumanInputFormRecord | None) -> HumanInputFormSubmissionRepository:
+        """Persist the form and recipient; None leaves the database empty for missing-token tests."""
+        if record is not None:
+            assert record.recipient_id is not None
+            assert record.recipient_type is not None
+            assert record.access_token is not None
+            form = HumanInputForm(
+                id=record.form_id,
+                tenant_id=record.tenant_id,
+                app_id=record.app_id,
+                workflow_run_id=record.workflow_run_id,
+                conversation_id=record.conversation_id,
+                node_id=record.node_id,
+                form_kind=record.form_kind,
+                form_definition=record.definition.model_dump_json(),
+                rendered_content=record.rendered_content,
+                created_at=record.created_at,
+                expiration_time=record.expiration_time,
+                status=record.status,
+                selected_action_id=record.selected_action_id,
+                submitted_data=json.dumps(record.submitted_data) if record.submitted_data is not None else None,
+                submitted_at=record.submitted_at,
+                submission_user_id=record.submission_user_id,
+                submission_end_user_id=record.submission_end_user_id,
+                completed_by_recipient_id=record.completed_by_recipient_id,
+            )
+            delivery = HumanInputDelivery(
+                id="delivery-id",
+                form_id=record.form_id,
+                delivery_method_type=DeliveryMethodType.WEBAPP,
+                channel_payload="{}",
+            )
+            recipient = HumanInputFormRecipient(
+                id=record.recipient_id,
+                form_id=record.form_id,
+                delivery_id=delivery.id,
+                recipient_type=record.recipient_type,
+                recipient_payload=json.dumps({"TYPE": record.recipient_type}),
+                access_token=record.access_token,
+            )
+            with sqlite_session_factory.begin() as session:
+                session.add_all([form, delivery, recipient])
+        repository = HumanInputFormSubmissionRepository()
+        mocker.spy(repository, "get_by_token")
+        mocker.spy(repository, "mark_submitted")
+        return repository
+
+    return build
+
+
 def test_enqueue_resume_dispatches_task_for_workflow(
     mocker: MockerFixture,
     sqlite_session_factory: sessionmaker[Session],
@@ -190,17 +254,17 @@ def test_enqueue_resume_skips_unsupported_app_mode(
 
 
 def test_get_form_definition_by_token_for_console_uses_repository(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     console_record = dataclasses.replace(sample_form_record, recipient_type=RecipientType.CONSOLE)
-    repo.get_by_token.return_value = console_record
+    repo = form_repository(console_record)
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     form = service.get_form_definition_by_token_for_console("token")
 
-    repo.get_by_token.assert_called_once_with("token")
+    cast(MagicMock, repo.get_by_token).assert_called_once_with("token")
     assert form is not None
     assert form.get_definition() == console_record.definition
 
@@ -273,13 +337,12 @@ def test_resolve_form_inputs_uses_runtime_select_options(
 
 
 def test_submit_form_by_token_calls_repository_and_enqueue(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = sample_form_record
-    repo.mark_submitted.return_value = sample_form_record
+    repo = form_repository(sample_form_record)
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
@@ -291,32 +354,36 @@ def test_submit_form_by_token_calls_repository_and_enqueue(
         submission_end_user_id="end-user-id",
     )
 
-    repo.get_by_token.assert_called_once_with("token")
-    repo.mark_submitted.assert_called_once()
-    call_kwargs = repo.mark_submitted.call_args.kwargs
+    cast(MagicMock, repo.get_by_token).assert_called_once_with("token")
+    cast(MagicMock, repo.mark_submitted).assert_called_once()
+    call_kwargs = cast(MagicMock, repo.mark_submitted).call_args.kwargs
     assert call_kwargs["form_id"] == sample_form_record.form_id
     assert call_kwargs["recipient_id"] == sample_form_record.recipient_id
     assert call_kwargs["selected_action_id"] == "submit"
     assert call_kwargs["form_data"] == {"field": "value"}
     assert call_kwargs["submission_end_user_id"] == "end-user-id"
+    persisted = repo.get_by_form_id(sample_form_record.form_id)
+    assert persisted is not None
+    assert persisted.status == HumanInputFormStatus.SUBMITTED
+    assert persisted.submitted_data == {"field": "value"}
+    assert persisted.submission_end_user_id == "end-user-id"
     enqueue_spy.assert_called_once_with(sample_form_record.workflow_run_id)
 
 
 def test_submit_form_by_token_enqueues_agent_app_resume_for_conversation_form(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
     # ENG-635: a conversation-owned (Agent v2 chat) form routes to the chat
     # resume, not the workflow resume.
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     conversation_record = dataclasses.replace(
         sample_form_record,
         workflow_run_id=None,
         conversation_id="conv-1",
     )
-    repo.get_by_token.return_value = conversation_record
-    repo.mark_submitted.return_value = conversation_record
+    repo = form_repository(conversation_record)
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     workflow_enqueue_spy = mocker.patch.object(service, "enqueue_resume")
     chat_enqueue_spy = mocker.patch.object(service, "enqueue_agent_app_resume")
@@ -334,18 +401,17 @@ def test_submit_form_by_token_enqueues_agent_app_resume_for_conversation_form(
 
 
 def test_submit_form_by_token_skips_enqueue_for_delivery_test(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     test_record = dataclasses.replace(
         sample_form_record,
         form_kind=HumanInputFormKind.DELIVERY_TEST,
         workflow_run_id=None,
     )
-    repo.get_by_token.return_value = test_record
-    repo.mark_submitted.return_value = test_record
+    repo = form_repository(test_record)
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
@@ -360,13 +426,12 @@ def test_submit_form_by_token_skips_enqueue_for_delivery_test(
 
 
 def test_submit_form_by_token_passes_submission_user_id(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = sample_form_record
-    repo.mark_submitted.return_value = sample_form_record
+    repo = form_repository(sample_form_record)
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
 
@@ -378,18 +443,18 @@ def test_submit_form_by_token_passes_submission_user_id(
         submission_user_id="account-id",
     )
 
-    call_kwargs = repo.mark_submitted.call_args.kwargs
+    call_kwargs = cast(MagicMock, repo.mark_submitted).call_args.kwargs
     assert call_kwargs["submission_user_id"] == "account-id"
     assert call_kwargs["submission_end_user_id"] is None
     enqueue_spy.assert_called_once_with(sample_form_record.workflow_run_id)
 
 
 def test_submit_form_by_token_invalid_action(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = dataclasses.replace(sample_form_record)
+    repo = form_repository(dataclasses.replace(sample_form_record))
     service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError) as exc_info:
@@ -401,14 +466,14 @@ def test_submit_form_by_token_invalid_action(
         )
 
     assert "Invalid action" in str(exc_info.value)
-    repo.mark_submitted.assert_not_called()
+    cast(MagicMock, repo.mark_submitted).assert_not_called()
 
 
 def test_submit_form_by_token_missing_inputs(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
 
     definition_with_input = FormDefinition(
         form_content="hello",
@@ -418,7 +483,7 @@ def test_submit_form_by_token_missing_inputs(
         expiration_time=sample_form_record.expiration_time,
     )
     form_with_input = dataclasses.replace(sample_form_record, definition=definition_with_input)
-    repo.get_by_token.return_value = form_with_input
+    repo = form_repository(form_with_input)
     service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError) as exc_info:
@@ -430,7 +495,7 @@ def test_submit_form_by_token_missing_inputs(
         )
 
     assert "Missing required inputs" in str(exc_info.value)
-    repo.mark_submitted.assert_not_called()
+    cast(MagicMock, repo.mark_submitted).assert_not_called()
 
 
 @pytest.mark.parametrize(
@@ -482,13 +547,13 @@ def test_submit_form_by_token_missing_inputs(
     ],
 )
 def test_validate_human_input_submission_rejects_invalid_select_and_file_payloads(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     input_definition: dict[str, JsonValue],
     submitted_value: JsonValue,
     expected_message: str,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition.model_validate(
         {
             "form_content": "Validate form data",
@@ -498,7 +563,7 @@ def test_validate_human_input_submission_rejects_invalid_select_and_file_payload
             "expiration_time": naive_utc_now() + timedelta(hours=1),
         }
     )
-    repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
+    repo = form_repository(dataclasses.replace(sample_form_record, definition=definition))
     service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError) as exc_info:
@@ -510,7 +575,7 @@ def test_validate_human_input_submission_rejects_invalid_select_and_file_payload
         )
 
     assert expected_message in str(exc_info.value)
-    repo.mark_submitted.assert_not_called()
+    cast(MagicMock, repo.mark_submitted).assert_not_called()
 
 
 def test_form_properties(sample_form_record: HumanInputFormRecord) -> None:
@@ -540,20 +605,22 @@ def test_human_input_service_init_with_engine(sqlite_engine: Engine) -> None:
     assert service._session_factory.kw["bind"] is sqlite_engine
 
 
-def test_get_form_by_token_none(unbound_session_factory: sessionmaker[Session]) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = None
+def test_get_form_by_token_none(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
+    unbound_session_factory: sessionmaker[Session],
+) -> None:
+    repo = form_repository(None)
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     assert service.get_form_by_token("invalid") is None
 
 
 def test_get_form_definition_by_token_mismatch(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = sample_form_record
+    repo = form_repository(sample_form_record)
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     # RecipientType mismatch
@@ -561,11 +628,11 @@ def test_get_form_definition_by_token_mismatch(
 
 
 def test_get_form_definition_by_token_success(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = sample_form_record
+    repo = form_repository(sample_form_record)
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     form = service.get_form_definition_by_token(RecipientType.STANDALONE_WEB_APP, "token")
@@ -574,21 +641,21 @@ def test_get_form_definition_by_token_success(
 
 
 def test_get_form_definition_by_token_for_console_mismatch(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = sample_form_record  # is STANDALONE_WEB_APP
+    repo = form_repository(sample_form_record)  # is STANDALONE_WEB_APP
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     assert service.get_form_definition_by_token_for_console("token") is None
 
 
 def test_submit_form_by_token_delivery_not_enabled(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = None
+    repo = form_repository(None)
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     with pytest.raises(human_input_service_module.WebAppDeliveryNotEnabledError):
@@ -596,16 +663,13 @@ def test_submit_form_by_token_delivery_not_enabled(
 
 
 def test_submit_form_by_token_no_workflow_run_id(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
-    repo.get_by_token.return_value = sample_form_record
-
-    # Return record with no workflow_run_id
-    result_record = dataclasses.replace(sample_form_record, workflow_run_id=None)
-    repo.mark_submitted.return_value = result_record
+    # Preserve the legacy ownerless-record case: submission must not enqueue a workflow.
+    repo = form_repository(dataclasses.replace(sample_form_record, workflow_run_id=None))
 
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     enqueue_spy = mocker.patch.object(service, "enqueue_resume")
@@ -707,11 +771,11 @@ def test_is_globally_expired_zero_timeout(
 
 
 def test_submit_form_by_token_normalizes_select_and_files(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
         inputs=[
@@ -727,8 +791,7 @@ def test_submit_form_by_token_normalizes_select_and_files(
         expiration_time=sample_form_record.expiration_time,
     )
     form_with_inputs = dataclasses.replace(sample_form_record, definition=definition)
-    repo.get_by_token.return_value = form_with_inputs
-    repo.mark_submitted.return_value = form_with_inputs
+    repo = form_repository(form_with_inputs)
     service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     single_file = File(
@@ -781,7 +844,7 @@ def test_submit_form_by_token_normalizes_select_and_files(
         },
     )
 
-    submitted_data = repo.mark_submitted.call_args.kwargs["form_data"]
+    submitted_data = cast(MagicMock, repo.mark_submitted).call_args.kwargs["form_data"]
     assert submitted_data["decision"] == "approve"
     assert submitted_data["attachment"]["filename"] == "resume.pdf"
     assert submitted_data["attachment"]["transfer_method"] == "local_file"
@@ -791,10 +854,10 @@ def test_submit_form_by_token_normalizes_select_and_files(
 
 
 def test_submit_form_by_token_invalid_select_value(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
         inputs=[
@@ -807,7 +870,7 @@ def test_submit_form_by_token_invalid_select_value(
         rendered_content="<p>hello</p>",
         expiration_time=sample_form_record.expiration_time,
     )
-    repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
+    repo = form_repository(dataclasses.replace(sample_form_record, definition=definition))
     service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(InvalidFormDataError, match="Invalid value for select input 'decision'"):
@@ -820,10 +883,10 @@ def test_submit_form_by_token_invalid_select_value(
 
 
 def test_submit_form_by_token_invalid_file_list_item(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
         inputs=[FileListInputConfig(output_variable_name="attachments", number_limits=2)],
@@ -831,7 +894,7 @@ def test_submit_form_by_token_invalid_file_list_item(
         rendered_content="<p>hello</p>",
         expiration_time=sample_form_record.expiration_time,
     )
-    repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
+    repo = form_repository(dataclasses.replace(sample_form_record, definition=definition))
     service = HumanInputService(unbound_session_factory, form_repository=repo)
 
     with pytest.raises(
@@ -847,11 +910,11 @@ def test_submit_form_by_token_invalid_file_list_item(
 
 
 def test_submit_form_by_token_rejects_cross_tenant_file(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
         inputs=[FileInputConfig(output_variable_name="attachment")],
@@ -859,7 +922,7 @@ def test_submit_form_by_token_rejects_cross_tenant_file(
         rendered_content="<p>hello</p>",
         expiration_time=sample_form_record.expiration_time,
     )
-    repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
+    repo = form_repository(dataclasses.replace(sample_form_record, definition=definition))
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     mocker.patch("services.human_input_service.build_from_mapping", side_effect=ValueError("Invalid upload file"))
 
@@ -877,15 +940,15 @@ def test_submit_form_by_token_rejects_cross_tenant_file(
             },
         )
 
-    repo.mark_submitted.assert_not_called()
+    cast(MagicMock, repo.mark_submitted).assert_not_called()
 
 
 def test_submit_form_by_token_rejects_cross_tenant_file_list(
+    form_repository: Callable[[HumanInputFormRecord | None], HumanInputFormSubmissionRepository],
     sample_form_record: HumanInputFormRecord,
     unbound_session_factory: sessionmaker[Session],
     mocker: MockerFixture,
 ) -> None:
-    repo = MagicMock(spec=HumanInputFormSubmissionRepository)
     definition = FormDefinition(
         form_content="hello",
         inputs=[FileListInputConfig(output_variable_name="attachments", number_limits=2)],
@@ -893,7 +956,7 @@ def test_submit_form_by_token_rejects_cross_tenant_file_list(
         rendered_content="<p>hello</p>",
         expiration_time=sample_form_record.expiration_time,
     )
-    repo.get_by_token.return_value = dataclasses.replace(sample_form_record, definition=definition)
+    repo = form_repository(dataclasses.replace(sample_form_record, definition=definition))
     service = HumanInputService(unbound_session_factory, form_repository=repo)
     mocker.patch("services.human_input_service.build_from_mappings", side_effect=ValueError("Invalid upload file"))
 
@@ -916,4 +979,4 @@ def test_submit_form_by_token_rejects_cross_tenant_file_list(
             },
         )
 
-    repo.mark_submitted.assert_not_called()
+    cast(MagicMock, repo.mark_submitted).assert_not_called()