test: use SQLite sessions in services core (#39112)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/tests/unit_tests/services/test_knowledge_retrieval_inner_service.py

@@ -3,8 +3,12 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from sqlalchemy.orm import Session
 
 from core.workflow.nodes.knowledge_retrieval.retrieval import Source, SourceMetadata
+from models.dataset import Dataset
+from models.enums import AppStatus
+from models.model import App, AppMode
 from services.entities.knowledge_retrieval_inner import InnerKnowledgeRetrieveRequest
 from services.errors.knowledge_retrieval import (
     InnerKnowledgeRetrieveAppNotFoundError,
@@ -14,17 +18,55 @@ from services.errors.knowledge_retrieval import (
 )
 from services.knowledge_retrieval_inner_service import InnerKnowledgeRetrievalService
 
+TENANT_ID = "11111111-1111-1111-1111-111111111111"
+OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222"
+USER_ID = "33333333-3333-3333-3333-333333333333"
+APP_ID = "44444444-4444-4444-4444-444444444444"
+DATASET_1_ID = "55555555-5555-5555-5555-555555555555"
+DATASET_2_ID = "66666666-6666-6666-6666-666666666666"
+
+
+def _app(*, tenant_id: str = TENANT_ID) -> App:
+    return App(
+        id=APP_ID,
+        tenant_id=tenant_id,
+        name="Test App",
+        description="",
+        mode=AppMode.WORKFLOW,
+        status=AppStatus.NORMAL,
+        enable_site=False,
+        enable_api=False,
+        max_active_requests=None,
+    )
+
+
+def _dataset(*, dataset_id: str, tenant_id: str = TENANT_ID, enable_api: bool = True) -> Dataset:
+    return Dataset(
+        id=dataset_id,
+        tenant_id=tenant_id,
+        name=f"Dataset {dataset_id[-1]}",
+        description="",
+        created_by=USER_ID,
+        enable_api=enable_api,
+    )
+
+
+def _persist_state(sqlite_session: Session, *models: App | Dataset) -> None:
+    sqlite_session.add_all(models)
+    sqlite_session.commit()
+    sqlite_session.expunge_all()
+
 
 def _build_request(**overrides):
     payload = {
         "caller": {
-            "tenant_id": "tenant-1",
+            "tenant_id": TENANT_ID,
-            "user_id": "user-1",
+            "user_id": USER_ID,
-            "app_id": "app-1",
+            "app_id": APP_ID,
             "user_from": "account",
             "invoke_from": "workflow",
         },
-        "dataset_ids": ["dataset-1", "dataset-2"],
+        "dataset_ids": [DATASET_1_ID, DATASET_2_ID],
         "query": "how to reset password",
         "retrieval": {
             "mode": "multiple",
@@ -73,15 +115,20 @@ def _build_source() -> Source:
 
 
 class TestInnerKnowledgeRetrievalService:
+    @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True)
     @patch("services.knowledge_retrieval_inner_service.DatasetRetrieval")
-    def test_retrieve_maps_multiple_request_and_skips_enable_api_check(self, mock_rag_cls):
+    def test_retrieve_maps_multiple_request_and_skips_enable_api_check(
+        self,
+        mock_rag_cls,
+        sqlite_session: Session,
+    ):
         request = _build_request()
-        mock_session = MagicMock()
+        _persist_state(
-        mock_app = MagicMock(id="app-1", tenant_id="tenant-1")
+            sqlite_session,
-        dataset_1 = MagicMock(id="dataset-1", tenant_id="tenant-1", enable_api=False)
+            _app(),
-        dataset_2 = MagicMock(id="dataset-2", tenant_id="tenant-1", enable_api=True)
+            _dataset(dataset_id=DATASET_1_ID, enable_api=False),
-        mock_session.scalar.return_value = mock_app
+            _dataset(dataset_id=DATASET_2_ID, enable_api=True),
-        mock_session.scalars.return_value.all.return_value = [dataset_1, dataset_2]
+        )
 
         rag = MagicMock()
         rag.knowledge_retrieval.return_value = [_build_source()]
@@ -103,13 +150,13 @@ class TestInnerKnowledgeRetrievalService:
         }
         mock_rag_cls.return_value = rag
 
-        response = InnerKnowledgeRetrievalService().retrieve(request, mock_session)
+        response = InnerKnowledgeRetrievalService().retrieve(request, sqlite_session)
 
         rag_request = rag.knowledge_retrieval.call_args.kwargs["request"]
-        assert rag_request.tenant_id == "tenant-1"
+        assert rag_request.tenant_id == TENANT_ID
-        assert rag_request.app_id == "app-1"
+        assert rag_request.app_id == APP_ID
-        assert rag_request.user_id == "user-1"
+        assert rag_request.user_id == USER_ID
-        assert rag_request.dataset_ids == ["dataset-1", "dataset-2"]
+        assert rag_request.dataset_ids == [DATASET_1_ID, DATASET_2_ID]
         assert rag_request.query == "how to reset password"
         assert rag_request.retrieval_mode == "multiple"
         assert rag_request.top_k == 4
@@ -127,11 +174,14 @@ class TestInnerKnowledgeRetrievalService:
         assert rag_request.attachment_ids == ["attachment-1"]
         assert response.results[0].title == "FAQ.md"
         assert response.usage.currency == "USD"
+        assert rag.knowledge_retrieval.call_args.kwargs["session"] is sqlite_session
+        assert sqlite_session.in_transaction()
 
+    @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True)
     @patch("services.knowledge_retrieval_inner_service.DatasetRetrieval")
-    def test_retrieve_maps_single_request(self, mock_rag_cls):
+    def test_retrieve_maps_single_request(self, mock_rag_cls, sqlite_session: Session):
         request = _build_request(
-            dataset_ids=["dataset-1"],
+            dataset_ids=[DATASET_1_ID],
             retrieval={
                 "mode": "single",
                 "model": {
@@ -152,9 +202,7 @@ class TestInnerKnowledgeRetrievalService:
             },
             attachment_ids=[],
         )
-        mock_session = MagicMock()
+        _persist_state(sqlite_session, _app(), _dataset(dataset_id=DATASET_1_ID))
-        mock_session.scalar.return_value = MagicMock(id="app-1", tenant_id="tenant-1")
-        mock_session.scalars.return_value.all.return_value = [MagicMock(id="dataset-1", tenant_id="tenant-1")]
 
         rag = MagicMock()
         rag.knowledge_retrieval.return_value = []
@@ -174,7 +222,7 @@ class TestInnerKnowledgeRetrievalService:
         }
         mock_rag_cls.return_value = rag
 
-        InnerKnowledgeRetrievalService().retrieve(request, mock_session)
+        InnerKnowledgeRetrievalService().retrieve(request, sqlite_session)
 
         rag_request = rag.knowledge_retrieval.call_args.kwargs["request"]
         assert rag_request.retrieval_mode == "single"
@@ -185,36 +233,39 @@ class TestInnerKnowledgeRetrievalService:
         assert rag_request.metadata_filtering_mode == "automatic"
         assert rag_request.metadata_model_config is not None
         assert rag_request.metadata_model_config.provider == "openai"
+        assert sqlite_session.in_transaction()
 
-    def test_retrieve_raises_when_app_missing(self):
+    @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True)
-        mock_session = MagicMock()
+    def test_retrieve_raises_when_app_missing(self, sqlite_session: Session):
-        mock_session.scalar.return_value = None
-
         with pytest.raises(InnerKnowledgeRetrieveAppNotFoundError):
