cetdev/dify
0
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

test: use SQLite sessions in core datasource (#39104)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Asuka Minato
2026-07-28 05:22:18 +00:00
committed by GitHub
co-authored by autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
parent b2d54cb2e9
commit 9e90b32991
2 changed files with 108 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/tests/unit_tests/core/datasource/test_datasource_manager.py
+36 -28
View File
@@ -1,10 +1,13 @@
import types
from collections.abc import Generator
from collections.abc import Generator, Iterator
import pytest
from pytest_mock import MockerFixture
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker
from contexts.wrapper import RecyclableContextVar
from core.datasource import datasource_manager as datasource_manager_module
from core.datasource.datasource_manager import DatasourceManager
from core.datasource.entities.datasource_entities import DatasourceMessage, DatasourceProviderType
from core.datasource.errors import DatasourceProviderNotFoundError
@@ -12,6 +15,34 @@ from core.workflow.file_reference import parse_file_reference
from graphon.enums import WorkflowNodeExecutionStatus
from graphon.file import File, FileTransferMethod, FileType
from graphon.node_events import StreamChunkEvent, StreamCompletedEvent
from models.base import TypeBase
from models.tools import ToolFile
@pytest.fixture
def tool_file_session(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]:
"""Bind datasource-owned lookups to a SQLite ToolFile table."""
TypeBase.metadata.create_all(sqlite_engine, tables=[TypeBase.metadata.tables[ToolFile.__tablename__]])
session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
monkeypatch.setattr(datasource_manager_module.session_factory, "create_session", session_maker)
with session_maker() as session:
yield session
def _persist_tool_file(session: Session, *, file_id: str, tenant_id: str) -> ToolFile:
tool_file = ToolFile(
user_id="user-1",
tenant_id=tenant_id,
conversation_id=None,
file_key="files/image.png",
mimetype="image/png",
name="image.png",
size=10,
)
tool_file.id = file_id
session.add(tool_file)
session.commit()
return tool_file
def _gen_messages_text_only(text: str) -> Generator[DatasourceMessage, None, None]:
@@ -373,7 +404,8 @@ def test_stream_node_events_emits_events_online_document(mocker: MockerFixture):
assert events[-1].node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED
def test_stream_node_events_builds_file_and_variables_from_messages(mocker: MockerFixture):
def test_stream_node_events_builds_file_and_variables_from_messages(mocker: MockerFixture, tool_file_session: Session):
_persist_tool_file(tool_file_session, file_id="tool_file_1", tenant_id="t1")
mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_messages_text_only("ignored"))
def _transformed(**_kwargs):
@@ -418,19 +450,6 @@ def test_stream_node_events_builds_file_and_variables_from_messages(mocker: Mock
side_effect=_transformed,
)
fake_tool_file = types.SimpleNamespace(mimetype="image/png")
class _Session:
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def scalar(self, _stmt):
return fake_tool_file
mocker.patch("core.datasource.datasource_manager.session_factory.create_session", return_value=_Session())
mocker.patch("core.datasource.datasource_manager.get_file_type_by_mime_type", return_value=FileType.IMAGE)
built = File(
file_type=FileType.IMAGE,
@@ -481,7 +500,8 @@ def test_stream_node_events_builds_file_and_variables_from_messages(mocker: Mock
assert events[-1].node_run_result.outputs["x"] == 1
def test_stream_node_events_raises_when_toolfile_missing(mocker: MockerFixture):
def test_stream_node_events_raises_when_toolfile_missing(mocker: MockerFixture, tool_file_session: Session):
_persist_tool_file(tool_file_session, file_id="missing", tenant_id="other-tenant")
mocker.patch.object(DatasourceManager, "stream_online_results", return_value=_gen_messages_text_only("ignored"))
def _transformed(**_kwargs):
@@ -496,18 +516,6 @@ def test_stream_node_events_raises_when_toolfile_missing(mocker: MockerFixture):
side_effect=_transformed,
)
class _Session:
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def scalar(self, _stmt):
return None
mocker.patch("core.datasource.datasource_manager.session_factory.create_session", return_value=_Session())
with pytest.raises(ValueError, match="ToolFile not found for file_id=missing, tenant_id=t1"):
list(
DatasourceManager.stream_node_events(
api/tests/unit_tests/core/datasource/test_notion_provider.py
+72 -11
View File
@@ -14,18 +14,64 @@ Tests follow the Arrange-Act-Assert pattern for clarity.
"""
import json
from collections.abc import Iterator
from dataclasses import dataclass
from typing import Any
from unittest.mock import Mock, patch
from uuid import uuid4
import httpx
import pytest
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session
from core.datasource.entities.datasource_entities import DatasourceProviderType
from core.datasource.online_document.online_document_provider import (
OnlineDocumentDatasourcePluginProviderController,
)
from core.rag.extractor import notion_extractor as notion_extractor_module
from core.rag.extractor.notion_extractor import NotionExtractor
from core.rag.models.document import Document
from models.base import TypeBase
from models.dataset import Document as DocumentModel
from models.enums import DataSourceType, DocumentCreatedFrom
@dataclass(frozen=True)
class _Database:
"""Expose the real SQLite session used by the extractor update."""
session: Session
@pytest.fixture
def database(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Iterator[_Database]:
"""Bind a real session for Notion document metadata persistence."""
TypeBase.metadata.create_all(sqlite_engine, tables=[DocumentModel.__table__])
with Session(sqlite_engine, expire_on_commit=False) as session:
database = _Database(session)
monkeypatch.setattr(notion_extractor_module, "db", database)
yield database
@pytest.fixture
def persisted_document(database: _Database) -> DocumentModel:
document = DocumentModel(
id=str(uuid4()),
tenant_id=str(uuid4()),
dataset_id=str(uuid4()),
position=1,
data_source_type=DataSourceType.NOTION_IMPORT,
data_source_info=json.dumps({"last_edited_time": "2024-01-01T00:00:00.000Z"}),
batch="batch",
name="Notion page",
created_from=DocumentCreatedFrom.WEB,
created_by=str(uuid4()),
)
database.session.add(document)
database.session.commit()
return document
class TestNotionExtractorAuthentication:
@@ -763,9 +809,14 @@ class TestNotionExtractorLastEditedTime:
call_args = mock_request.call_args
assert "databases/database-789" in call_args[0][1]
@patch("core.rag.extractor.notion_extractor.db")
@patch("httpx.request")
def test_update_last_edited_time(self, mock_request, mock_db, extractor_page, mock_document_model):
def test_update_last_edited_time(
self,
mock_request: Mock,
extractor_page: NotionExtractor,
database: _Database,
persisted_document: DocumentModel,
):
"""Test updating document model with last edited time."""
# Arrange
mock_response = Mock()
@@ -777,11 +828,11 @@ class TestNotionExtractorLastEditedTime:
mock_request.return_value = mock_response
# Act
extractor_page.update_last_edited_time(mock_document_model)
extractor_page.update_last_edited_time(persisted_document)
# Assert
assert mock_document_model.data_source_info_dict["last_edited_time"] == "2024-11-27T18:00:00.000Z"
mock_db.session.commit.assert_called_once()
database.session.expire(persisted_document)
assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T18:00:00.000Z"
def test_update_last_edited_time_no_document(self, extractor_page):
"""Test update_last_edited_time with None document model."""
@@ -807,9 +858,10 @@ class TestNotionExtractorIntegration:
mock_doc.data_source_info_dict = {"last_edited_time": "2024-01-01T00:00:00.000Z"}
return mock_doc
@patch("core.rag.extractor.notion_extractor.db")
@patch("httpx.request")
def test_extract_page_complete_workflow(self, mock_request, mock_db, mock_document_model):
def test_extract_page_complete_workflow(
self, mock_request: Mock, database: _Database, persisted_document: DocumentModel
):
"""Test complete page extraction workflow."""
# Arrange
extractor = NotionExtractor(
@@ -818,7 +870,7 @@ class TestNotionExtractorIntegration:
notion_page_type="page",
tenant_id="tenant-789",
notion_access_token="test-token",
document_model=mock_document_model,
document_model=persisted_document,
)
# Mock last edited time request
@@ -869,11 +921,18 @@ class TestNotionExtractorIntegration:
assert isinstance(documents[0], Document)
assert "# Test Page" in documents[0].page_content
assert "Test content" in documents[0].page_content
database.session.expire(persisted_document)
assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T20:00:00.000Z"
@patch("core.rag.extractor.notion_extractor.db")
@patch("httpx.post")
@patch("httpx.request")
def test_extract_database_complete_workflow(self, mock_request, mock_post, mock_db, mock_document_model):
def test_extract_database_complete_workflow(
self,
mock_request: Mock,
mock_post: Mock,
database: _Database,
persisted_document: DocumentModel,
):
"""Test complete database extraction workflow."""
# Arrange
extractor = NotionExtractor(
@@ -882,7 +941,7 @@ class TestNotionExtractorIntegration:
notion_page_type="database",
tenant_id="tenant-789",
notion_access_token="test-token",
document_model=mock_document_model,
document_model=persisted_document,
)
# Mock last edited time request
@@ -921,6 +980,8 @@ class TestNotionExtractorIntegration:
assert isinstance(documents[0], Document)
assert "Name:Item 1" in documents[0].page_content
assert "Status:Active" in documents[0].page_content
database.session.expire(persisted_document)
assert persisted_document.data_source_info_dict["last_edited_time"] == "2024-11-27T20:00:00.000Z"
def test_extract_invalid_page_type(self):
"""Test extract with invalid page type."""