mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 02:28:30 +08:00
test: use sqlite3 session in test_remote_fetcher (#38773)
This commit is contained in:
Asuka Minato
GitHub
parent
81c3aa2d7e
commit
e21071cbb8
@ -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()
|
||||
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user