diff --git a/api/tests/unit_tests/controllers/console/test_files.py b/api/tests/unit_tests/controllers/console/test_files.py index f894e04f481..60bfb92d1c1 100644 --- a/api/tests/unit_tests/controllers/console/test_files.py +++ b/api/tests/unit_tests/controllers/console/test_files.py @@ -1,8 +1,10 @@ import io -from unittest.mock import MagicMock, patch +from datetime import datetime +from unittest.mock import patch import pytest from flask import Flask +from sqlalchemy import Engine from werkzeug.exceptions import Forbidden from configs import dify_config @@ -21,8 +23,11 @@ from controllers.console.files import ( FileSupportTypeApi, upload_file_from_request, ) +from extensions.storage.storage_type import StorageType from models import Account from models.account import AccountStatus, TenantAccountRole +from models.enums import CreatorUserRole +from models.model import UploadFile def unwrap(func): @@ -34,6 +39,24 @@ def unwrap(func): return func +def _upload_file(*, file_id: str = "file-id-123", size: int = 1024) -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-123", + storage_type=StorageType.LOCAL, + key=f"upload/{file_id}/test.txt", + name="test.txt", + size=size, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-123", + created_at=datetime(2024, 1, 1), + used=False, + ) + upload_file.id = file_id + return upload_file + + @pytest.fixture def app(): app = Flask(__name__) @@ -69,9 +92,9 @@ def mock_account_context(mock_current_user): @pytest.fixture -def mock_db(): +def mock_db(sqlite_engine: Engine): with patch("controllers.console.files.db") as db_mock: - db_mock.engine = MagicMock() + db_mock.engine = sqlite_engine yield db_mock @@ -161,25 +184,7 @@ class TestFileApiPost: api = FileApi() post_method = unwrap(api.post) - mock_file = MagicMock() - mock_file.id = "file-id-123" - mock_file.filename = "test.txt" - mock_file.name = "test.txt" - mock_file.size = 1024 - mock_file.extension = "txt" - mock_file.mime_type = "text/plain" - mock_file.created_by = "user-123" - mock_file.created_at = 1234567890 - mock_file.preview_url = "http://example.com/preview/file-id-123" - mock_file.source_url = "http://example.com/source/file-id-123" - mock_file.original_url = None - mock_file.reference = None - mock_file.user_id = "user-123" - mock_file.tenant_id = "tenant-123" - mock_file.conversation_id = None - mock_file.file_key = "file-key-123" - - mock_file_service.upload_file.return_value = mock_file + mock_file_service.upload_file.return_value = _upload_file() data = { "file": (io.BytesIO(b"hello"), "test.txt"), @@ -193,7 +198,7 @@ class TestFileApiPost: assert response["name"] == "test.txt" def test_upload_with_resource_tenant(self, app: Flask, mock_account_context, mock_file_service): - upload_file = MagicMock() + upload_file = _upload_file() mock_file_service.upload_file.return_value = upload_file with app.test_request_context( @@ -214,7 +219,7 @@ class TestFileApiPost: mock_account_context, mock_file_service, ): - upload_file = MagicMock() + upload_file = _upload_file() mock_file_service.upload_file.return_value = upload_file with ( @@ -240,26 +245,7 @@ class TestFileApiPost: api = FileApi() post_method = unwrap(api.post) - # Create a properly structured mock file object - mock_file = MagicMock() - mock_file.id = "file-id-456" - mock_file.filename = "test.txt" - mock_file.name = "test.txt" - mock_file.size = 512 - mock_file.extension = "txt" - mock_file.mime_type = "text/plain" - mock_file.created_by = "user-456" - mock_file.created_at = 1234567890 - mock_file.preview_url = None - mock_file.source_url = None - mock_file.original_url = None - mock_file.reference = None - mock_file.user_id = "user-456" - mock_file.tenant_id = "tenant-456" - mock_file.conversation_id = None - mock_file.file_key = "file-key-456" - - mock_file_service.upload_file.return_value = mock_file + mock_file_service.upload_file.return_value = _upload_file(file_id="file-id-456", size=512) data = { "file": (io.BytesIO(b"content"), "test.txt"), diff --git a/api/tests/unit_tests/controllers/console/test_files_security.py b/api/tests/unit_tests/controllers/console/test_files_security.py index 370bf63fdbc..eb090ff8ea5 100644 --- a/api/tests/unit_tests/controllers/console/test_files_security.py +++ b/api/tests/unit_tests/controllers/console/test_files_security.py @@ -13,6 +13,8 @@ from controllers.common.errors import ( TooManyFilesError, UnsupportedFileTypeError, ) +from models import Account +from models.account import AccountStatus, TenantAccountRole from services.errors.file import FileTooLargeError as ServiceFileTooLargeError from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError @@ -106,11 +108,8 @@ class TestFileUploadSecurity: # Test 3: Permission validation def test_should_validate_dataset_permissions(self): """Test dataset upload permission logic""" - - class MockUser: - is_dataset_editor = False - - user = MockUser() + user = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) + user.role = TenantAccountRole.NORMAL source = "datasets" # Simulate the permission check in FileApi.post() @@ -120,11 +119,8 @@ class TestFileUploadSecurity: def test_should_allow_general_upload_without_permission(self): """Test general upload doesn't require dataset permission""" - - class MockUser: - is_dataset_editor = False - - user = MockUser() + user = Account(name="Viewer", email="viewer@example.com", status=AccountStatus.ACTIVE) + user.role = TenantAccountRole.NORMAL source = None # General upload # This should not raise an exception diff --git a/api/tests/unit_tests/controllers/console/test_knowledge_fs_proxy.py b/api/tests/unit_tests/controllers/console/test_knowledge_fs_proxy.py index 34a8ca9fec2..1a00bb7a508 100644 --- a/api/tests/unit_tests/controllers/console/test_knowledge_fs_proxy.py +++ b/api/tests/unit_tests/controllers/console/test_knowledge_fs_proxy.py @@ -28,6 +28,7 @@ from controllers.console.knowledge_fs_proxy import ( proxy_knowledge_fs_write, ) from controllers.console.wraps import RBACPermission +from models.account import Account, TenantAccountRole from services.knowledge_fs_operations import ( KnowledgeFSMethod, KnowledgeFSOperation, @@ -84,12 +85,17 @@ def _set_current_workspace( has_edit_permission: bool = True, admin_or_owner: bool = True, ) -> None: - account = MagicMock( - id="account-1", - has_edit_permission=has_edit_permission, - is_admin_or_owner=admin_or_owner, - is_dataset_editor=editor, - ) + if admin_or_owner: + role = TenantAccountRole.ADMIN + elif editor and has_edit_permission: + role = TenantAccountRole.EDITOR + elif editor: + role = TenantAccountRole.DATASET_OPERATOR + else: + role = TenantAccountRole.NORMAL + account = Account(name="Knowledge User", email="knowledge@example.com") + account.id = "account-1" + account.role = role monkeypatch.setattr( "controllers.console.knowledge_fs_proxy.current_account_with_tenant", lambda: (account, "tenant-1"), @@ -286,9 +292,11 @@ def test_read_post_applies_knowledge_rate_limit_once( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr("controllers.common.wraps.dify_config.RBAC_ENABLED", False) - account = MagicMock(id="account-1", is_dataset_editor=True) + account = Account(name="Knowledge User", email="knowledge@example.com") + account.id = "account-1" + account.role = TenantAccountRole.DATASET_OPERATOR - def current_workspace() -> tuple[MagicMock, str]: + def current_workspace() -> tuple[Account, str]: return account, "tenant-1" monkeypatch.setattr("controllers.console.knowledge_fs_proxy.current_account_with_tenant", current_workspace) @@ -323,9 +331,11 @@ def test_denied_write_does_not_consume_the_workspace_rate_limit( app: Flask, monkeypatch: pytest.MonkeyPatch, ) -> None: - account = MagicMock(id="account-1", is_dataset_editor=False) + account = Account(name="Knowledge Viewer", email="viewer@example.com") + account.id = "account-1" + account.role = TenantAccountRole.NORMAL - def current_workspace() -> tuple[MagicMock, str]: + def current_workspace() -> tuple[Account, str]: return account, "tenant-1" monkeypatch.setattr( @@ -460,9 +470,11 @@ def test_generic_write_forwards_through_the_authorized_production_path( app: Flask, monkeypatch: pytest.MonkeyPatch, ) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True) + account = Account(name="Knowledge User", email="knowledge@example.com") + account.id = "account-1" + account.role = TenantAccountRole.DATASET_OPERATOR - def current_workspace() -> tuple[MagicMock, str]: + def current_workspace() -> tuple[Account, str]: return account, "tenant-1" monkeypatch.setattr("controllers.console.knowledge_fs_proxy.current_account_with_tenant", current_workspace) diff --git a/api/tests/unit_tests/controllers/console/test_remote_files.py b/api/tests/unit_tests/controllers/console/test_remote_files.py index f50d60e390d..6eeb68f79ac 100644 --- a/api/tests/unit_tests/controllers/console/test_remote_files.py +++ b/api/tests/unit_tests/controllers/console/test_remote_files.py @@ -12,8 +12,11 @@ from flask import Flask from controllers.common.errors import FileTooLargeError, RemoteFileUploadError, UnsupportedFileTypeError from controllers.console import remote_files as remote_files_module +from extensions.storage.storage_type import StorageType from models import Account from models.account import AccountStatus, TenantAccountRole +from models.enums import CreatorUserRole +from models.model import UploadFile from services.errors.file import FileTooLargeError as ServiceFileTooLargeError from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError @@ -29,6 +32,32 @@ def _make_account(account_id: str = "u1") -> Account: return account +def _upload_file( + *, + file_id: str = "file-1", + name: str = "report.txt", + size: int = 16, + extension: str = ".txt", + mime_type: str = "text/plain", + created_at: datetime | None = None, +) -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key=f"upload/{name}", + name=name, + size=size, + extension=extension, + mime_type=mime_type, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="u1", + created_at=created_at or datetime(2024, 1, 1, tzinfo=UTC), + used=False, + ) + upload_file.id = file_id + return upload_file + + class _FakeResponse: def __init__( self, @@ -160,15 +189,7 @@ def test_remote_file_upload_success_when_fetch_falls_back_to_get(app: Flask, mon monkeypatch.setattr(remote_files_module.remote_fetcher, "make_request", make_request) file_service_cls, current_user = _mock_upload_dependencies(monkeypatch) - upload_file = SimpleNamespace( - id="file-1", - name="report.txt", - size=16, - extension=".txt", - mime_type="text/plain", - created_by="u1", - created_at=datetime(2024, 1, 1, tzinfo=UTC), - ) + upload_file = _upload_file() file_service_cls.return_value.upload_file.return_value = upload_file with app.test_request_context(method="POST", json={"url": url}): @@ -197,15 +218,7 @@ def test_remote_file_upload_assigns_resource_tenant(app: Flask, monkeypatch: pyt monkeypatch.setattr(remote_files_module.remote_fetcher, "make_request", MagicMock(return_value=response)) file_service_cls, current_user = _mock_upload_dependencies(monkeypatch) - file_service_cls.return_value.upload_file.return_value = SimpleNamespace( - id="file-1", - name="report.txt", - size=7, - extension=".txt", - mime_type="text/plain", - created_by="u1", - created_at=datetime(2024, 1, 1, tzinfo=UTC), - ) + file_service_cls.return_value.upload_file.return_value = _upload_file(size=7) with app.test_request_context(method="POST", json={"url": url}): remote_files_module.upload_remote_file_from_request( @@ -236,13 +249,12 @@ def test_remote_file_upload_fetches_content_with_second_get_when_head_succeeds( monkeypatch.setattr(remote_files_module.remote_fetcher, "make_request", make_request) file_service_cls, current_user = _mock_upload_dependencies(monkeypatch) - upload_file = SimpleNamespace( - id="file-2", + upload_file = _upload_file( + file_id="file-2", name="photo.jpg", size=18, extension=".jpg", mime_type="image/jpeg", - created_by="u1", created_at=datetime(2024, 1, 2, tzinfo=UTC), ) file_service_cls.return_value.upload_file.return_value = upload_file diff --git a/api/tests/unit_tests/controllers/web/test_files.py b/api/tests/unit_tests/controllers/web/test_files.py index a3921b03731..d25ef232a04 100644 --- a/api/tests/unit_tests/controllers/web/test_files.py +++ b/api/tests/unit_tests/controllers/web/test_files.py @@ -2,12 +2,13 @@ from __future__ import annotations +from datetime import datetime from io import BytesIO -from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from flask import Flask +from sqlalchemy import Engine from controllers.common.errors import ( FilenameNotExistsError, @@ -16,14 +17,50 @@ from controllers.common.errors import ( TooManyFilesError, ) from controllers.web.files import FileApi +from extensions.storage.storage_type import StorageType +from models.enums import CreatorUserRole, EndUserType +from models.model import App, AppMode, EndUser, UploadFile -def _app_model() -> SimpleNamespace: - return SimpleNamespace(id="app-1") +def _app_model() -> App: + return App( + id="app-1", + tenant_id="tenant-1", + name="Web App", + description="", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) -def _end_user() -> SimpleNamespace: - return SimpleNamespace(id="eu-1") +def _end_user() -> EndUser: + return EndUser( + id="eu-1", + tenant_id="tenant-1", + app_id="app-1", + type=EndUserType.BROWSER, + session_id="session-1", + ) + + +def _upload_file() -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key="upload/test.txt", + name="test.txt", + size=100, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.END_USER, + created_by="eu-1", + created_at=datetime(2024, 1, 1), + used=False, + ) + upload_file.id = "file-1" + return upload_file class TestFileApi: @@ -50,20 +87,11 @@ class TestFileApi: @patch("controllers.web.files.FileService") @patch("controllers.web.files.db") - def test_upload_success(self, mock_db: MagicMock, mock_file_svc_cls: MagicMock, app: Flask) -> None: - mock_db.engine = "engine" - from datetime import datetime - - upload_file = SimpleNamespace( - id="file-1", - name="test.txt", - size=100, - extension="txt", - mime_type="text/plain", - created_by="eu-1", - created_at=datetime(2024, 1, 1), - ) - mock_file_svc_cls.return_value.upload_file.return_value = upload_file + def test_upload_success( + self, mock_db: MagicMock, mock_file_svc_cls: MagicMock, app: Flask, sqlite_engine: Engine + ) -> None: + mock_db.engine = sqlite_engine + mock_file_svc_cls.return_value.upload_file.return_value = _upload_file() data = {"file": (BytesIO(b"content"), "test.txt")} with app.test_request_context("/files/upload", method="POST", data=data, content_type="multipart/form-data"): @@ -75,10 +103,12 @@ class TestFileApi: @patch("controllers.web.files.FileService") @patch("controllers.web.files.db") - def test_file_too_large_from_service(self, mock_db: MagicMock, mock_file_svc_cls: MagicMock, app: Flask) -> None: + def test_file_too_large_from_service( + self, mock_db: MagicMock, mock_file_svc_cls: MagicMock, app: Flask, sqlite_engine: Engine + ) -> None: import services.errors.file - mock_db.engine = "engine" + mock_db.engine = sqlite_engine mock_file_svc_cls.return_value.upload_file.side_effect = services.errors.file.FileTooLargeError( description="max 10MB" ) diff --git a/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py b/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py index a01a74b5353..81e361d494c 100644 --- a/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py +++ b/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock import pytest from flask import Flask +from sqlalchemy import Engine import controllers.web.human_input_file_upload as upload_module from controllers.common.errors import NoFileUploadedError @@ -17,6 +18,12 @@ from controllers.web.human_input_file_upload import ( InvalidUploadTokenForbiddenError, InvalidUploadTokenUnauthorizedError, ) +from extensions.storage.storage_type import StorageType +from models import Account +from models.account import AccountStatus +from models.enums import CreatorUserRole +from models.model import UploadFile +from services.human_input_file_upload_service import HumanInputUploadContext @pytest.fixture @@ -26,26 +33,40 @@ def app() -> Flask: return app -def _upload_context() -> SimpleNamespace: - return SimpleNamespace( +def _account() -> Account: + account = Account(name="Form Owner", email="owner@example.com", status=AccountStatus.ACTIVE) + account.id = "owner-1" + return account + + +def _upload_context() -> HumanInputUploadContext: + return HumanInputUploadContext( + tenant_id="tenant-1", + app_id="app-1", form_id="form-1", + recipient_id="recipient-1", upload_token_id="token-row-1", - owner=SimpleNamespace(id="owner-1", current_tenant_id="tenant-1"), + owner=_account(), ) -def _upload_file() -> SimpleNamespace: - return SimpleNamespace( - id="file-1", +def _upload_file() -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key="upload/sample.txt", name="sample.txt", size=7, extension="txt", mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, created_by="end-user-1", created_at=datetime(2024, 1, 1), - tenant_id="tenant-1", + used=False, source_url="signed-source-url", ) + upload_file.id = "file-1" + return upload_file def _patch_upload_service(monkeypatch: pytest.MonkeyPatch, service: MagicMock) -> tuple[MagicMock, dict[str, object]]: @@ -90,7 +111,9 @@ def test_local_upload_requires_authorization_before_reading_files(app: Flask) -> HumanInputFileUploadApi().post() -def test_local_upload_ignores_source_and_records_form_file_link(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: +def test_local_upload_ignores_source_and_records_form_file_link( + monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine +) -> None: service = MagicMock() service.validate_upload_token.return_value = _upload_context() repo_factory, captured = _patch_upload_service(monkeypatch, service) @@ -99,7 +122,7 @@ def test_local_upload_ignores_source_and_records_form_file_link(monkeypatch: pyt file_service.upload_file.return_value = _upload_file() file_service_cls = MagicMock(return_value=file_service) monkeypatch.setattr(upload_module, "FileService", file_service_cls) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) data = { "file": (BytesIO(b"content"), "sample.txt"), @@ -127,11 +150,13 @@ def test_local_upload_ignores_source_and_records_form_file_link(monkeypatch: pyt ) -def test_local_upload_missing_file_raises_after_valid_token(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: +def test_local_upload_missing_file_raises_after_valid_token( + monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine +) -> None: service = MagicMock() service.validate_upload_token.return_value = _upload_context() _patch_upload_service(monkeypatch, service) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) with app.test_request_context( "/api/human-input-forms/files", @@ -145,11 +170,13 @@ def test_local_upload_missing_file_raises_after_valid_token(monkeypatch: pytest. service.validate_upload_token.assert_called_once_with("hitl_upload_token-1") -def test_remote_upload_validates_token_before_fetching_remote_url(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: +def test_remote_upload_validates_token_before_fetching_remote_url( + monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine +) -> None: service = MagicMock() service.validate_upload_token.side_effect = InvalidUploadTokenForbiddenError() _patch_upload_service(monkeypatch, service) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) ssrf_proxy = MagicMock() monkeypatch.setattr(upload_module, "ssrf_proxy", ssrf_proxy) @@ -167,11 +194,13 @@ def test_remote_upload_validates_token_before_fetching_remote_url(monkeypatch: p ssrf_proxy.get.assert_not_called() -def test_remote_upload_records_form_file_link(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: +def test_remote_upload_records_form_file_link( + monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine +) -> None: service = MagicMock() service.validate_upload_token.return_value = _upload_context() _patch_upload_service(monkeypatch, service) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) response = MagicMock() response.status_code = 200 diff --git a/api/tests/unit_tests/controllers/web/test_remote_files.py b/api/tests/unit_tests/controllers/web/test_remote_files.py index 3c829e74816..ae912489e14 100644 --- a/api/tests/unit_tests/controllers/web/test_remote_files.py +++ b/api/tests/unit_tests/controllers/web/test_remote_files.py @@ -3,22 +3,60 @@ from __future__ import annotations import urllib.parse +from datetime import datetime from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from flask import Flask +from sqlalchemy import Engine from controllers.common.errors import FileTooLargeError, RemoteFileUploadError from controllers.web.remote_files import RemoteFileInfoApi, RemoteFileUploadApi +from extensions.storage.storage_type import StorageType +from models.enums import CreatorUserRole, EndUserType +from models.model import App, AppMode, EndUser, UploadFile -def _app_model() -> SimpleNamespace: - return SimpleNamespace(id="app-1") +def _app_model() -> App: + return App( + id="app-1", + tenant_id="tenant-1", + name="Web App", + description="", + mode=AppMode.CHAT, + enable_site=True, + enable_api=True, + max_active_requests=0, + ) -def _end_user() -> SimpleNamespace: - return SimpleNamespace(id="eu-1") +def _end_user() -> EndUser: + return EndUser( + id="eu-1", + tenant_id="tenant-1", + app_id="app-1", + type=EndUserType.BROWSER, + session_id="session-1", + ) + + +def _upload_file() -> UploadFile: + upload_file = UploadFile( + tenant_id="tenant-1", + storage_type=StorageType.LOCAL, + key="upload/file.pdf", + name="file.pdf", + size=100, + extension="pdf", + mime_type="application/pdf", + created_by_role=CreatorUserRole.END_USER, + created_by="eu-1", + created_at=datetime(2024, 1, 1), + used=False, + ) + upload_file.id = "f-1" + return upload_file # --------------------------------------------------------------------------- @@ -107,8 +145,9 @@ class TestRemoteFileUploadApi: mock_file_svc_cls: MagicMock, mock_signed: MagicMock, app: Flask, + sqlite_engine: Engine, ) -> None: - mock_db.engine = "engine" + mock_db.engine = sqlite_engine mock_ns.payload = {"url": "https://example.com/file.pdf"} head_resp = MagicMock() head_resp.status_code = 200 @@ -123,18 +162,7 @@ class TestRemoteFileUploadApi: ) mock_file_svc_cls.is_file_size_within_limit.return_value = True - from datetime import datetime - - upload_file = SimpleNamespace( - id="f-1", - name="file.pdf", - size=100, - extension="pdf", - mime_type="application/pdf", - created_by="eu-1", - created_at=datetime(2024, 1, 1), - ) - mock_file_svc_cls.return_value.upload_file.return_value = upload_file + mock_file_svc_cls.return_value.upload_file.return_value = _upload_file() with app.test_request_context("/remote-files/upload", method="POST"): result, status = RemoteFileUploadApi().post(_app_model(), _end_user()) diff --git a/api/tests/unit_tests/services/test_file_service.py b/api/tests/unit_tests/services/test_file_service.py index f345a2ec012..ffd507fd06a 100644 --- a/api/tests/unit_tests/services/test_file_service.py +++ b/api/tests/unit_tests/services/test_file_service.py @@ -3,7 +3,7 @@ import hashlib import os from collections.abc import Iterator from datetime import UTC, datetime -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from sqlalchemy import Engine @@ -20,6 +20,12 @@ from services.errors.file import BlockedFileExtensionError, FileTooLargeError, U from services.file_service import FileService +def _account() -> Account: + account = Account(name="Test Account", email="test@example.com") + account.id = "user_id" + return account + + class TestFileService: @pytest.fixture def sqlite_session_maker(self, sqlite_engine: Engine) -> sessionmaker[Session]: @@ -150,7 +156,7 @@ class TestFileService: def test_upload_file_invalid_characters(self, file_service): with pytest.raises(ValueError, match="Filename contains invalid characters"): - file_service.upload_file(filename="invalid/file.txt", content=b"", mimetype="text/plain", user=MagicMock()) + file_service.upload_file(filename="invalid/file.txt", content=b"", mimetype="text/plain", user=_account()) def test_upload_file_long_filename(self, file_service: FileService, db_session: Session): # Setup @@ -173,13 +179,13 @@ class TestFileService: with patch.object(dify_config, "inner_UPLOAD_FILE_EXTENSION_BLACKLIST", "exe"): with pytest.raises(BlockedFileExtensionError): file_service.upload_file( - filename="test.exe", content=b"", mimetype="application/octet-stream", user=MagicMock() + filename="test.exe", content=b"", mimetype="application/octet-stream", user=_account() ) def test_upload_file_unsupported_type_for_datasets(self, file_service): with pytest.raises(UnsupportedFileTypeError): file_service.upload_file( - filename="test.jpg", content=b"", mimetype="image/jpeg", user=MagicMock(), source="datasets" + filename="test.jpg", content=b"", mimetype="image/jpeg", user=_account(), source="datasets" ) def test_upload_file_too_large(self, file_service): @@ -187,7 +193,7 @@ class TestFileService: content = b"a" * (16 * 1024 * 1024) with patch.object(dify_config, "UPLOAD_IMAGE_FILE_SIZE_LIMIT", 15): with pytest.raises(FileTooLargeError): - file_service.upload_file(filename="test.jpg", content=content, mimetype="image/jpeg", user=MagicMock()) + file_service.upload_file(filename="test.jpg", content=content, mimetype="image/jpeg", user=_account()) def test_upload_file_end_user(self, file_service: FileService, db_session: Session): user = EndUser( diff --git a/api/tests/unit_tests/services/test_knowledge_fs_proxy.py b/api/tests/unit_tests/services/test_knowledge_fs_proxy.py index 3ba26fb9733..4d7e67f6d60 100644 --- a/api/tests/unit_tests/services/test_knowledge_fs_proxy.py +++ b/api/tests/unit_tests/services/test_knowledge_fs_proxy.py @@ -10,6 +10,7 @@ from pydantic import SecretStr from core.helper import ssrf_proxy from core.rbac import RBACPermission from core.tools.errors import ToolSSRFError +from models.account import Account, TenantAccountRole from services.knowledge_fs_operations import ( KNOWLEDGE_FS_CONSOLE_OPERATIONS, KnowledgeFSMethod, @@ -33,6 +34,23 @@ from services.knowledge_fs_proxy import ( _JWT_SECRET = "production-secret-with-at-least-32-bytes" + +def _account(*, role: TenantAccountRole = TenantAccountRole.DATASET_OPERATOR) -> Account: + account = Account(name="Knowledge User", email="knowledge@example.com") + account.id = "account-1" + account.role = role + return account + + +def _account_for_legacy_role(role: str) -> Account: + roles = { + "admin": TenantAccountRole.ADMIN, + "dataset_editor": TenantAccountRole.DATASET_OPERATOR, + "reader": TenantAccountRole.NORMAL, + } + return _account(role=roles[role]) + + _HAPPY_PATH_OPERATION_IDS = ( "listKnowledgeSpaces", "createKnowledgeSpace", @@ -322,7 +340,7 @@ def test_list_and_create_forward_raw_request(monkeypatch: pytest.MonkeyPatch, me def test_proxy_forwards_only_registry_declared_headers(monkeypatch: pytest.MonkeyPatch) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True) + account = _account() upstream = MagicMock() forward = MagicMock(return_value=upstream) monkeypatch.setattr( @@ -344,7 +362,7 @@ def test_proxy_forwards_only_registry_declared_headers(monkeypatch: pytest.Monke def test_authorized_proxy_does_not_repeat_workspace_rbac(monkeypatch: pytest.MonkeyPatch) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True) + account = _account() check_access = MagicMock(return_value=True) forward = MagicMock(return_value=MagicMock()) monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access) @@ -374,7 +392,7 @@ def test_authorization_capability_cannot_be_constructed_directly() -> None: def test_authorization_resolves_the_canonical_operation_policy(monkeypatch: pytest.MonkeyPatch) -> None: - account = MagicMock(id="account-1", is_dataset_editor=False, is_admin_or_owner=False) + account = _account(role=TenantAccountRole.NORMAL) check_access = MagicMock(return_value=True) monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access) @@ -402,7 +420,7 @@ def test_authorization_capability_binding_cannot_be_mutated( attribute: str, value: object, ) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True) + account = _account() monkeypatch.setattr( "services.knowledge_fs_proxy.RBACService.CheckAccess.check", MagicMock(return_value=True), @@ -423,7 +441,7 @@ def test_authorization_capability_binding_cannot_be_mutated( def test_authorization_capability_cannot_be_reused(monkeypatch: pytest.MonkeyPatch) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True) + account = _account() forward = MagicMock(return_value=MagicMock()) monkeypatch.setattr( "services.knowledge_fs_proxy.RBACService.CheckAccess.check", @@ -449,7 +467,7 @@ def test_authorization_rejects_workspace_rbac_denial( monkeypatch: pytest.MonkeyPatch, operation: KnowledgeFSOperation, ) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True) + account = _account_for_legacy_role(operation.legacy_role) check_access = MagicMock(return_value=False) monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access) @@ -478,7 +496,7 @@ def test_dataset_editor_operations_reject_legacy_viewers_before_rbac( monkeypatch: pytest.MonkeyPatch, operation: KnowledgeFSOperation, ) -> None: - account = MagicMock(id="account-1", is_dataset_editor=False, is_admin_or_owner=False) + account = _account(role=TenantAccountRole.NORMAL) check_access = MagicMock(return_value=True) monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access) @@ -494,7 +512,7 @@ def test_dataset_editor_operations_reject_legacy_viewers_before_rbac( def test_admin_operation_rejects_legacy_editors_before_rbac(monkeypatch: pytest.MonkeyPatch) -> None: - account = MagicMock(id="account-1", is_dataset_editor=True, is_admin_or_owner=False) + account = _account() check_access = MagicMock(return_value=True) monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access) operation = get_knowledge_fs_operation( @@ -513,7 +531,7 @@ def test_admin_operation_rejects_legacy_editors_before_rbac(monkeypatch: pytest. def test_authorization_uses_the_declared_reader_policy(monkeypatch: pytest.MonkeyPatch) -> None: - account = MagicMock(id="account-1", is_dataset_editor=False, is_admin_or_owner=False) + account = _account(role=TenantAccountRole.NORMAL) check_access = MagicMock(return_value=True) monkeypatch.setattr("services.knowledge_fs_proxy.RBACService.CheckAccess.check", check_access) operation = get_knowledge_fs_operation("GET", "knowledge-spaces")