Merge branch 'main' into feat/dify-builder

api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py

@@ -1,4 +1,3 @@
-from collections.abc import Callable
 from datetime import datetime
 from inspect import getsource, unwrap
 from types import SimpleNamespace
@@ -309,15 +308,9 @@ def account_id() -> str:
 
 
 def test_agent_app_list_and_create_use_agent_route(
-    app: Flask,
-    monkeypatch: pytest.MonkeyPatch,
-    account_id: str,
-    sqlite_session: Session,
-    config_overrides: Callable[..., None],
+    app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str, sqlite_session: Session
 ) -> None:
     captured: dict[str, object] = {}
-    replace_whitelist = MagicMock()
-    initialize_access = MagicMock()
 
     class FakeAppService:
         def get_app(self, app_obj: object, *, session: object) -> object:
@@ -403,9 +396,7 @@ def test_agent_app_list_and_create_use_agent_route(
         lambda _self, **kwargs: {"agent-list": "debug-conversation-list"},
     )
     monkeypatch.setattr(
-        roster_controller.AgentRosterService,
-        "count_agent_app_debug_conversation_messages",
-        lambda _self, **kwargs: 0,
+        roster_controller.AgentRosterService, "count_agent_app_debug_conversation_messages", lambda _self, **kwargs: 0
     )
 
     def get_or_create_debug_conversation(_self: object, **kwargs: object) -> str:
@@ -422,13 +413,6 @@ def test_agent_app_list_and_create_use_agent_route(
         "get_system_features",
         lambda: SimpleNamespace(webapp_auth=SimpleNamespace(enabled=False)),
     )
-    config_overrides(RBAC_ENABLED=True)
-    monkeypatch.setattr(
-        roster_controller.enterprise_rbac_service.RBACService.AppAccess,
-        "replace_whitelist",
-        replace_whitelist,
-    )
-    monkeypatch.setattr(roster_controller.initialize_created_app_rbac_access_task, "delay", initialize_access)
     with app.test_request_context(
         "/console/api/agent?page=1&limit=10&mode=workflow&sort_by=recently_created"
         "&is_created_by_me=true&publication_status=published"
@@ -469,22 +453,12 @@ def test_agent_app_list_and_create_use_agent_route(
     assert count_params.agent_is_published is True
     with app.test_request_context(
         "/console/api/agent",
-        json={
-            "name": "Iris",
-            "description": "Agent app",
-            "role": "Coordinator",
-            "icon_type": "emoji",
-            "icon": "robot",
-        },
+        json={"name": "Iris", "description": "Agent app", "role": "Coordinator", "icon_type": "emoji", "icon": "robot"},
     ):
         created, status = unwrap(AgentAppListApi.post)(
             AgentAppListApi(),
             AgentAppCreatePayload(
-                name="Iris",
-                description="Agent app",
-                role="Coordinator",
-                icon_type="emoji",
-                icon="robot",
+                name="Iris", description="Agent app", role="Coordinator", icon_type="emoji", icon="robot"
             ),
             sqlite_session,
             "tenant-1",
@@ -507,81 +481,6 @@ def test_agent_app_list_and_create_use_agent_route(
         "account_id": account_id,
         "commit": False,
     }
-    replace_whitelist.assert_called_once()
-    assert replace_whitelist.call_args.args[:3] == ("tenant-1", account_id, "app-created")
-    replace_payload = replace_whitelist.call_args.args[3]
-    assert replace_payload.automatic_include_workspace_members is True
-    initialize_access.assert_called_once_with("tenant-1", account_id, app_id="app-created")
-
-
-def test_agent_app_create_skips_rbac_access_initialization_when_rbac_is_disabled(
-    app: Flask,
-    monkeypatch: pytest.MonkeyPatch,
-    account_id: str,
-    sqlite_session: Session,
-    config_overrides: Callable[..., None],
-) -> None:
-    replace_whitelist = MagicMock()
-    initialize_access = MagicMock()
-
-    class FakeAppService:
-        def get_app(self, app_obj: object, *, session: object) -> object:
-            return app_obj
-
-        def create_app(self, tenant_id: str, params, current_user: object, *, session: object) -> object:
-            return _app_detail_obj(id="app-created", bound_agent_id="agent-created")
-
-    monkeypatch.setattr(roster_controller, "AppService", FakeAppService)
-    monkeypatch.setattr(
-        roster_controller.AgentRosterService,
-        "get_app_backing_agent",
-        lambda _self, **kwargs: Agent(
-            id="agent-created",
-            app_id="app-created",
-            backing_app_id=None,
-            role="Created role",
-            active_config_snapshot_id=None,
-        ),
-    )
-    monkeypatch.setattr(
-        roster_controller.AgentRosterService,
-        "get_or_create_build_conversation",
-        lambda _self, **kwargs: "debug-conversation-created",
-    )
-    monkeypatch.setattr(
-        roster_controller.AgentRosterService, "count_agent_app_debug_conversation_messages", lambda _self, **kwargs: 0
-    )
-    monkeypatch.setattr(
-        roster_controller.FeatureService,
-        "get_system_features",
-        lambda: SimpleNamespace(webapp_auth=SimpleNamespace(enabled=False)),
-    )
-    config_overrides(RBAC_ENABLED=False)
-    monkeypatch.setattr(
-        roster_controller.enterprise_rbac_service.RBACService.AppAccess,
-        "replace_whitelist",
-        replace_whitelist,
-    )
-    monkeypatch.setattr(roster_controller.initialize_created_app_rbac_access_task, "delay", initialize_access)
-
-    with app.test_request_context(
-        "/console/api/agent",
-        json={"name": "Iris", "description": "Agent app", "role": "Coordinator", "icon_type": "emoji", "icon": "robot"},
-    ):
-        created, status = unwrap(AgentAppListApi.post)(
-            AgentAppListApi(),
-            AgentAppCreatePayload(
-                name="Iris", description="Agent app", role="Coordinator", icon_type="emoji", icon="robot"
-            ),
-            sqlite_session,
-            "tenant-1",
-            _account(account_id=account_id),
-        )
-
-    assert status == 201
-    assert created["id"] == "agent-created"
-    replace_whitelist.assert_not_called()
-    initialize_access.assert_not_called()
 
 
 def test_agent_app_create_payload_allows_optional_role() -> None:

api/tests/unit_tests/controllers/console/datasets/test_datasets.py

@@ -8,6 +8,7 @@ from unittest.mock import ANY, MagicMock, PropertyMock, call, patch
 
 import pytest
 from flask import Flask
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import BadRequest, Forbidden, NotFound
 
 import services
@@ -44,7 +45,7 @@ from core.rag.index_processor.constant.index_type import IndexStructureType
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from extensions.storage.storage_type import StorageType
 from models.account import Account, TenantAccountRole
-from models.dataset import Dataset, DatasetQuery, Document
+from models.dataset import AppDatasetJoin, Dataset, DatasetPermission, DatasetQuery, Document, DocumentSegment
 from models.enums import CreatorUserRole, DataSourceType, DocumentCreatedFrom, IndexingStatus
 from models.model import ApiToken, App, AppMode, IconType, UploadFile
 from services.dataset_ref_service import DatasetRef
@@ -170,7 +171,29 @@ def make_document_status(**overrides) -> Document:
     return Document(**base)
 
 
-class TestDatasetList:
+def make_document_segment(*, position: int, completed: bool) -> DocumentSegment:
+    return DocumentSegment(
+        tenant_id="tenant-1",
+        dataset_id="dataset-1",
+        document_id="doc-1",
+        position=position,
+        content=f"segment {position}",
+        word_count=2,
+        tokens=2,
+        created_by="account-1",
+        completed_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC) if completed else None,
+    )
+
+
+class _UsesSQLiteSession:
+    session: Session
+
+    @pytest.fixture(autouse=True)
+    def _inject_sqlite_session(self, sqlite_session: Session) -> None:
+        self.session = sqlite_session
+
+
+class TestDatasetList(_UsesSQLiteSession):
     def _mock_user(self):
         user = make_account()
         return user
@@ -185,7 +208,7 @@ class TestDatasetList:
                 patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         assert status == 200
         assert resp["total"] == 1
         assert resp["data"][0]["embedding_available"] is True
@@ -201,7 +224,7 @@ class TestDatasetList:
         method = unwrap(api.get)
         current_user = self._mock_user()
         dataset = make_dataset()
-        session = MagicMock()
+        session = self.session
         with app.test_request_context("/datasets"):
             with (
                 patch.object(DatasetService, "get_datasets", return_value=([dataset], 1)),
@@ -222,7 +245,7 @@ class TestDatasetList:
                 patch.object(DatasetService, "get_datasets_by_ids", return_value=(datasets, 2)) as by_ids_mock,
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         by_ids_mock.assert_called_once()
         assert status == 200
         assert resp["total"] == 2
@@ -251,7 +274,7 @@ class TestDatasetList:
                     return_value=permissions,
                 ) as get_permissions,
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         get_permissions.assert_called_once_with("tenant-1", current_user.id, session=ANY)
         assert status == 200
         assert resp["data"][0]["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"]
@@ -281,7 +304,7 @@ class TestDatasetList:
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                method(api, MagicMock(), "tenant-1", current_user)
+                method(api, self.session, "tenant-1", current_user)
         assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == []
         assert get_datasets.call_args.kwargs["include_own_datasets"] is False
 
@@ -308,7 +331,7 @@ class TestDatasetList:
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                method(api, MagicMock(), "tenant-1", current_user)
+                method(api, self.session, "tenant-1", current_user)
         assert get_datasets.call_args.kwargs["accessible_dataset_ids"] is None
 
     def test_get_restricted_whitelist_overrides_default_read_permission(
@@ -374,7 +397,7 @@ class TestDatasetList:
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                method(api, MagicMock(), "tenant-1", current_user)
+                method(api, self.session, "tenant-1", current_user)
         assert get_datasets.call_args.kwargs["accessible_dataset_ids"] == [
             "dataset-whitelist-only",
         ]
@@ -399,9 +422,9 @@ class TestDatasetList:
                 ),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                method(api, MagicMock(), "tenant-1", current_user)
+                method(api, self.session, "tenant-1", current_user)
         session = get_datasets_by_ids.call_args.kwargs["session"]
-        assert isinstance(session, MagicMock)
+        assert session is self.session
         assert get_datasets_by_ids.call_args.args == (["dataset-1"], "tenant-1")
         assert get_datasets_by_ids.call_args.kwargs == {
             "user": current_user,
@@ -420,7 +443,7 @@ class TestDatasetList:
                 patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         assert status == 200
 
     def test_get_allows_legacy_weighted_score_without_weight_type(self, app: Flask):
@@ -453,7 +476,7 @@ class TestDatasetList:
                 patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         assert status == 200
         assert resp["data"][0]["retrieval_model_dict"]["weights"]["weight_type"] is None
 
@@ -467,7 +490,7 @@ class TestDatasetList:
                 patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)),
                 patch.object(ProviderManager, "get_configurations", return_value=MagicMock(get_models=lambda **_: [])),
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         assert status == 200
         retrieval_model = resp["data"][0]["retrieval_model_dict"]
         assert retrieval_model["search_method"] == "semantic_search"
@@ -491,7 +514,7 @@ class TestDatasetList:
                 patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)),
                 patch.object(ProviderManager, "get_configurations", return_value=config),
             ):
