From c12d4662dbf87d767114b29d57ff08e17802594a Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 14 Jul 2026 15:39:33 +0900 Subject: [PATCH] test: use sqlite3 session in test_file_tenant_validator (#38756) --- .../agent_v2/test_file_tenant_validator.py | 191 +++++++++++++----- 1 file changed, 136 insertions(+), 55 deletions(-) diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_file_tenant_validator.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_file_tenant_validator.py index 2d25968fe71..80581a4630b 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_file_tenant_validator.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_file_tenant_validator.py @@ -8,34 +8,65 @@ the current tenant should simply return False. from __future__ import annotations -from unittest.mock import patch +from collections.abc import Iterator +from datetime import UTC, datetime import pytest +from sqlalchemy import Engine, event +from sqlalchemy.exc import SQLAlchemyError +from sqlalchemy.orm import Session, sessionmaker from core.workflow.nodes.agent_v2.file_tenant_validator import UploadFileTenantValidator +from extensions.storage.storage_type import StorageType from graphon.file import FileTransferMethod +from models import CreatorUserRole, ToolFile, UploadFile + +TABLES = (ToolFile, UploadFile) +TENANT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222" +USER_ID = "33333333-3333-3333-3333-333333333333" -def test_empty_inputs_return_false_without_db_hit(): +@pytest.fixture(autouse=True) +def _bind_sqlite_session_factory(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> None: + """Bind the validator's service-owned sessions to the isolated SQLite engine.""" + sqlite_session_maker = sessionmaker(bind=sqlite_engine, expire_on_commit=False) + monkeypatch.setattr("core.db.session_factory._session_maker", sqlite_session_maker) + + +@pytest.fixture +def executed_statements(sqlite_engine: Engine) -> Iterator[list[str]]: + statements: list[str] = [] + + def record_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", record_statement) + try: + yield statements + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_statement) + + +def test_empty_inputs_return_false_without_db_hit(executed_statements: list[str]): validator = UploadFileTenantValidator() - with patch("core.workflow.nodes.agent_v2.file_tenant_validator.session_factory") as factory: - assert ( - validator.is_accessible_file_mapping( - file_id="", - tenant_id="tenant-1", - transfer_method=FileTransferMethod.LOCAL_FILE, - ) - is False + assert ( + validator.is_accessible_file_mapping( + file_id="", + tenant_id="tenant-1", + transfer_method=FileTransferMethod.LOCAL_FILE, ) - assert ( - validator.is_accessible_file_mapping( - file_id="abc", - tenant_id="", - transfer_method=FileTransferMethod.LOCAL_FILE, - ) - is False + is False + ) + assert ( + validator.is_accessible_file_mapping( + file_id="abc", + tenant_id="", + transfer_method=FileTransferMethod.LOCAL_FILE, ) - factory.create_session.assert_not_called() + is False + ) + assert executed_statements == [] @pytest.mark.parametrize( @@ -49,29 +80,35 @@ def test_empty_inputs_return_false_without_db_hit(): "550e8400-e29b-41d4-a716-446655440000-trailing", ], ) -def test_non_uuid_file_ids_return_false_without_db_hit(bad_file_id: str): +def test_non_uuid_file_ids_return_false_without_db_hit(bad_file_id: str, executed_statements: list[str]): validator = UploadFileTenantValidator() - with patch("core.workflow.nodes.agent_v2.file_tenant_validator.session_factory") as factory: - assert ( - validator.is_accessible_file_mapping( - file_id=bad_file_id, - tenant_id="tenant-1", - transfer_method=FileTransferMethod.LOCAL_FILE, - ) - is False + assert ( + validator.is_accessible_file_mapping( + file_id=bad_file_id, + tenant_id="tenant-1", + transfer_method=FileTransferMethod.LOCAL_FILE, ) - factory.create_session.assert_not_called() + is False + ) + assert executed_statements == [] -def test_db_error_swallowed_and_returns_false(): +def test_db_error_swallowed_and_returns_false(sqlite_engine: Engine): """Any DB-level fault (timeout, dialect quirk, connection drop) must reject the file rather than crash the workflow node.""" - from sqlalchemy.exc import SQLAlchemyError - validator = UploadFileTenantValidator() valid_uuid = "550e8400-e29b-41d4-a716-446655440000" - with patch("core.workflow.nodes.agent_v2.file_tenant_validator.session_factory") as factory: - factory.create_session.return_value.__enter__.return_value.scalar.side_effect = SQLAlchemyError("boom") + rollbacks: list[bool] = [] + + def fail_query(*_args) -> None: + raise SQLAlchemyError("boom") + + def record_rollback(_connection) -> None: + rollbacks.append(True) + + event.listen(sqlite_engine, "before_cursor_execute", fail_query) + event.listen(sqlite_engine, "rollback", record_rollback) + try: assert ( validator.is_accessible_file_mapping( file_id=valid_uuid, @@ -80,29 +117,73 @@ def test_db_error_swallowed_and_returns_false(): ) is False ) + finally: + event.remove(sqlite_engine, "before_cursor_execute", fail_query) + event.remove(sqlite_engine, "rollback", record_rollback) + + assert rollbacks == [True] -def test_accessible_file_mapping_checks_transfer_method_family(): +@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True) +def test_accessible_file_mapping_checks_transfer_method_family(sqlite_session: Session): validator = UploadFileTenantValidator() - valid_uuid = "550e8400-e29b-41d4-a716-446655440000" - with patch("core.workflow.nodes.agent_v2.file_tenant_validator.session_factory") as factory: - factory.create_session.return_value.__enter__.return_value.scalar.return_value = None - assert ( - validator.is_accessible_file_mapping( - file_id=valid_uuid, - tenant_id="tenant-1", - transfer_method=FileTransferMethod.LOCAL_FILE, - ) - is False - ) + upload_file = UploadFile( + tenant_id=TENANT_ID, + storage_type=StorageType.LOCAL, + key="uploads/report.pdf", + name="report.pdf", + size=42, + extension="pdf", + mime_type="application/pdf", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=USER_ID, + created_at=datetime(2026, 1, 1, tzinfo=UTC), + used=False, + ) + tool_file = ToolFile( + user_id=USER_ID, + tenant_id=TENANT_ID, + conversation_id=None, + file_key="tools/chart.png", + mimetype="image/png", + name="chart.png", + size=99, + ) + sqlite_session.add_all([upload_file, tool_file]) + sqlite_session.commit() - with patch("core.workflow.nodes.agent_v2.file_tenant_validator.session_factory") as factory: - factory.create_session.return_value.__enter__.return_value.scalar.return_value = "tenant-1" - assert ( - validator.is_accessible_file_mapping( - file_id=valid_uuid, - tenant_id="tenant-1", - transfer_method=FileTransferMethod.TOOL_FILE, - ) - is True - ) + assert validator.is_accessible_file_mapping( + file_id=upload_file.id, + tenant_id=TENANT_ID, + transfer_method=FileTransferMethod.LOCAL_FILE, + ) + assert validator.is_accessible_file_mapping( + file_id=upload_file.id, + tenant_id=TENANT_ID, + transfer_method=FileTransferMethod.DATASOURCE_FILE, + ) + assert validator.is_accessible_file_mapping( + file_id=tool_file.id, + tenant_id=TENANT_ID, + transfer_method=FileTransferMethod.TOOL_FILE, + ) + assert not validator.is_accessible_file_mapping( + file_id=upload_file.id, + tenant_id=OTHER_TENANT_ID, + transfer_method=FileTransferMethod.LOCAL_FILE, + ) + assert not validator.is_accessible_file_mapping( + file_id=tool_file.id, + tenant_id=OTHER_TENANT_ID, + transfer_method=FileTransferMethod.TOOL_FILE, + ) + assert not validator.is_accessible_file_mapping( + file_id=upload_file.id, + tenant_id=TENANT_ID, + transfer_method=FileTransferMethod.TOOL_FILE, + ) + assert not validator.is_accessible_file_mapping( + file_id=tool_file.id, + tenant_id=TENANT_ID, + transfer_method=FileTransferMethod.LOCAL_FILE, + )