opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-26 06:08:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

test: use sqlite3 session in test_build_from_mapping (#38769)

Browse Source
This commit is contained in:
Asuka Minato 2026-07-21 17:21:36 +09:00 committed by GitHub
parent a5194ff17b
commit 150613cc05
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 123 additions and 124 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

247
api/tests/unit_tests/factories/test_build_from_mapping.py
View File

@ -1,36 +1,28 @@
import uuid
from unittest.mock import MagicMock, patch
from collections.abc import Iterator
from dataclasses import dataclass
from datetime import UTC, datetime
from unittest.mock import patch
import pytest
from httpx import Response
from sqlalchemy import Engine
from sqlalchemy.orm import Session, sessionmaker
from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom
from core.app.file_access import DatabaseFileAccessController, FileAccessScope, bind_file_access_scope
from core.workflow.file_reference import build_file_reference, parse_file_reference, resolve_file_record_id
from extensions.storage.storage_type import StorageType
from factories.file_factory.builders import build_from_mapping as _build_from_mapping
from graphon.file import File, FileTransferMethod, FileType, FileUploadConfig
from models import ToolFile, UploadFile
def _make_session_ctx_mock(scalar_return=None):
"""Return a mock usable as the ``session_factory.create_session()`` context manager.
Patch ``factories.file_factory.builders.session_factory`` and set
``mock_sf.create_session.return_value = <this mock>`` to intercept DB calls
without requiring a live Flask app or database engine.
"""
session = MagicMock()
session.__enter__.return_value = session
session.__exit__.return_value = False
session.scalar.return_value = scalar_return
return session
from models import CreatorUserRole, ToolFile, UploadFile
# Test Data
TEST_TENANT_ID = "test_tenant_id"
TEST_UPLOAD_FILE_ID = str(uuid.uuid4())
TEST_TOOL_FILE_ID = str(uuid.uuid4())
TEST_REMOTE_URL = "http://example.com/test.jpg"
TEST_END_USER_ID = "end-user-id"
TEST_ACCESS_CONTROLLER = DatabaseFileAccessController()
# Test Config
@ -52,39 +44,52 @@ def build_from_mapping(*, mapping, tenant_id, config=None, strict_type_validatio
)
@dataclass(frozen=True)
class FileRecords:
session: Session
upload_file: UploadFile
tool_file: ToolFile
# Fixtures
@pytest.fixture
def mock_upload_file():
mock = MagicMock(spec=UploadFile)
mock.id = TEST_UPLOAD_FILE_ID
mock.tenant_id = TEST_TENANT_ID
mock.name = "test.jpg"
mock.extension = "jpg"
mock.mime_type = "image/jpeg"
mock.source_url = TEST_REMOTE_URL
mock.size = 1024
mock.key = "test_key"
session = _make_session_ctx_mock(scalar_return=mock)
with patch("factories.file_factory.builders.session_factory") as mock_sf:
mock_sf.create_session.return_value = session
# yield session.scalar so callers can inspect call_args and mutate return_value
yield session.scalar
@pytest.fixture(autouse=True)
def file_records(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> Iterator[FileRecords]:
"""Persist authorized file rows and bind builder-owned sessions to SQLite."""
UploadFile.metadata.create_all(sqlite_engine, tables=[UploadFile.__table__, ToolFile.__table__])
sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False)
monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker)
upload_file = UploadFile(
tenant_id=TEST_TENANT_ID,
storage_type=StorageType.LOCAL,
key="test_key",
name="test.jpg",
size=1024,
extension="jpg",
mime_type="image/jpeg",
created_by_role=CreatorUserRole.END_USER,
created_by=TEST_END_USER_ID,
created_at=datetime(2026, 1, 1, tzinfo=UTC),
used=False,
source_url=TEST_REMOTE_URL,
)
upload_file.id = TEST_UPLOAD_FILE_ID
tool_file = ToolFile(
user_id=TEST_END_USER_ID,
tenant_id=TEST_TENANT_ID,
conversation_id=None,
file_key="tool_file.pdf",
mimetype="application/pdf",
original_url="http://example.com/tool.pdf",
name="tool_file.pdf",
size=2048,
)
tool_file.id = TEST_TOOL_FILE_ID
@pytest.fixture
def mock_tool_file():
mock = MagicMock(spec=ToolFile)
mock.id = TEST_TOOL_FILE_ID
mock.tenant_id = TEST_TENANT_ID
mock.name = "tool_file.pdf"
mock.file_key = "tool_file.pdf"
mock.mimetype = "application/pdf"
mock.original_url = "http://example.com/tool.pdf"
mock.size = 2048
session = _make_session_ctx_mock(scalar_return=mock)
with patch("factories.file_factory.builders.session_factory") as mock_sf:
mock_sf.create_session.return_value = session
yield mock
with sqlite_session_maker() as session:
session.add_all([upload_file, tool_file])
session.commit()
yield FileRecords(session=session, upload_file=upload_file, tool_file=tool_file)
@pytest.fixture
@ -122,7 +127,7 @@ def tool_file_mapping(file_type="document"):
# Tests
def test_build_from_mapping_backward_compatibility(mock_upload_file):
def test_build_from_mapping_backward_compatibility():
mapping = local_file_mapping(file_type="image")
file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
assert isinstance(file, File)
@ -133,7 +138,7 @@ def test_build_from_mapping_backward_compatibility(mock_upload_file):
assert file.storage_key == "test_key"
def test_build_from_mapping_accepts_opaque_reference_for_local_file(mock_upload_file):
def test_build_from_mapping_accepts_opaque_reference_for_local_file():
mapping = {
"transfer_method": "local_file",
"reference": build_file_reference(record_id=TEST_UPLOAD_FILE_ID, storage_key="test_key"),
@ -148,7 +153,7 @@ def test_build_from_mapping_accepts_opaque_reference_for_local_file(mock_upload_
assert resolve_file_record_id(file.reference) == TEST_UPLOAD_FILE_ID
def test_build_from_mapping_accepts_opaque_related_id_for_tool_file(mock_tool_file):
def test_build_from_mapping_accepts_opaque_related_id_for_tool_file():
mapping = {
"transfer_method": "tool_file",
"related_id": build_file_reference(record_id=TEST_TOOL_FILE_ID, storage_key="tool_file.pdf"),
@ -165,10 +170,11 @@ def test_build_from_mapping_accepts_opaque_related_id_for_tool_file(mock_tool_fi
assert file.storage_key == "tool_file.pdf"
def test_build_from_mapping_prefers_tool_filename_extension_over_mimetype(mock_tool_file):
mock_tool_file.name = "report.docx"
mock_tool_file.file_key = "tools/test_tenant_id/file.bin"
mock_tool_file.mimetype = "application/octet-stream"
def test_build_from_mapping_prefers_tool_filename_extension_over_mimetype(file_records: FileRecords):
file_records.tool_file.name = "report.docx"
file_records.tool_file.file_key = "tools/test_tenant_id/file.bin"
file_records.tool_file.mimetype = "application/octet-stream"
file_records.session.commit()
mapping = tool_file_mapping(file_type="document")
file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
@ -186,7 +192,7 @@ def test_build_from_mapping_prefers_tool_filename_extension_over_mimetype(mock_t
("document", False, "Detected file type does not match"),
],
)
def test_build_from_local_file_strict_validation(mock_upload_file, file_type, should_pass, expected_error):
def test_build_from_local_file_strict_validation(file_type, should_pass, expected_error):
mapping = local_file_mapping(file_type=file_type)
if should_pass:
file = build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID, strict_type_validation=True)
@ -203,7 +209,7 @@ def test_build_from_local_file_strict_validation(mock_upload_file, file_type, sh
("image", False, "Detected file type does not match"),
],
)
def test_build_from_tool_file_strict_validation(mock_tool_file, file_type, should_pass, expected_error):
def test_build_from_tool_file_strict_validation(file_type, should_pass, expected_error):
"""Strict type validation for tool_file."""
mapping = tool_file_mapping(file_type=file_type)
if should_pass:
@ -282,27 +288,27 @@ def test_build_from_remote_url_without_strict_validation(mock_http_head):
assert file.filename == "remote_test.jpg"
def test_tool_file_not_found():
def test_tool_file_not_found(file_records: FileRecords):
"""Test ToolFile not found in database."""
session = _make_session_ctx_mock(scalar_return=None)
with patch("factories.file_factory.builders.session_factory") as mock_sf:
mock_sf.create_session.return_value = session
mapping = tool_file_mapping()
with pytest.raises(ValueError, match=f"ToolFile {TEST_TOOL_FILE_ID} not found"):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
file_records.session.delete(file_records.tool_file)
file_records.session.commit()
mapping = tool_file_mapping()
with pytest.raises(ValueError, match=f"ToolFile {TEST_TOOL_FILE_ID} not found"):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
def test_local_file_not_found():
def test_local_file_not_found(file_records: FileRecords):
"""Test UploadFile not found in database."""
session = _make_session_ctx_mock(scalar_return=None)
with patch("factories.file_factory.builders.session_factory") as mock_sf:
mock_sf.create_session.return_value = session
mapping = local_file_mapping()
with pytest.raises(ValueError, match="Invalid upload file"):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
file_records.session.delete(file_records.upload_file)
file_records.session.commit()
mapping = local_file_mapping()
with pytest.raises(ValueError, match="Invalid upload file"):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
def test_build_without_type_specification(mock_upload_file):
def test_build_without_type_specification():
"""Test the situation where no file type is specified"""
mapping = {
"transfer_method": "local_file",
@ -321,7 +327,7 @@ def test_build_without_type_specification(mock_upload_file):
("video", False, "File validation failed"),
],
)
def test_file_validation_with_config(mock_upload_file, file_type, should_pass, expected_error):
def test_file_validation_with_config(file_type, should_pass, expected_error):
"""Test the validation of files and configurations"""
mapping = local_file_mapping(file_type=file_type)
if should_pass:
@ -354,73 +360,65 @@ def test_invalid_uuid_format():
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
def test_tenant_mismatch():
def test_tenant_mismatch(file_records: FileRecords):
"""Test that tenant mismatch raises security error."""
# Create a mock upload file with a different tenant_id
mock_file = MagicMock(spec=UploadFile)
mock_file.id = TEST_UPLOAD_FILE_ID
mock_file.tenant_id = "different_tenant_id"
mock_file.name = "test.jpg"
mock_file.extension = "jpg"
mock_file.mime_type = "image/jpeg"
mock_file.source_url = TEST_REMOTE_URL
mock_file.size = 1024
mock_file.key = "test_key"
file_records.upload_file.tenant_id = "different_tenant_id"
file_records.session.commit()
# Mock the database query to return None (no file found for this tenant)
session = _make_session_ctx_mock(scalar_return=None)
with patch("factories.file_factory.builders.session_factory") as mock_sf:
mock_sf.create_session.return_value = session
mapping = local_file_mapping()
with pytest.raises(ValueError, match="Invalid upload file"):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
mapping = local_file_mapping()
with pytest.raises(ValueError, match="Invalid upload file"):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID)
def test_build_from_mapping_scopes_upload_file_to_end_user(mock_upload_file):
def test_build_from_mapping_scopes_upload_file_to_end_user():
scope = FileAccessScope(
tenant_id=TEST_TENANT_ID,
user_id="end-user-id",
user_id=TEST_END_USER_ID,
user_from=UserFrom.END_USER,
invoke_from=InvokeFrom.WEB_APP,
)
with bind_file_access_scope(scope):
build_from_mapping(mapping=local_file_mapping(), tenant_id=TEST_TENANT_ID)
file = build_from_mapping(mapping=local_file_mapping(), tenant_id=TEST_TENANT_ID)
stmt = mock_upload_file.call_args.args[0]
whereclause = str(stmt.whereclause)
assert "upload_files.created_by_role" in whereclause
assert "upload_files.created_by" in whereclause
assert resolve_file_record_id(file.reference) == TEST_UPLOAD_FILE_ID
unauthorized_scope = FileAccessScope(
tenant_id=TEST_TENANT_ID,
user_id="different-end-user",
user_from=UserFrom.END_USER,
invoke_from=InvokeFrom.WEB_APP,
)
with bind_file_access_scope(unauthorized_scope):
with pytest.raises(ValueError, match="Invalid upload file"):
build_from_mapping(mapping=local_file_mapping(), tenant_id=TEST_TENANT_ID)
def test_build_from_mapping_scopes_tool_file_to_end_user():
tool_file = MagicMock(spec=ToolFile)
tool_file.id = TEST_TOOL_FILE_ID
tool_file.tenant_id = TEST_TENANT_ID
tool_file.name = "tool_file.pdf"
tool_file.file_key = "tool_file.pdf"
tool_file.mimetype = "application/pdf"
tool_file.original_url = "http://example.com/tool.pdf"
tool_file.size = 2048
scope = FileAccessScope(
tenant_id=TEST_TENANT_ID,
user_id="end-user-id",
user_id=TEST_END_USER_ID,
user_from=UserFrom.END_USER,
invoke_from=InvokeFrom.WEB_APP,
)
session = _make_session_ctx_mock(scalar_return=tool_file)
with patch("factories.file_factory.builders.session_factory") as mock_sf:
mock_sf.create_session.return_value = session
with bind_file_access_scope(scope):
with bind_file_access_scope(scope):
file = build_from_mapping(mapping=tool_file_mapping(), tenant_id=TEST_TENANT_ID)
assert resolve_file_record_id(file.reference) == TEST_TOOL_FILE_ID
unauthorized_scope = FileAccessScope(
tenant_id=TEST_TENANT_ID,
user_id="different-end-user",
user_from=UserFrom.END_USER,
invoke_from=InvokeFrom.WEB_APP,
)
with bind_file_access_scope(unauthorized_scope):
with pytest.raises(ValueError, match=f"ToolFile {TEST_TOOL_FILE_ID} not found"):
build_from_mapping(mapping=tool_file_mapping(), tenant_id=TEST_TENANT_ID)
stmt = session.scalar.call_args.args[0]
whereclause = str(stmt.whereclause)
assert "tool_files.user_id" in whereclause
def test_disallowed_file_types(mock_upload_file):
def test_disallowed_file_types():
"""Test that disallowed file types are rejected."""
# Config that only allows image and document types
restricted_config = FileUploadConfig(
@ -433,12 +431,12 @@ def test_disallowed_file_types(mock_upload_file):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID, config=restricted_config)
def test_disallowed_extensions(mock_upload_file):
def test_disallowed_extensions(file_records: FileRecords):
"""Test that disallowed file extensions are rejected for custom type."""
# Mock a file with .exe extension
mock_upload_file.return_value.extension = "exe"
mock_upload_file.return_value.name = "malicious.exe"
mock_upload_file.return_value.mime_type = "application/x-msdownload"
file_records.upload_file.extension = "exe"
file_records.upload_file.name = "malicious.exe"
file_records.upload_file.mime_type = "application/x-msdownload"
file_records.session.commit()
# Config that only allows specific extensions for custom files
restricted_config = FileUploadConfig(
@ -456,11 +454,12 @@ def test_disallowed_extensions(mock_upload_file):
build_from_mapping(mapping=mapping, tenant_id=TEST_TENANT_ID, config=restricted_config)
def test_custom_file_type_uses_extension_validation_under_strict_mode(mock_upload_file):
def test_custom_file_type_uses_extension_validation_under_strict_mode(file_records: FileRecords):
"""Custom form uploads are classified by the configured extension list."""
mock_upload_file.return_value.extension = "txt"
mock_upload_file.return_value.name = "notes.txt"
mock_upload_file.return_value.mime_type = "text/plain"
file_records.upload_file.extension = "txt"
file_records.upload_file.name = "notes.txt"
file_records.upload_file.mime_type = "text/plain"
file_records.session.commit()
custom_config = FileUploadConfig(
allowed_file_types=[FileType.CUSTOM],