-                resp, status = method(api, MagicMock(), "tenant-1", current_user)
+                resp, status = method(api, self.session, "tenant-1", current_user)
         assert resp["data"][0]["embedding_available"] is False
 
     def test_partial_members_permission(self, app: Flask):
@@ -499,8 +522,9 @@ class TestDatasetList:
         method = unwrap(api.get)
         current_user = self._mock_user()
         datasets = [make_dataset(permission="partial_members")]
-        session = MagicMock()
-        session.execute.return_value.all.return_value = [("ds-1", "u1")]
+        session = self.session
+        session.add(DatasetPermission(dataset_id="ds-1", account_id="u1", tenant_id="tenant-1"))
+        session.flush()
         with app.test_request_context("/datasets"):
             with (
                 patch.object(DatasetService, "get_datasets", return_value=(datasets, 1)),
@@ -510,7 +534,7 @@ class TestDatasetList:
         assert resp["data"][0]["partial_member_list"] == ["u1"]
 
 
-class TestDatasetListApiPost:
+class TestDatasetListApiPost(_UsesSQLiteSession):
     def test_post_success(self, app: Flask):
         api = DatasetListApi()
         method = unwrap(api.post)
@@ -522,7 +546,7 @@ class TestDatasetListApiPost:
             patch.object(type(console_ns), "payload", payload),
             patch.object(DatasetService, "create_empty_dataset", return_value=dataset),
         ):
-            _, status = method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user)
+            _, status = method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user)
         assert status == 201
 
     def test_post_forbidden(self, app: Flask):
@@ -532,7 +556,7 @@ class TestDatasetListApiPost:
         user = make_account(TenantAccountRole.NORMAL)
         with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload):
             with pytest.raises(Forbidden):
-                method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user)
+                method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user)
 
     def test_post_duplicate_name(self, app: Flask):
         api = DatasetListApi()