-            InnerKnowledgeRetrievalService().retrieve(_build_request(), mock_session)
+            InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session)
+        assert sqlite_session.in_transaction()
 
-    def test_retrieve_raises_when_app_belongs_to_other_tenant(self):
+    @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True)
-        mock_session = MagicMock()
+    def test_retrieve_raises_when_app_belongs_to_other_tenant(self, sqlite_session: Session):
-        mock_session.scalar.return_value = MagicMock(id="app-1", tenant_id="tenant-2")
+        _persist_state(sqlite_session, _app(tenant_id=OTHER_TENANT_ID))
 
         with pytest.raises(InnerKnowledgeRetrieveAppTenantMismatchError):
-            InnerKnowledgeRetrievalService().retrieve(_build_request(), mock_session)
+            InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session)
+        assert sqlite_session.in_transaction()
 
-    def test_retrieve_raises_when_dataset_missing(self):
+    @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True)
-        mock_session = MagicMock()
+    def test_retrieve_raises_when_dataset_missing(self, sqlite_session: Session):
-        mock_session.scalar.return_value = MagicMock(id="app-1", tenant_id="tenant-1")
+        _persist_state(sqlite_session, _app(), _dataset(dataset_id=DATASET_1_ID))
-        mock_session.scalars.return_value.all.return_value = [MagicMock(id="dataset-1", tenant_id="tenant-1")]
 
         with pytest.raises(InnerKnowledgeRetrieveDatasetNotFoundError):
-            InnerKnowledgeRetrievalService().retrieve(_build_request(), mock_session)
+            InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session)
+        assert sqlite_session.in_transaction()
 
-    def test_retrieve_raises_when_dataset_belongs_to_other_tenant(self):
+    @pytest.mark.parametrize("sqlite_session", [(App, Dataset)], indirect=True)
-        mock_session = MagicMock()
+    def test_retrieve_raises_when_dataset_belongs_to_other_tenant(self, sqlite_session: Session):
-        mock_session.scalar.return_value = MagicMock(id="app-1", tenant_id="tenant-1")
+        _persist_state(
-        mock_session.scalars.return_value.all.return_value = [
+            sqlite_session,
-            MagicMock(id="dataset-1", tenant_id="tenant-1"),
+            _app(),
-            MagicMock(id="dataset-2", tenant_id="tenant-2"),
+            _dataset(dataset_id=DATASET_1_ID),
-        ]
+            _dataset(dataset_id=DATASET_2_ID, tenant_id=OTHER_TENANT_ID),
+        )
 
         with pytest.raises(InnerKnowledgeRetrieveDatasetTenantMismatchError):
-            InnerKnowledgeRetrievalService().retrieve(_build_request(), mock_session)
+            InnerKnowledgeRetrievalService().retrieve(_build_request(), sqlite_session)
+        assert sqlite_session.in_transaction()