From e21071cbb8f65ed88f169d9dc86176e9e1f856ca Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 14 Jul 2026 16:01:23 +0900 Subject: [PATCH] test: use sqlite3 session in test_remote_fetcher (#38773) --- .../core/file/test_remote_fetcher.py | 357 ++++++++---------- 1 file changed, 162 insertions(+), 195 deletions(-) diff --git a/api/tests/unit_tests/core/file/test_remote_fetcher.py b/api/tests/unit_tests/core/file/test_remote_fetcher.py index c9612a062ac..8e414e22064 100644 --- a/api/tests/unit_tests/core/file/test_remote_fetcher.py +++ b/api/tests/unit_tests/core/file/test_remote_fetcher.py @@ -2,20 +2,104 @@ import base64 import hashlib import hmac import urllib.parse -from types import SimpleNamespace +from collections.abc import Iterator +from dataclasses import dataclass +from datetime import UTC, datetime from unittest.mock import MagicMock import httpx import pytest +from sqlalchemy import Engine, event +from sqlalchemy.orm import Session, sessionmaker from core.datasource.datasource_file_manager import DatasourceFileManager from core.file import remote_fetcher from core.tools.signature import sign_tool_file, sign_upload_file_preview_url from core.tools.tool_file_manager import ToolFileManager +from extensions.storage.storage_type import StorageType +from models import CreatorUserRole, ToolFile, UploadFile UPLOAD_FILE_ID = "1602650a-4fe4-423c-85a2-af76c083e3c4" TOOL_FILE_ID = "2602650a-4fe4-423c-85a2-af76c083e3c4" DATASOURCE_FILE_ID = "3602650a-4fe4-423c-85a2-af76c083e3c4" +TENANT_ID = "tenant-1" +USER_ID = "user-1" + + +@dataclass(frozen=True) +class FileDatabase: + session: Session + upload_file: UploadFile + tool_file: ToolFile + datasource_upload_file: UploadFile + statements: list[str] + + +@pytest.fixture(autouse=True) +def file_database(monkeypatch: pytest.MonkeyPatch, sqlite_engine: Engine) -> Iterator[FileDatabase]: + """Persist file records and bind remote-fetcher 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=TENANT_ID, + storage_type=StorageType.LOCAL, + key="upload_files/tenant/hello.txt", + name="hello.txt", + size=5, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=USER_ID, + created_at=datetime(2026, 1, 1, tzinfo=UTC), + used=False, + ) + upload_file.id = UPLOAD_FILE_ID + tool_file = ToolFile( + user_id=USER_ID, + tenant_id=TENANT_ID, + conversation_id=None, + file_key="tools/tenant/result.txt", + mimetype="text/plain", + name="result.txt", + size=6, + ) + tool_file.id = TOOL_FILE_ID + datasource_upload_file = UploadFile( + tenant_id=TENANT_ID, + storage_type=StorageType.LOCAL, + key="datasources/tenant/data.txt", + name="data.txt", + size=4, + extension="txt", + mime_type="text/plain", + created_by_role=CreatorUserRole.ACCOUNT, + created_by=USER_ID, + created_at=datetime(2026, 1, 1, tzinfo=UTC), + used=False, + ) + datasource_upload_file.id = DATASOURCE_FILE_ID + + statements: list[str] = [] + + def record_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: + statements.append(statement) + + with sqlite_session_maker() as session: + session.add_all([upload_file, tool_file, datasource_upload_file]) + session.commit() + event.listen(sqlite_engine, "before_cursor_execute", record_statement) + try: + yield FileDatabase( + session=session, + upload_file=upload_file, + tool_file=tool_file, + datasource_upload_file=datasource_upload_file, + statements=statements, + ) + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_statement) def _signed_url(*, base_url: str, path: str, payload: str, secret: str = "test-secret") -> str: @@ -44,15 +128,6 @@ def _patch_file_fetcher_config(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(remote_fetcher.time, "time", lambda: 1700000100) -def _patch_session(monkeypatch: pytest.MonkeyPatch): - session = MagicMock() - session_cm = MagicMock() - session_cm.__enter__.return_value = session - session_cm.__exit__.return_value = False - monkeypatch.setattr(remote_fetcher.session_factory, "create_session", MagicMock(return_value=session_cm)) - return session - - def _patch_ssrf_make_request(monkeypatch: pytest.MonkeyPatch, response=None): make_request = MagicMock(return_value=response) if response is not None else MagicMock() monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", make_request) @@ -67,18 +142,7 @@ def _patch_signer_times(monkeypatch: pytest.MonkeyPatch): def test_get_signed_upload_file_url_reads_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch): _patch_file_fetcher_config(monkeypatch) - session = _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=UPLOAD_FILE_ID, - key="upload_files/tenant/hello.txt", - name="hello.txt", - mime_type="text/plain", - size=5, - extension="txt", - ) - get_upload_file = MagicMock(return_value=upload_file) load_once = MagicMock(return_value=b"hello") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( @@ -94,29 +158,22 @@ def test_get_signed_upload_file_url_reads_storage_without_ssrf(monkeypatch: pyte assert response.headers["Content-Type"] == "text/plain" assert response.headers["Content-Length"] == "5" assert response.request.method == "GET" - get_upload_file.assert_called_once_with( - session=session, - file_id=UPLOAD_FILE_ID, - ) load_once.assert_called_once_with("upload_files/tenant/hello.txt") ssrf_make_request.assert_not_called() -def test_make_request_resolves_upload_preview_url_generated_by_signer(monkeypatch: pytest.MonkeyPatch): +def test_make_request_resolves_upload_preview_url_generated_by_signer( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) _patch_signer_times(monkeypatch) - session = _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=UPLOAD_FILE_ID, - key="upload_files/tenant/image.png", - name="image.png", - mime_type="image/png", - size=6, - extension=".png", - ) - get_upload_file = MagicMock(return_value=upload_file) + file_database.upload_file.key = "upload_files/tenant/image.png" + file_database.upload_file.name = "image.png" + file_database.upload_file.mime_type = "image/png" + file_database.upload_file.size = 6 + file_database.upload_file.extension = "png" + file_database.session.commit() load_once = MagicMock(return_value=b"image!") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = MagicMock() monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request) @@ -127,25 +184,21 @@ def test_make_request_resolves_upload_preview_url_generated_by_signer(monkeypatc assert response.status_code == 200 assert response.content == b"image!" assert response.headers["Content-Type"] == "image/png" - get_upload_file.assert_called_once_with(session=session, file_id=UPLOAD_FILE_ID) load_once.assert_called_once_with("upload_files/tenant/image.png") ssrf_make_request.assert_not_called() -def test_make_request_resolves_sign_tool_file_url_with_empty_extension(monkeypatch: pytest.MonkeyPatch): +def test_make_request_resolves_sign_tool_file_url_with_empty_extension( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) _patch_signer_times(monkeypatch) - session = _patch_session(monkeypatch) - tool_file = SimpleNamespace( - id=TOOL_FILE_ID, - file_key="tools/tenant/no-extension", - name="no-extension", - mimetype="application/octet-stream", - size=8, - ) - get_tool_file = MagicMock(return_value=tool_file) + file_database.tool_file.file_key = "tools/tenant/no-extension" + file_database.tool_file.name = "no-extension" + file_database.tool_file.mimetype = "application/octet-stream" + file_database.tool_file.size = 8 + file_database.session.commit() load_once = MagicMock(return_value=b"tooldata") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = MagicMock() monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request) @@ -156,25 +209,21 @@ def test_make_request_resolves_sign_tool_file_url_with_empty_extension(monkeypat assert response.status_code == 200 assert response.content == b"tooldata" assert response.headers["Content-Type"] == "application/octet-stream" - get_tool_file.assert_called_once_with(session=session, file_id=TOOL_FILE_ID) load_once.assert_called_once_with("tools/tenant/no-extension") ssrf_make_request.assert_not_called() -def test_make_request_resolves_tool_manager_url_with_empty_extension(monkeypatch: pytest.MonkeyPatch): +def test_make_request_resolves_tool_manager_url_with_empty_extension( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) _patch_signer_times(monkeypatch) - session = _patch_session(monkeypatch) - tool_file = SimpleNamespace( - id=TOOL_FILE_ID, - file_key="tools/tenant/manager-file", - name="manager-file", - mimetype="application/octet-stream", - size=12, - ) - get_tool_file = MagicMock(return_value=tool_file) + file_database.tool_file.file_key = "tools/tenant/manager-file" + file_database.tool_file.name = "manager-file" + file_database.tool_file.mimetype = "application/octet-stream" + file_database.tool_file.size = 12 + file_database.session.commit() load_once = MagicMock(return_value=b"manager-data") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = MagicMock() monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request) @@ -184,28 +233,22 @@ def test_make_request_resolves_tool_manager_url_with_empty_extension(monkeypatch assert response.status_code == 200 assert response.content == b"manager-data" - get_tool_file.assert_called_once_with(session=session, file_id=TOOL_FILE_ID) load_once.assert_called_once_with("tools/tenant/manager-file") ssrf_make_request.assert_not_called() -def test_make_request_resolves_datasource_manager_url_with_empty_extension(monkeypatch: pytest.MonkeyPatch): +def test_make_request_resolves_datasource_manager_url_with_empty_extension( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) _patch_signer_times(monkeypatch) - _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=DATASOURCE_FILE_ID, - key="datasources/tenant/no-extension", - name="no-extension", - mime_type="application/octet-stream", - size=10, - extension="", - ) - get_upload_file = MagicMock(return_value=upload_file) - get_tool_file = MagicMock() + file_database.datasource_upload_file.key = "datasources/tenant/no-extension" + file_database.datasource_upload_file.name = "no-extension" + file_database.datasource_upload_file.mime_type = "application/octet-stream" + file_database.datasource_upload_file.size = 10 + file_database.datasource_upload_file.extension = "" + file_database.session.commit() load_once = MagicMock(return_value=b"datasource") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = MagicMock() monkeypatch.setattr(remote_fetcher.ssrf_proxy, "make_request", ssrf_make_request) @@ -216,25 +259,12 @@ def test_make_request_resolves_datasource_manager_url_with_empty_extension(monke assert response.status_code == 200 assert response.content == b"datasource" assert response.headers["Content-Type"] == "application/octet-stream" - get_upload_file.assert_called_once() - get_tool_file.assert_not_called() load_once.assert_called_once_with("datasources/tenant/no-extension") ssrf_make_request.assert_not_called() def test_head_signed_upload_file_url_returns_metadata_without_storage_content(monkeypatch: pytest.MonkeyPatch): _patch_file_fetcher_config(monkeypatch) - session = _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=UPLOAD_FILE_ID, - key="upload_files/tenant/hello.txt", - name="hello.txt", - mime_type="text/plain", - size=5, - extension="txt", - ) - get_upload_file = MagicMock(return_value=upload_file) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) load_once = MagicMock(return_value=b"hello") monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) @@ -251,28 +281,13 @@ def test_head_signed_upload_file_url_returns_metadata_without_storage_content(mo assert response.headers["Content-Type"] == "text/plain" assert response.headers["Content-Length"] == "5" assert response.request.method == "HEAD" - get_upload_file.assert_called_once_with( - session=session, - file_id=UPLOAD_FILE_ID, - ) load_once.assert_not_called() ssrf_make_request.assert_not_called() def test_make_request_get_signed_upload_file_url_reads_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=UPLOAD_FILE_ID, - key="upload_files/tenant/hello.txt", - name="hello.txt", - mime_type="text/plain", - size=5, - extension="txt", - ) - get_upload_file = MagicMock(return_value=upload_file) load_once = MagicMock(return_value=b"hello") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( @@ -286,24 +301,12 @@ def test_make_request_get_signed_upload_file_url_reads_storage_without_ssrf(monk assert response.status_code == 200 assert response.content == b"hello" assert response.request.method == "GET" - get_upload_file.assert_called_once() load_once.assert_called_once_with("upload_files/tenant/hello.txt") ssrf_make_request.assert_not_called() def test_make_request_head_signed_upload_file_url_returns_metadata_without_ssrf(monkeypatch: pytest.MonkeyPatch): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=UPLOAD_FILE_ID, - key="upload_files/tenant/hello.txt", - name="hello.txt", - mime_type="text/plain", - size=5, - extension="txt", - ) - get_upload_file = MagicMock(return_value=upload_file) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) load_once = MagicMock(return_value=b"hello") monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) @@ -320,15 +323,14 @@ def test_make_request_head_signed_upload_file_url_returns_metadata_without_ssrf( assert response.headers["Content-Type"] == "text/plain" assert response.headers["Content-Length"] == "5" assert response.request.method == "HEAD" - get_upload_file.assert_called_once() load_once.assert_not_called() ssrf_make_request.assert_not_called() -def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(monkeypatch: pytest.MonkeyPatch): +def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - get_upload_file = MagicMock() - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) url = f"http://localhost:5001/files/{UPLOAD_FILE_ID}/file-preview?timestamp=1700000000&nonce=nonce" proxy_response = httpx.Response(403, request=httpx.Request("GET", url)) ssrf_make_request = _patch_ssrf_make_request(monkeypatch, proxy_response) @@ -336,7 +338,7 @@ def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(monkeypatch: response = remote_fetcher.make_request("GET", url, timeout=3) assert response is proxy_response - get_upload_file.assert_not_called() + assert file_database.statements == [] ssrf_make_request.assert_called_once_with( method="GET", url=url, @@ -345,10 +347,10 @@ def test_make_request_get_unsigned_dify_url_delegates_to_ssrf_proxy(monkeypatch: ) -def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(monkeypatch: pytest.MonkeyPatch): +def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - get_upload_file = MagicMock() - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) proxy_response = httpx.Response(201, request=httpx.Request("POST", f"http://localhost:5001/files/{UPLOAD_FILE_ID}")) ssrf_make_request = _patch_ssrf_make_request(monkeypatch, proxy_response) url = _signed_url( @@ -360,7 +362,7 @@ def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(monkey response = remote_fetcher.make_request("POST", url, json={"name": "ignored"}) assert response is proxy_response - get_upload_file.assert_not_called() + assert file_database.statements == [] ssrf_make_request.assert_called_once_with( method="POST", url=url, @@ -369,20 +371,17 @@ def test_make_request_post_signed_upload_file_url_delegates_to_ssrf_proxy(monkey ) -def test_get_signed_image_preview_url_uses_image_preview_signature(monkeypatch: pytest.MonkeyPatch): +def test_get_signed_image_preview_url_uses_image_preview_signature( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=UPLOAD_FILE_ID, - key="upload_files/tenant/image.png", - name="image.png", - mime_type="image/png", - size=6, - extension="png", - ) - get_upload_file = MagicMock(return_value=upload_file) + file_database.upload_file.key = "upload_files/tenant/image.png" + file_database.upload_file.name = "image.png" + file_database.upload_file.mime_type = "image/png" + file_database.upload_file.size = 6 + file_database.upload_file.extension = "png" + file_database.session.commit() load_once = MagicMock(return_value=b"image!") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( @@ -396,7 +395,6 @@ def test_get_signed_image_preview_url_uses_image_preview_signature(monkeypatch: assert response.status_code == 200 assert response.content == b"image!" assert response.headers["Content-Type"] == "image/png" - get_upload_file.assert_called_once() load_once.assert_called_once_with("upload_files/tenant/image.png") ssrf_make_request.assert_not_called() @@ -597,11 +595,12 @@ def test_invalid_configured_file_origin_delegates_to_ssrf_proxy(monkeypatch: pyt ) -def test_signed_upload_file_url_returns_404_when_record_missing(monkeypatch: pytest.MonkeyPatch): +def test_signed_upload_file_url_returns_404_when_record_missing( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - get_upload_file = MagicMock(return_value=None) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) + file_database.session.delete(file_database.upload_file) + file_database.session.commit() ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( base_url="http://localhost:5001", @@ -613,23 +612,12 @@ def test_signed_upload_file_url_returns_404_when_record_missing(monkeypatch: pyt assert response.status_code == 404 assert response.content == b"" - get_upload_file.assert_called_once() ssrf_make_request.assert_not_called() def test_get_signed_tool_file_url_reads_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch): _patch_file_fetcher_config(monkeypatch) - session = _patch_session(monkeypatch) - tool_file = SimpleNamespace( - id=TOOL_FILE_ID, - file_key="tools/tenant/result.txt", - name="result.txt", - mimetype="text/plain", - size=6, - ) - get_tool_file = MagicMock(return_value=tool_file) load_once = MagicMock(return_value=b"result") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( @@ -643,19 +631,16 @@ def test_get_signed_tool_file_url_reads_storage_without_ssrf(monkeypatch: pytest assert response.status_code == 200 assert response.content == b"result" assert response.headers["Content-Type"] == "text/plain" - get_tool_file.assert_called_once_with( - session=session, - file_id=TOOL_FILE_ID, - ) load_once.assert_called_once_with("tools/tenant/result.txt") ssrf_make_request.assert_not_called() -def test_signed_tool_file_url_returns_404_when_record_missing(monkeypatch: pytest.MonkeyPatch): +def test_signed_tool_file_url_returns_404_when_record_missing( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - get_tool_file = MagicMock(return_value=None) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) + file_database.session.delete(file_database.tool_file) + file_database.session.commit() ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( base_url="http://localhost:5001", @@ -667,26 +652,12 @@ def test_signed_tool_file_url_returns_404_when_record_missing(monkeypatch: pytes assert response.status_code == 404 assert response.content == b"" - get_tool_file.assert_called_once() ssrf_make_request.assert_not_called() def test_get_signed_datasource_file_url_reads_upload_storage_without_ssrf(monkeypatch: pytest.MonkeyPatch): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - upload_file = SimpleNamespace( - id=DATASOURCE_FILE_ID, - key="datasources/tenant/data.txt", - name="data.txt", - mime_type="text/plain", - size=4, - extension="txt", - ) - get_upload_file = MagicMock(return_value=upload_file) - get_tool_file = MagicMock(return_value=None) load_once = MagicMock(return_value=b"data") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( @@ -699,27 +670,28 @@ def test_get_signed_datasource_file_url_reads_upload_storage_without_ssrf(monkey assert response.status_code == 200 assert response.content == b"data" - get_upload_file.assert_called_once() - get_tool_file.assert_not_called() load_once.assert_called_once_with("datasources/tenant/data.txt") ssrf_make_request.assert_not_called() -def test_get_signed_datasource_file_url_reads_tool_storage_when_upload_missing(monkeypatch: pytest.MonkeyPatch): +def test_get_signed_datasource_file_url_reads_tool_storage_when_upload_missing( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - tool_file = SimpleNamespace( - id=DATASOURCE_FILE_ID, + file_database.session.delete(file_database.datasource_upload_file) + datasource_tool_file = ToolFile( + user_id=USER_ID, + tenant_id=TENANT_ID, + conversation_id=None, file_key="datasources/tenant/tool-data.txt", name="tool-data.txt", mimetype="text/plain", size=9, ) - get_upload_file = MagicMock(return_value=None) - get_tool_file = MagicMock(return_value=tool_file) + datasource_tool_file.id = DATASOURCE_FILE_ID + file_database.session.add(datasource_tool_file) + file_database.session.commit() load_once = MagicMock(return_value=b"tool-data") - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) monkeypatch.setattr(remote_fetcher.storage, "load_once", load_once) ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( @@ -734,19 +706,16 @@ def test_get_signed_datasource_file_url_reads_tool_storage_when_upload_missing(m assert response.content == b"tool-data" assert response.headers["Content-Type"] == "text/plain" assert response.headers["Content-Length"] == "9" - get_upload_file.assert_called_once() - get_tool_file.assert_called_once() load_once.assert_called_once_with("datasources/tenant/tool-data.txt") ssrf_make_request.assert_not_called() -def test_signed_datasource_file_url_returns_404_when_records_missing(monkeypatch: pytest.MonkeyPatch): +def test_signed_datasource_file_url_returns_404_when_records_missing( + monkeypatch: pytest.MonkeyPatch, file_database: FileDatabase +): _patch_file_fetcher_config(monkeypatch) - _patch_session(monkeypatch) - get_upload_file = MagicMock(return_value=None) - get_tool_file = MagicMock(return_value=None) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_upload_file", get_upload_file) - monkeypatch.setattr(remote_fetcher._file_access_controller, "get_tool_file", get_tool_file) + file_database.session.delete(file_database.datasource_upload_file) + file_database.session.commit() ssrf_make_request = _patch_ssrf_make_request(monkeypatch) url = _signed_url( base_url="http://localhost:5001", @@ -758,8 +727,6 @@ def test_signed_datasource_file_url_returns_404_when_records_missing(monkeypatch assert response.status_code == 404 assert response.content == b"" - get_upload_file.assert_called_once() - get_tool_file.assert_called_once() ssrf_make_request.assert_not_called()