@@ -547,14 +571,14 @@ class TestDatasetListApiPost:
             ),
         ):
             with pytest.raises(DatasetNameDuplicateError):
-                method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", user)
+                method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", user)
 
     def test_post_invalid_payload_missing_name(self, app: Flask):
         api = DatasetListApi()
         method = unwrap(api.post)
         with app.test_request_context("/datasets", json={}), patch.object(type(console_ns), "payload", {}):
             with pytest.raises(ValueError):
-                method(api, DatasetCreatePayload(), MagicMock(), "tenant-1", make_account())
+                method(api, DatasetCreatePayload(), self.session, "tenant-1", make_account())
 
     def test_post_invalid_indexing_technique(self, app: Flask):
         api = DatasetListApi()
@@ -562,7 +586,7 @@ class TestDatasetListApiPost:
         payload = {"name": "bad", "indexing_technique": "invalid-tech"}
         with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload):
             with pytest.raises(ValueError, match="Invalid indexing technique"):
-                method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account())
+                method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account())
 
     def test_post_invalid_provider(self, app: Flask):
         api = DatasetListApi()
@@ -570,10 +594,10 @@ class TestDatasetListApiPost:
         payload = {"name": "bad", "provider": "unknown"}
         with app.test_request_context("/datasets", json=payload), patch.object(type(console_ns), "payload", payload):
             with pytest.raises(ValueError, match="Invalid provider"):
-                method(api, DatasetCreatePayload(**payload), MagicMock(), "tenant-1", make_account())
+                method(api, DatasetCreatePayload(**payload), self.session, "tenant-1", make_account())
 
 
-class TestDatasetApiGet:
+class TestDatasetApiGet(_UsesSQLiteSession):
     def test_get_success_basic(self, app: Flask):
         api = DatasetApi()
         method = unwrap(api.get)
@@ -588,7 +612,7 @@ class TestDatasetApiGet:
             patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock,
         ):
             provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = []
-            data, status = method(api, MagicMock(), tenant_id, user, dataset_id)
+            data, status = method(api, self.session, tenant_id, user, dataset_id)
         assert status == 200
         assert data["embedding_available"] is True
 
@@ -597,7 +621,7 @@ class TestDatasetApiGet:
         api = DatasetApi()
         method = unwrap(api.get)
         dataset_id = "123e4567-e89b-12d3-a456-426614174000"
-        user = MagicMock(id="account-1")
+        user = make_account()
         tenant_id = "tenant-1"
         dataset = make_dataset(id=dataset_id)
         with (
@@ -619,7 +643,7 @@ class TestDatasetApiGet:
             patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock,
         ):
             provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = []
-            data, status = method(api, MagicMock(), tenant_id, user, dataset_id)
+            data, status = method(api, self.session, tenant_id, user, dataset_id)
         get_permissions.assert_called_once_with(tenant_id, user.id, dataset_id=dataset_id, session=ANY)
         assert status == 200
         assert data["permission_keys"] == ["dataset.acl.readonly", "dataset.acl.edit"]
@@ -636,7 +660,7 @@ class TestDatasetApiGet:
             patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock,
         ):
             provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = []
-            data, status = method(api, MagicMock(), "tenant", make_account(), dataset_id)
+            data, status = method(api, self.session, "tenant", make_account(), dataset_id)
         assert status == 200
         assert data["external_retrieval_model"] == {"top_k": 2, "score_threshold": 0.0, "score_threshold_enabled": None}
 
@@ -649,7 +673,7 @@ class TestDatasetApiGet:
             patch.object(DatasetService, "get_dataset", return_value=None),
         ):
             with pytest.raises(NotFound, match="Dataset not found"):
-                method(api, MagicMock(), "tenant", make_account(), dataset_id)
+                method(api, self.session, "tenant", make_account(), dataset_id)
 
     def test_get_permission_denied(self, app: Flask):
         api = DatasetApi()
@@ -666,7 +690,7 @@ class TestDatasetApiGet:
             ),
         ):
             with pytest.raises(Forbidden, match="no access"):
-                method(api, MagicMock(), "tenant", make_account(), dataset_id)
+                method(api, self.session, "tenant", make_account(), dataset_id)
 
     def test_get_high_quality_embedding_unavailable(self, app: Flask):
         api = DatasetApi()
@@ -687,7 +711,7 @@ class TestDatasetApiGet:
             patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock,
         ):
             provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = []
-            data, _ = method(api, MagicMock(), tenant_id, user, dataset_id)
+            data, _ = method(api, self.session, tenant_id, user, dataset_id)
         assert data["embedding_available"] is False
 
     def test_get_partial_members_permission(self, app: Flask):
@@ -704,11 +728,11 @@ class TestDatasetApiGet:
             patch("controllers.console.datasets.datasets.create_plugin_provider_manager") as provider_manager_mock,
         ):
             provider_manager_mock.return_value.get_configurations.return_value.get_models.return_value = []
-            data, _ = method(api, MagicMock(), "tenant", make_account(), dataset_id)
+            data, _ = method(api, self.session, "tenant", make_account(), dataset_id)
         assert data["partial_member_list"] == partial_members
 
 
-class TestDatasetApiPatch:
+class TestDatasetApiPatch(_UsesSQLiteSession):
     def test_patch_success_basic(self, app: Flask):
         api = DatasetApi()
         method = unwrap(api.patch)
@@ -725,7 +749,7 @@ class TestDatasetApiPatch:
             patch.object(DatasetService, "update_dataset", return_value=dataset),
             patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]),
         ):
-            result, status = method(api, DatasetUpdatePayload(), MagicMock(), tenant_id, user, dataset_id)
+            result, status = method(api, DatasetUpdatePayload(), self.session, tenant_id, user, dataset_id)
         assert status == 200
         assert result["partial_member_list"] == []
 
@@ -737,7 +761,7 @@ class TestDatasetApiPatch:
             patch.object(DatasetService, "get_dataset", return_value=None),
         ):
             with pytest.raises(NotFound, match="Dataset not found"):
-                method(api, DatasetUpdatePayload(), MagicMock(), "tenant-1", make_account(), "missing")
+                method(api, DatasetUpdatePayload(), self.session, "tenant-1", make_account(), "missing")
 
     def test_patch_permission_denied(self, app: Flask):
         api = DatasetApi()
@@ -752,7 +776,7 @@ class TestDatasetApiPatch:
             patch.object(DatasetPermissionService, "check_permission", side_effect=Forbidden("no permission")),
         ):
             with pytest.raises(Forbidden):
-                method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id)
+                method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id)
 
     def test_patch_partial_members_update(self, app: Flask):
         api = DatasetApi()
