From 150613cc053052609d97f8b9c4a14a514e8b63e7 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 21 Jul 2026 17:21:36 +0900 Subject: [PATCH] test: use sqlite3 session in test_build_from_mapping (#38769) --- .../factories/test_build_from_mapping.py | 247 +++++++++--------- 1 file changed, 123 insertions(+), 124 deletions(-) diff --git a/api/tests/unit_tests/factories/test_build_from_mapping.py b/api/tests/unit_tests/factories/test_build_from_mapping.py index 90d92989c61..30d69144f45 100644 --- a/api/tests/unit_tests/factories/test_build_from_mapping.py +++ b/api/tests/unit_tests/factories/test_build_from_mapping.py @@ -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 = `` 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],