test: migrate RAG residual entities to real ORM models (#40629)

api/tests/unit_tests/core/rag/extractor/test_extract_processor.py

@@ -1,12 +1,35 @@
+from datetime import datetime
 from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
+from sqlalchemy.orm import Session
 
 import core.rag.extractor.extract_processor as processor_module
 from core.rag.extractor.entity.datasource_type import DatasourceType
 from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.models.document import Document
+from extensions.storage.storage_type import StorageType
+from models.enums import CreatorUserRole
+from models.model import UploadFile
+
+
+def _upload_file(*, key: str, file_id: str = "upload-file-1") -> UploadFile:
+    upload_file = UploadFile(
+        tenant_id="tenant-1",
+        storage_type=StorageType.LOCAL,
+        key=key,
+        name=Path(key).name,
+        size=1,
+        extension=Path(key).suffix.lstrip("."),
+        mime_type="application/octet-stream",
+        created_by_role=CreatorUserRole.ACCOUNT,
+        created_by="user-1",
+        created_at=datetime(2025, 1, 1),
+        used=True,
+    )
+    upload_file.id = file_id
+    return upload_file
 
 
 class _ExtractorFactory:
@@ -68,7 +91,7 @@ class TestExtractProcessorLoaders:
             ],
         )
 
-        upload_file = SimpleNamespace(key="file.txt")
+        upload_file = _upload_file(key="file.txt")
 
         docs = ExtractProcessor.load_from_upload_file(upload_file=upload_file, return_text=False)
         text = ExtractProcessor.load_from_upload_file(upload_file=upload_file, return_text=True)
@@ -154,12 +177,7 @@ class TestExtractProcessorFileRouting:
 
         setting = SimpleNamespace(
             datasource_type=DatasourceType.FILE,
-            upload_file=SimpleNamespace(
-                id="upload-file-1",
-                key=f"uploaded{extension}",
-                tenant_id="tenant-1",
-                created_by="user-1",
-            ),
+            upload_file=_upload_file(key=f"uploaded{extension}"),
         )
 
         docs = ExtractProcessor.extract(setting, is_automatic=is_automatic, session=session)
@@ -229,13 +247,13 @@ class TestExtractProcessorFileRouting:
 
     @pytest.mark.parametrize("extension", [".pdf", ".docx"])
     def test_extract_passes_session_to_database_backed_file_extractors(
-        self, monkeypatch: pytest.MonkeyPatch, extension: str
+        self, monkeypatch: pytest.MonkeyPatch, extension: str, unbound_session: Session
     ):
-        session = object()
+        _, _, kwargs = self._run_extract_for_extension(
+            monkeypatch, extension, etl_type="SelfHosted", session=unbound_session
+        )
 
-        _, _, kwargs = self._run_extract_for_extension(monkeypatch, extension, etl_type="SelfHosted", session=session)
-
-        assert kwargs["session"] is session
+        assert kwargs["session"] is unbound_session
 
     def test_extract_requires_upload_file_when_file_path_not_provided(self):
         setting = SimpleNamespace(datasource_type=DatasourceType.FILE, upload_file=None)

api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py

@@ -1,14 +1,22 @@
 """Focused tests for attachment-aware dataset retrieval entry behavior."""
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
+
+from sqlalchemy.orm import Session
 
 from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
 from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest
+from models.dataset import Dataset
 
 
-def test_knowledge_retrieval_allows_attachment_only_requests() -> None:
+def test_knowledge_retrieval_allows_attachment_only_requests(unbound_session: Session) -> None:
     retrieval = DatasetRetrieval()
-    available_dataset = MagicMock(id="dataset-1")
+    available_dataset = Dataset(
+        id="dataset-1",
+        tenant_id="tenant-1",
+        name="Dataset",
+        created_by="user-1",
+    )
 
     request = KnowledgeRetrievalRequest(
         tenant_id="tenant-1",
@@ -30,7 +38,7 @@ def test_knowledge_retrieval_allows_attachment_only_requests() -> None:
         patch.object(retrieval, "_get_available_datasets", return_value=[available_dataset]),
         patch.object(retrieval, "multiple_retrieve", return_value=[]) as mock_multiple,
     ):
-        result = retrieval.knowledge_retrieval(MagicMock(), request)
+        result = retrieval.knowledge_retrieval(unbound_session, request)
 
     assert result == []
     mock_multiple.assert_called_once()