@@ -769,7 +793,7 @@ class TestDatasetApiPatch:
             patch.object(DatasetPermissionService, "update_partial_member_list", return_value=None),
             patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=["u1", "u2"]),
         ):
-            result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id)
+            result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id)
         assert result["partial_member_list"] == ["u1", "u2"]
 
     def test_patch_clear_partial_members(self, app: Flask):
@@ -787,11 +811,11 @@ class TestDatasetApiPatch:
             patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None),
             patch.object(DatasetPermissionService, "get_dataset_partial_member_list", return_value=[]),
         ):
-            result, _ = method(api, DatasetUpdatePayload(), MagicMock(), "tenant", make_account(), dataset_id)
+            result, _ = method(api, DatasetUpdatePayload(), self.session, "tenant", make_account(), dataset_id)
         assert result["partial_member_list"] == []
 
 
-class TestDatasetApiDelete:
+class TestDatasetApiDelete(_UsesSQLiteSession):
     def test_delete_success(self, app: Flask):
         api = DatasetApi()
         method = unwrap(api.delete)
@@ -802,7 +826,7 @@ class TestDatasetApiDelete:
             patch.object(DatasetService, "delete_dataset", return_value=True),
             patch.object(DatasetPermissionService, "clear_partial_member_list", return_value=None),
         ):
-            result, status = method(api, MagicMock(), user, dataset_id)
+            result, status = method(api, self.session, user, dataset_id)
         assert status == 204
         assert result == ""
 
@@ -813,7 +837,7 @@ class TestDatasetApiDelete:
         user = make_account(TenantAccountRole.NORMAL)
         with app.test_request_context(f"/datasets/{dataset_id}"):
             with pytest.raises(Forbidden):
-                method(api, MagicMock(), user, dataset_id)
+                method(api, self.session, user, dataset_id)
 
     def test_delete_dataset_not_found(self, app: Flask):
         api = DatasetApi()
@@ -825,7 +849,7 @@ class TestDatasetApiDelete:
             patch.object(DatasetService, "delete_dataset", return_value=False),
         ):
             with pytest.raises(NotFound, match="Dataset not found"):
-                method(api, MagicMock(), user, dataset_id)
+                method(api, self.session, user, dataset_id)
 
     def test_delete_dataset_in_use(self, app: Flask):
         api = DatasetApi()
@@ -837,10 +861,10 @@ class TestDatasetApiDelete:
             patch.object(DatasetService, "delete_dataset", side_effect=services.errors.dataset.DatasetInUseError()),
         ):
             with pytest.raises(DatasetInUseError):
-                method(api, MagicMock(), user, dataset_id)
+                method(api, self.session, user, dataset_id)
 
 
-class TestDatasetUseCheckApi:
+class TestDatasetUseCheckApi(_UsesSQLiteSession):
     @pytest.mark.parametrize("is_using", [True, False])
     def test_get_use_check(self, app: Flask, is_using: bool):
         api = DatasetUseCheckApi()
@@ -848,7 +872,7 @@ class TestDatasetUseCheckApi:
         dataset_id = "dataset-id"
         dataset = make_dataset(id=dataset_id)
         current_user = make_account()
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context(f"/datasets/{dataset_id}/use-check"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset,
@@ -867,7 +891,7 @@ class TestDatasetUseCheckApi:
         api = DatasetUseCheckApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-id")
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/datasets/dataset-id/use-check"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset),
@@ -884,11 +908,11 @@ class TestDatasetUseCheckApi:
     "api_cls",
     [DatasetUseCheckApi, DatasetIndexingStatusApi, DatasetErrorDocs, DatasetAutoDisableLogApi],
 )
-def test_dataset_scoped_read_permission_denied(app: Flask, api_cls):
+def test_dataset_scoped_read_permission_denied(app: Flask, api_cls, sqlite_session: Session):
     api = api_cls()
     method = unwrap(api.get)
     dataset = make_dataset(id="dataset-1")
-    session = MagicMock()
+    session = sqlite_session
     with (
         app.test_request_context("/"),
         patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset),
@@ -902,7 +926,7 @@ def test_dataset_scoped_read_permission_denied(app: Flask, api_cls):
             method(api, session, "tenant-1", make_account(), "dataset-1")
 
 
-class TestDatasetQueryApi:
+class TestDatasetQueryApi(_UsesSQLiteSession):
     def _query_record(self, index: int = 1) -> DatasetQuery:
         query = DatasetQuery(
             dataset_id="dataset-id",
@@ -929,7 +953,7 @@ class TestDatasetQueryApi:
             patch.object(DatasetService, "check_dataset_permission", return_value=None),
             patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 2)),
         ):
-            response, status = method(api, MagicMock(), current_user, dataset_id)
+            response, status = method(api, self.session, current_user, dataset_id)
         assert status == 200
         assert response["total"] == 2
         assert response["page"] == 1
@@ -952,24 +976,30 @@ class TestDatasetQueryApi:
         dataset = make_dataset(id="dataset-id")
         query = self._query_record()
         query.content = json.dumps([{"content_type": "image_query", "content": "file-1"}])
-        upload_file = SimpleNamespace(
-            id="file-1",
+        upload_file = UploadFile(
+            tenant_id="tenant-1",
+            storage_type=StorageType.LOCAL,
+            key="image.png",
             name="image.png",
             size=10,
             extension="png",
             mime_type="image/png",
+            created_by_role=CreatorUserRole.ACCOUNT,
+            created_by="account-1",
+            created_at=datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC),
+            used=False,
         )
-        session = MagicMock()
-        session.scalar.return_value = upload_file
+        upload_file.id = "file-1"
+        session = self.session
+        session.add(upload_file)
+        session.flush()
         with (
             app.test_request_context("/datasets/queries"),
             patch.object(DatasetService, "get_dataset", return_value=dataset),
             patch.object(DatasetService, "check_dataset_permission", return_value=None),
             patch.object(DatasetService, "get_dataset_queries", return_value=([query], 1)),
-            patch("models.dataset.db") as db_mock,
             patch("models.dataset.sign_upload_file_preview_url", return_value="signed-url"),
         ):
-            db_mock.session.scalar.return_value = upload_file
             response, status = method(api, session, make_account(), "dataset-id")
 
         assert status == 200
@@ -987,8 +1017,7 @@ class TestDatasetQueryApi:
                 },
             }
         ]
-        session.scalar.assert_called_once()
-        db_mock.session.scalar.assert_not_called()
+        assert session.get(UploadFile, "file-1") is upload_file
 
     def test_get_queries_dataset_not_found(self, app: Flask):
         api = DatasetQueryApi()
