diff --git a/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py b/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py index 0b73d488fa0..369e63e57e3 100644 --- a/api/tests/unit_tests/core/rag/extractor/test_extract_processor.py +++ b/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) diff --git a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py index 1e41827d316..a5a91a01545 100644 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval_attachment_entry.py +++ b/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()