@@ -1000,7 +1029,7 @@ class TestDatasetQueryApi:
             patch.object(DatasetService, "get_dataset", return_value=None),
         ):
             with pytest.raises(NotFound, match="Dataset not found"):
-                method(api, MagicMock(), current_user, dataset_id)
+                method(api, self.session, current_user, dataset_id)
 
     def test_get_queries_permission_denied(self, app: Flask):
         api = DatasetQueryApi()
@@ -1018,7 +1047,7 @@ class TestDatasetQueryApi:
             ),
         ):
             with pytest.raises(Forbidden):
-                method(api, MagicMock(), current_user, dataset_id)
+                method(api, self.session, current_user, dataset_id)
 
     def test_get_queries_pagination_has_more(self, app: Flask):
         api = DatasetQueryApi()
@@ -1033,13 +1062,13 @@ class TestDatasetQueryApi:
             patch.object(DatasetService, "check_dataset_permission", return_value=None),
             patch.object(DatasetService, "get_dataset_queries", return_value=(queries, 40)),
         ):
-            response, status = method(api, MagicMock(), current_user, dataset_id)
+            response, status = method(api, self.session, current_user, dataset_id)
         assert status == 200
         assert response["has_more"] is True
         assert len(response["data"]) == 20
 
 
-class TestDatasetIndexingEstimateApi:
+class TestDatasetIndexingEstimateApi(_UsesSQLiteSession):
     def _upload_file(self, *, tenant_id: str = "tenant-1", file_id: str = "file-1") -> UploadFile:
         upload_file = UploadFile(
             tenant_id=tenant_id,
@@ -1072,8 +1101,9 @@ class TestDatasetIndexingEstimateApi:
         method = unwrap(api.post)
         payload = self._base_payload()
         mock_file = self._upload_file()
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [mock_file]
+        session = self.session
+        session.add(mock_file)
+        session.flush()
 
         mock_response = IndexingEstimate(total_segments=100, preview=[])
 
@@ -1102,8 +1132,7 @@ class TestDatasetIndexingEstimateApi:
         api = DatasetIndexingEstimateApi()
         method = unwrap(api.post)
         payload = self._base_payload()
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = None
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
@@ -1122,8 +1151,9 @@ class TestDatasetIndexingEstimateApi:
         method = unwrap(api.post)
         mock_file = self._upload_file()
         payload = self._base_payload()
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [mock_file]
+        session = self.session
+        session.add(mock_file)
+        session.flush()
         with (
             app.test_request_context("/"),
             patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
@@ -1146,8 +1176,9 @@ class TestDatasetIndexingEstimateApi:
         method = unwrap(api.post)
         mock_file = self._upload_file()
         payload = self._base_payload()
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [mock_file]
+        session = self.session
+        session.add(mock_file)
+        session.flush()
         with (
             app.test_request_context("/"),
             patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
@@ -1170,8 +1201,9 @@ class TestDatasetIndexingEstimateApi:
         method = unwrap(api.post)
         mock_file = self._upload_file()
         payload = self._base_payload()
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [mock_file]
+        session = self.session
+        session.add(mock_file)
+        session.flush()
         with (
             app.test_request_context("/"),
             patch.object(type(console_ns), "payload", new_callable=PropertyMock, return_value=payload),
@@ -1189,16 +1221,16 @@ class TestDatasetIndexingEstimateApi:
                 )
 
 
-class TestDatasetRelatedAppListApi:
+class TestDatasetRelatedAppListApi(_UsesSQLiteSession):
     def test_get_success(self, app: Flask):
         api = DatasetRelatedAppListApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
         app1 = make_related_app(id="app-1", name="App 1")
         app2 = make_related_app(id="app-2", name="App 2")
-        join1 = MagicMock(app_id="app-1")
-        join2 = MagicMock(app_id="app-2")
-        session = MagicMock()
+        join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1")
+        join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1")
+        session = self.session
         with (
             app.test_request_context("/"),
             patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset),
@@ -1251,7 +1283,7 @@ class TestDatasetRelatedAppListApi:
             patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=None),
         ):
             with pytest.raises(NotFound):
-                method(api, MagicMock(), make_account(), "dataset-1")
+                method(api, self.session, make_account(), "dataset-1")
 
     def test_get_permission_denied(self, app: Flask):
         api = DatasetRelatedAppListApi()
@@ -1266,16 +1298,16 @@ class TestDatasetRelatedAppListApi:
             ),
         ):
             with pytest.raises(Forbidden):
-                method(api, MagicMock(), make_account(), "dataset-1")
+                method(api, self.session, make_account(), "dataset-1")
 
     def test_get_filters_none_apps(self, app: Flask):
         api = DatasetRelatedAppListApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
         app1 = make_related_app()
-        join1 = MagicMock(app_id="app-1")
-        join2 = MagicMock(app_id="app-2")
-        session = MagicMock()
+        join1 = AppDatasetJoin(app_id="app-1", dataset_id="dataset-1")
+        join2 = AppDatasetJoin(app_id="app-2", dataset_id="dataset-1")
+        session = self.session
         with (
             app.test_request_context("/"),
             patch("controllers.console.datasets.datasets.DatasetService.get_dataset", return_value=dataset),
@@ -1303,26 +1335,17 @@ class TestDatasetRelatedAppListApi:
         ]
 
 
-class TestDatasetIndexingStatusApi:
+class TestDatasetIndexingStatusApi(_UsesSQLiteSession):
     def test_get_success_with_documents(self, app: Flask):
         api = DatasetIndexingStatusApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
         current_user = make_account()
-        document = MagicMock()
-        document.id = "doc-1"
-        document.indexing_status = "completed"
-        document.processing_started_at = None
-        document.parsing_completed_at = None
-        document.cleaning_completed_at = None
-        document.splitting_completed_at = None
-        document.completed_at = None
-        document.paused_at = None
-        document.error = None
-        document.stopped_at = None
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [document]
-        session.scalar.return_value = 3
+        document = make_document_status()
+        session = self.session
+        session.add(document)
+        session.add_all([make_document_segment(position=position, completed=True) for position in range(1, 4)])
+        session.flush()
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset,
@@ -1337,16 +1360,13 @@ class TestDatasetIndexingStatusApi:
         assert item["total_segments"] == 3
         get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session)
         check_permission.assert_called_once_with(dataset, current_user, session)
-        assert {"dataset-1", "tenant-1"} <= set(session.scalars.call_args.args[0].compile().params.values())
-        for segment_count_call in session.scalar.call_args_list:
-            assert {"dataset-1", "tenant-1", "doc-1"} <= set(segment_count_call.args[0].compile().params.values())
+        assert session.get(Document, "doc-1") is document
 
     def test_get_success_no_documents(self, app: Flask):
         api = DatasetIndexingStatusApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = []
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset),
@@ -1360,20 +1380,11 @@ class TestDatasetIndexingStatusApi:
         api = DatasetIndexingStatusApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
-        document = MagicMock()
-        document.id = "doc-1"
-        document.indexing_status = "indexing"
-        document.processing_started_at = None
-        document.parsing_completed_at = None
-        document.cleaning_completed_at = None
-        document.splitting_completed_at = None
-        document.completed_at = None
-        document.paused_at = None
-        document.error = None
-        document.stopped_at = None
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [document]
-        session.scalar.side_effect = [2, 5]
+        document = make_document_status(indexing_status=IndexingStatus.INDEXING)
+        session = self.session
+        session.add(document)
+        session.add_all([make_document_segment(position=position, completed=position <= 2) for position in range(1, 6)])
+        session.flush()
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset),
@@ -1386,7 +1397,7 @@ class TestDatasetIndexingStatusApi:
         assert item["total_segments"] == 5
 
 
-class TestDatasetApiKeyApi:
+class TestDatasetApiKeyApi(_UsesSQLiteSession):
     def test_get_api_keys_success(self, app: Flask):
         api = DatasetApiKeyApi()
         method = unwrap(api.get)
@@ -1404,8 +1415,11 @@ class TestDatasetApiKeyApi:
             last_used_at=None,
             created_at=None,
         )
-        session = MagicMock()
-        session.scalars.return_value.all.return_value = [mock_key_1, mock_key_2]
+        session = self.session
+        mock_key_1.tenant_id = "tenant-1"
+        mock_key_2.tenant_id = "tenant-1"
+        session.add_all([mock_key_1, mock_key_2])
+        session.flush()
         with app.test_request_context("/"):
             response = method(api, session, "tenant-1")
         assert "data" in response
@@ -1418,30 +1432,31 @@ class TestDatasetApiKeyApi:
     def test_post_create_api_key_success(self, app: Flask):
         api = DatasetApiKeyApi()
         method = unwrap(api.post)
-        mock_token = MagicMock()
-        mock_token.id = "new-key-id"
-        mock_token.last_used_at = None
-        mock_token.created_at = datetime.datetime(2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC)
-        mock_api_token_cls = MagicMock()
-        mock_api_token_cls.return_value = mock_token
-        mock_api_token_cls.generate_api_key.return_value = "dataset-abc123"
-        session = MagicMock()
-        session.scalar.return_value = 3
-        with app.test_request_context("/"), patch("controllers.console.datasets.datasets.ApiToken", mock_api_token_cls):
+        session = self.session
+        with (
+            app.test_request_context("/"),
+            patch.object(ApiToken, "generate_api_key", return_value="dataset-abc123") as generate_api_key,
+        ):
             response, status = method(api, session, "tenant-1")
         assert status == 200
         assert isinstance(response, dict)
-        assert response["id"] == "new-key-id"
         assert response["token"] == "dataset-abc123"
         assert response["type"] == "dataset"
         assert response["created_at"] is not None
-        mock_api_token_cls.generate_api_key.assert_called_once_with("dataset-", 24, session=session)
+        generate_api_key.assert_called_once_with("dataset-", 24, session=session)
+        assert session.get(ApiToken, response["id"]).token == "dataset-abc123"
 
     def test_post_exceed_max_keys(self, app: Flask):
         api = DatasetApiKeyApi()
         method = unwrap(api.post)
-        session = MagicMock()
-        session.scalar.return_value = 10
+        session = self.session
+        session.add_all(
+            [
+                ApiToken(id=f"key-{index}", tenant_id="tenant-1", type="dataset", token=f"ds-{index}")
+                for index in range(10)
+            ]
+        )
+        session.flush()
         with app.test_request_context("/"):
             with pytest.raises(BadRequest) as exc_info:
                 method(api, session, "tenant-1")
@@ -1452,36 +1467,42 @@ class TestDatasetApiKeyApi:
         }
 
 
-class TestDatasetApiDeleteApi:
+class TestDatasetApiDeleteApi(_UsesSQLiteSession):
     def test_delete_success(self, app: Flask):
         api = DatasetApiDeleteApi()
         method = unwrap(api.delete)
-        mock_key = MagicMock()
-        session = MagicMock()
-        session.scalar.return_value = mock_key
-        with app.test_request_context("/"):
+        session = self.session
+        key = ApiToken(id="api-key-id", tenant_id="tenant-1", type="dataset", token="dataset-secret")
+        session.add(key)
+        session.flush()
+        with (
+            app.test_request_context("/"),
+            patch("controllers.console.datasets.datasets.ApiTokenCache.delete") as delete_cache,
+        ):
             response, status = method(api, session, "tenant-1", "api-key-id")
         assert status == 204
         assert response == ""
+        delete_cache.assert_called_once()
+        session.flush()
+        assert session.get(ApiToken, "api-key-id") is None
 
     def test_delete_key_not_found(self, app: Flask):
         api = DatasetApiDeleteApi()
         method = unwrap(api.delete)
-        session = MagicMock()
-        session.scalar.return_value = None
+        session = self.session
         with app.test_request_context("/"):
             with pytest.raises(NotFound):
                 method(api, session, "tenant-1", "api-key-id")
 
 
-class TestDatasetEnableApiApi:
+class TestDatasetEnableApiApi(_UsesSQLiteSession):
     @pytest.mark.parametrize(("status_value", "enabled"), [("enable", True), ("disable", False)])
     def test_update_api_status(self, app: Flask, status_value: str, enabled: bool):
         api = DatasetEnableApiApi()
         method = unwrap(api.post)
         dataset = make_dataset(id="dataset-1")
         current_user = make_account()
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset,
@@ -1499,7 +1520,7 @@ class TestDatasetEnableApiApi:
         api = DatasetEnableApiApi()
         method = unwrap(api.post)
         dataset = make_dataset(id="dataset-1")
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset),
@@ -1582,7 +1603,7 @@ class TestDatasetRetrievalSettingApi:
         ]
 
 
-class TestDatasetRetrievalSettingMockApi:
+class TestDatasetRetrievalSettingMockApi(_UsesSQLiteSession):
     def test_get_success(self, app: Flask):
         api = DatasetRetrievalSettingMockApi()
         method = unwrap(api.get)
@@ -1597,14 +1618,14 @@ class TestDatasetRetrievalSettingMockApi:
         assert response["retrieval_method"] == ["semantic"]
 
 
-class TestDatasetErrorDocs:
+class TestDatasetErrorDocs(_UsesSQLiteSession):
     def test_get_success(self, app: Flask):
         api = DatasetErrorDocs()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
         error_doc = make_document_status(id="error-doc", indexing_status=IndexingStatus.ERROR, error="failed")
         current_user = make_account()
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset,
@@ -1624,7 +1645,7 @@ class TestDatasetErrorDocs:
     def test_get_dataset_not_found(self, app: Flask):
         api = DatasetErrorDocs()
         method = unwrap(api.get)
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset,
@@ -1634,7 +1655,7 @@ class TestDatasetErrorDocs:
         get_dataset.assert_called_once_with("dataset-1", "tenant-1", session=session)
 
 
-class TestDatasetPermissionUserListApi:
+class TestDatasetPermissionUserListApi(_UsesSQLiteSession):
     def test_get_success(self, app: Flask):
         api = DatasetPermissionUserListApi()
         method = unwrap(api.get)
@@ -1649,7 +1670,7 @@ class TestDatasetPermissionUserListApi:
                 return_value=users,
             ),
         ):
-            response, status = method(api, MagicMock(), make_account(), "dataset-1")
+            response, status = method(api, self.session, make_account(), "dataset-1")
         assert status == 200
         assert response["data"] == users
 
@@ -1666,17 +1687,17 @@ class TestDatasetPermissionUserListApi:
             ),
         ):
             with pytest.raises(Forbidden):
-                method(api, MagicMock(), make_account(), "dataset-1")
+                method(api, self.session, make_account(), "dataset-1")
 
 
-class TestDatasetAutoDisableLogApi:
+class TestDatasetAutoDisableLogApi(_UsesSQLiteSession):
     def test_get_success(self, app: Flask):
         api = DatasetAutoDisableLogApi()
         method = unwrap(api.get)
         dataset = make_dataset(id="dataset-1")
         logs = {"document_ids": ["doc-1"], "count": 1}
         current_user = make_account()
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=dataset) as get_dataset,
@@ -1693,7 +1714,7 @@ class TestDatasetAutoDisableLogApi:
     def test_get_dataset_not_found(self, app: Flask):
         api = DatasetAutoDisableLogApi()
         method = unwrap(api.get)
-        session = MagicMock()
+        session = self.session
         with (
             app.test_request_context("/"),
             patch.object(DatasetService, "get_dataset_for_tenant", return_value=None) as get_dataset,

api/tests/unit_tests/controllers/console/workspace/test_model_providers.py

@@ -1,11 +1,13 @@
+from collections.abc import Callable
 from inspect import unwrap
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
-from flask import Flask, g
+from flask import Flask, g, request
 from pydantic_core import ValidationError
 from sqlalchemy.orm import Session
-from werkzeug.exceptions import Forbidden
+from werkzeug.exceptions import Forbidden, UnprocessableEntity
 
 from configs import dify_config
 from controllers.console.workspace.model_providers import (
@@ -17,6 +19,14 @@ from controllers.console.workspace.model_providers import (
     ModelProviderPaymentCheckoutUrlApi,
     ModelProviderSummaryListApi,
     ModelProviderValidateApi,
+    ParserCredentialCreate,
+    ParserCredentialDelete,
+    ParserCredentialId,
+    ParserCredentialSwitch,
+    ParserCredentialUpdate,
+    ParserCredentialValidate,
+    ParserModelList,
+    ParserPreferredProviderType,
     PreferredProviderTypeUpdateApi,
 )
 from core.entities.provider_entities import CredentialConfiguration
@@ -26,6 +36,7 @@ from graphon.model_runtime.entities.model_entities import ModelType
 from graphon.model_runtime.entities.provider_entities import ConfigurateMethod
 from graphon.model_runtime.errors.validate import CredentialsValidateFailedError
 from models import Account
+from models.account import TenantAccountRole
 from models.provider import ProviderType
 from services.entities.model_provider_entities import (
     CustomConfigurationResponse,
@@ -116,6 +127,19 @@ def expected_provider_payload() -> dict[str, object]:
     }
 
 
+def _payload() -> dict[str, object]:
+    """Mirror the source the ``@model_validate`` decorator reads for the request in scope.
+
+    The tests build their contexts with ``test_request_context``, which defaults to GET even for
+    handlers mounted on POST, so fall back to the JSON body whenever the query string is empty.
+    """
+    args: dict[str, object] = dict(request.args.to_dict(flat=True))
+    if args:
+        return args
+    body = request.get_json(silent=True)
+    return dict(body) if isinstance(body, dict) else {}
+
+
 class TestModelProviderListApi:
     def test_get_success(self, app: Flask):
         api = ModelProviderListApi()
@@ -129,7 +153,7 @@ class TestModelProviderListApi:
                 return_value=[provider],
             ) as get_provider_list,
         ):
-            result = method(api, "tenant1")
+            result = method(api, ParserModelList.model_validate(_payload()), "tenant1")
 
         get_provider_list.assert_called_once_with(tenant_id="tenant1", model_type=ModelType.LLM)
         assert result == {"data": [expected_provider_payload()]}
@@ -145,7 +169,7 @@ class TestModelProviderListApi:
                 return_value=[],
             ) as get_provider_list,
         ):
-            result = method(api, "tenant1")
+            result = method(api, ParserModelList.model_validate(_payload()), "tenant1")
 
         get_provider_list.assert_called_once_with(tenant_id="tenant1", model_type=None)
         assert result == {"data": []}
@@ -297,7 +321,7 @@ class TestModelProviderCredentialApi:
                 },
             ) as get_provider_credential,
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserCredentialId.model_validate(_payload()), "tenant1", provider="openai")
 
         get_provider_credential.assert_called_once_with(
             tenant_id="tenant1", provider="openai", credential_id=VALID_UUID
@@ -321,7 +345,7 @@ class TestModelProviderCredentialApi:
                 return_value=None,
             ) as get_provider_credential,
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserCredentialId.model_validate(_payload()), "tenant1", provider="openai")
 
         get_provider_credential.assert_called_once_with(tenant_id="tenant1", provider="openai", credential_id=None)
         assert result == {"credentials": None}
@@ -332,7 +356,7 @@ class TestModelProviderCredentialApi:
 
         with app.test_request_context(f"/?credential_id={INVALID_UUID}"):
             with pytest.raises(ValidationError):
-                method(api, "tenant1", provider="openai")
+                method(api, ParserCredentialId.model_validate(_payload()), "tenant1", provider="openai")
 
     def test_post_create_success(self, app: Flask):
         api = ModelProviderCredentialApi()
@@ -347,7 +371,9 @@ class TestModelProviderCredentialApi:
                 return_value=None,
             ) as create_provider_credential,
         ):
-            result, status = method(api, "tenant1", provider="openai")
+            result, status = method(
+                api, ParserCredentialCreate.model_validate(_payload()), "tenant1", provider="openai"
+            )
 
         create_provider_credential.assert_called_once_with(
             tenant_id="tenant1",
@@ -372,7 +398,7 @@ class TestModelProviderCredentialApi:
             ),
         ):
             with pytest.raises(ValueError):
-                method(api, "tenant1", provider="openai")
+                method(api, ParserCredentialCreate.model_validate(_payload()), "tenant1", provider="openai")
 
     def test_put_update_success(self, app: Flask):
         api = ModelProviderCredentialApi()
@@ -387,7 +413,7 @@ class TestModelProviderCredentialApi:
                 return_value=None,
             ) as update_provider_credential,
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserCredentialUpdate.model_validate(_payload()), "tenant1", provider="openai")
 
         update_provider_credential.assert_called_once_with(
             tenant_id="tenant1",
@@ -406,7 +432,7 @@ class TestModelProviderCredentialApi:
 
         with app.test_request_context("/", json=payload):
             with pytest.raises(ValidationError):
-                method(api, "tenant1", provider="openai")
+                method(api, ParserCredentialUpdate.model_validate(_payload()), "tenant1", provider="openai")
 
     def test_delete_success(self, app: Flask):
         api = ModelProviderCredentialApi()
@@ -421,7 +447,9 @@ class TestModelProviderCredentialApi:
                 return_value=None,
             ) as remove_provider_credential,
         ):
-            result, status = method(api, "tenant1", provider="openai")
+            result, status = method(
+                api, ParserCredentialDelete.model_validate(_payload()), "tenant1", provider="openai"
+            )
 
         remove_provider_credential.assert_called_once_with(
             tenant_id="tenant1", provider="openai", credential_id=VALID_UUID
@@ -444,7 +472,7 @@ class TestModelProviderCredentialSwitchApi:
                 return_value=None,
             ) as switch_active_provider_credential,
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserCredentialSwitch.model_validate(_payload()), "tenant1", provider="openai")
 
         switch_active_provider_credential.assert_called_once_with(
             tenant_id="tenant1",
@@ -461,7 +489,7 @@ class TestModelProviderCredentialSwitchApi:
 
         with app.test_request_context("/", json=payload):
             with pytest.raises(ValidationError):
-                method(api, "tenant1", provider="openai")
+                method(api, ParserCredentialSwitch.model_validate(_payload()), "tenant1", provider="openai")
 
 
 class TestModelProviderValidateApi:
@@ -478,7 +506,7 @@ class TestModelProviderValidateApi:
                 return_value=None,
             ) as validate_provider_credentials,
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserCredentialValidate.model_validate(_payload()), "tenant1", provider="openai")
 
         validate_provider_credentials.assert_called_once_with(
             tenant_id="tenant1", provider="openai", credentials={"a": "b"}
@@ -498,7 +526,7 @@ class TestModelProviderValidateApi:
                 side_effect=CredentialsValidateFailedError("bad"),
             ),
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserCredentialValidate.model_validate(_payload()), "tenant1", provider="openai")
 
         assert result == {"result": "error", "error": "bad"}
 
@@ -549,7 +577,7 @@ class TestPreferredProviderTypeUpdateApi:
                 return_value=None,
             ) as switch_preferred_provider,
         ):
-            result = method(api, "tenant1", provider="openai")
+            result = method(api, ParserPreferredProviderType.model_validate(_payload()), "tenant1", provider="openai")
 
         switch_preferred_provider.assert_called_once_with(
             tenant_id="tenant1", provider="openai", preferred_provider_type="custom"
@@ -564,7 +592,7 @@ class TestPreferredProviderTypeUpdateApi:
 
         with app.test_request_context("/", json=payload):
             with pytest.raises(ValidationError):
-                method(api, "tenant1", provider="openai")
+                method(api, ParserPreferredProviderType.model_validate(_payload()), "tenant1", provider="openai")
 
 
 class TestModelProviderPaymentCheckoutUrlApi:
@@ -618,3 +646,57 @@ class TestModelProviderPaymentCheckoutUrlApi:
                 api.get(provider="anthropic")
 
         get_model_provider_payment_link.assert_not_called()
+
+
+class TestModelValidateDecorator:
+    """The tests above unwrap the view, so this is what covers the decorators themselves."""
+
+    EMPTY_BODY: dict[str, object] = {}
+    EMPTY_KWARGS: dict[str, object] = {}
+
+    @pytest.mark.parametrize(
+        ("verb", "url", "body", "call"),
+        [
+            ("GET", "/?model_type=not-a-model-type", None, lambda: ModelProviderListApi().get()),
+            (
+                "GET",
+                "/?credential_id=not-a-uuid",
+                None,
+                lambda: ModelProviderCredentialApi().get(provider="openai"),
+            ),
+            ("POST", "/", EMPTY_BODY, lambda: ModelProviderCredentialApi().post(provider="openai")),
+            ("PUT", "/", EMPTY_BODY, lambda: ModelProviderCredentialApi().put(provider="openai")),
+            ("DELETE", "/", EMPTY_BODY, lambda: ModelProviderCredentialApi().delete(provider="openai")),
+            ("POST", "/", EMPTY_BODY, lambda: ModelProviderCredentialSwitchApi().post(provider="openai")),
+            ("POST", "/", EMPTY_BODY, lambda: ModelProviderValidateApi().post(provider="openai")),
+            ("POST", "/", EMPTY_BODY, lambda: PreferredProviderTypeUpdateApi().post(provider="openai")),
+        ],
+    )
+    def test_invalid_input_is_rejected_before_the_handler_runs(
+        self,
+        app: Flask,
+        verb: str,
+        url: str,
+        body: dict[str, object] | None,
+        call: Callable[[], object],
+    ) -> None:
+        account = make_account()
+        account.role = TenantAccountRole.OWNER
+
+        with (
+            app.test_request_context(url, method=verb, json=body),
+            config_overrides_context(LOGIN_DISABLED=True, RBAC_ENABLED=False),
+            patch("controllers.console.wraps._is_setup_completed", return_value=True),
+            patch(
+                "controllers.console.wraps.current_account_with_tenant",
+                return_value=(account, "tenant1"),
+            ),
+            patch("libs.login.current_user", SimpleNamespace(_get_current_object=lambda: account)),
+            patch("controllers.console.workspace.model_providers.ModelProviderService") as service,
+        ):
+            g._login_user = account
+            with pytest.raises(UnprocessableEntity) as exc_info:
+                call()
+
+        assert exc_info.value.code == 422
+        service.assert_not_called()